Skip to content

Commit 8e10846

Browse files
committed
further checks on kb
1 parent ed2bc08 commit 8e10846

4 files changed

Lines changed: 249 additions & 14 deletions

File tree

apps/sim/app/api/knowledge/secret-provenance.test.ts

Lines changed: 118 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,17 +10,132 @@ import {
1010
PRIVATE_SECRET_PROVENANCE_FIELD,
1111
PRIVATE_SECRET_PROVENANCE_HEADER,
1212
} from '@/lib/execution/private-tool-metadata'
13-
import { resolveKnowledgeWriteSecretProvenance } from '@/app/api/knowledge/secret-provenance'
13+
import {
14+
resolveKnowledgeDocumentWriteSecretProvenance,
15+
resolveKnowledgeWriteSecretProvenance,
16+
} from '@/app/api/knowledge/secret-provenance'
17+
18+
function createRequest(
19+
payload: Record<string, unknown>,
20+
provenanceHeader = PRIVATE_SECRET_PROVENANCE_BUNDLE_V1
21+
): NextRequest {
22+
return new NextRequest('http://localhost/api/knowledge/kb/documents', {
23+
method: 'POST',
24+
headers: { [PRIVATE_SECRET_PROVENANCE_HEADER]: provenanceHeader },
25+
body: JSON.stringify(payload),
26+
})
27+
}
1428

15-
function createRequest(payload: Record<string, unknown>): NextRequest {
29+
function createHeaderlessRequest(payload: Record<string, unknown>): NextRequest {
1630
return new NextRequest('http://localhost/api/knowledge/kb/documents', {
1731
method: 'POST',
18-
headers: { [PRIVATE_SECRET_PROVENANCE_HEADER]: PRIVATE_SECRET_PROVENANCE_BUNDLE_V1 },
1932
body: JSON.stringify(payload),
2033
})
2134
}
2235

2336
describe('knowledge write secret provenance', () => {
37+
it('does not track durable provenance for a headerless external chunk write', () => {
38+
const payload = { content: 'manual content' }
39+
40+
const result = resolveKnowledgeWriteSecretProvenance({
41+
request: createHeaderlessRequest(payload),
42+
payload,
43+
authType: AuthType.API_KEY,
44+
userId: 'user-1',
45+
workspaceId: 'workspace-1',
46+
selectionKeys: ['chunk-content'],
47+
})
48+
49+
expect(result).toEqual({ success: true })
50+
})
51+
52+
it('does not track durable provenance for a headerless external document write', () => {
53+
const payload = {
54+
filename: 'manual.txt',
55+
documentTagsData: JSON.stringify([{ tagName: 'source', tagValue: 'manual' }]),
56+
}
57+
58+
const result = resolveKnowledgeDocumentWriteSecretProvenance({
59+
request: createHeaderlessRequest(payload),
60+
payload,
61+
authType: AuthType.SESSION,
62+
userId: 'user-1',
63+
workspaceId: 'workspace-1',
64+
documents: [payload],
65+
})
66+
67+
expect(result).toEqual({ success: true })
68+
})
69+
70+
it('does not track durable provenance for a legacy headerless internal write', () => {
71+
const payload = { content: 'legacy workflow content' }
72+
73+
const result = resolveKnowledgeWriteSecretProvenance({
74+
request: createHeaderlessRequest(payload),
75+
payload,
76+
authType: AuthType.INTERNAL_JWT,
77+
userId: 'user-1',
78+
workspaceId: 'workspace-1',
79+
selectionKeys: ['chunk-content'],
80+
})
81+
82+
expect(result).toEqual({ success: true })
83+
})
84+
85+
it('tracks exact-empty provenance only when an internal write supplies a verified envelope', () => {
86+
const bundle = {
87+
version: 1 as const,
88+
complete: true,
89+
selections: [
90+
{
91+
key: 'chunk-content',
92+
provenance: { version: 1 as const, complete: true, entries: [] },
93+
},
94+
],
95+
}
96+
const payload = { content: 'workflow content', [PRIVATE_SECRET_PROVENANCE_FIELD]: bundle }
97+
98+
const result = resolveKnowledgeWriteSecretProvenance({
99+
request: createRequest(payload),
100+
payload,
101+
authType: AuthType.INTERNAL_JWT,
102+
userId: 'user-1',
103+
workspaceId: 'workspace-1',
104+
selectionKeys: ['chunk-content'],
105+
})
106+
107+
expect(result).toEqual({
108+
success: true,
109+
provenances: [{ status: 'exact', entries: [] }],
110+
})
111+
})
112+
113+
it('rejects a private provenance envelope from an external caller', () => {
114+
const bundle = {
115+
version: 1 as const,
116+
complete: true,
117+
selections: [
118+
{
119+
key: 'chunk-content',
120+
provenance: { version: 1 as const, complete: true, entries: [] },
121+
},
122+
],
123+
}
124+
const payload = { content: 'external content', [PRIVATE_SECRET_PROVENANCE_FIELD]: bundle }
125+
126+
const result = resolveKnowledgeWriteSecretProvenance({
127+
request: createRequest(payload),
128+
payload,
129+
authType: AuthType.API_KEY,
130+
userId: 'user-1',
131+
workspaceId: 'workspace-1',
132+
selectionKeys: ['chunk-content'],
133+
})
134+
135+
expect(result.success).toBe(false)
136+
if (!result.success) expect(result.response.status).toBe(400)
137+
})
138+
24139
it('rejects an unavailable verified selection before a write can start', () => {
25140
const bundle = {
26141
version: 1 as const,

apps/sim/app/api/knowledge/secret-provenance.ts

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import {
44
createDurableSecretProvenanceRegistry,
55
type DurableSecretProvenance,
66
durableSecretProvenanceFromPrivateBundle,
7-
EXACT_EMPTY_DURABLE_SECRET_PROVENANCE,
87
} from '@/lib/execution/durable-secret-provenance'
98
import {
109
inspectPrivateSecretProvenanceRequest,
@@ -49,12 +48,7 @@ export function resolveKnowledgeWriteSecretProvenance(options: {
4948
const { request } = options
5049
const inspection = inspectPrivateSecretProvenanceRequest(request.headers, options.payload)
5150
if (inspection.status === 'unsupported') {
52-
return options.authType === AuthType.INTERNAL_JWT
53-
? { success: true }
54-
: {
55-
success: true,
56-
provenances: options.selectionKeys.map(() => EXACT_EMPTY_DURABLE_SECRET_PROVENANCE),
57-
}
51+
return { success: true }
5852
}
5953
if (inspection.status !== 'verified' || options.authType !== AuthType.INTERNAL_JWT) {
6054
return { success: false, response: invalidKnowledgeProvenanceResponse() }

apps/sim/lib/knowledge/secret-provenance.ts

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -166,10 +166,41 @@ export function rebindKnowledgeDocumentSecretProvenance(
166166
provenance: DurableSecretProvenance,
167167
previousSource: KnowledgeDocumentSourceValue,
168168
nextSource: KnowledgeDocumentSourceValue
169+
): DurableSecretProvenance {
170+
return rebindKnowledgeDocumentSecretProvenanceFields(
171+
provenance,
172+
previousSource,
173+
nextSource,
174+
new Set()
175+
)
176+
}
177+
178+
/** Rebinds unchanged source fields while dropping provenance for metadata fields being cleared. */
179+
export function rebindKnowledgeDocumentSecretProvenanceAfterMetadataClear(
180+
provenance: DurableSecretProvenance,
181+
previousSource: KnowledgeDocumentSourceValue,
182+
nextSource: KnowledgeDocumentSourceValue,
183+
clearedFields: ReadonlySet<KnowledgeDocumentMetadataField>
184+
): DurableSecretProvenance {
185+
return rebindKnowledgeDocumentSecretProvenanceFields(
186+
provenance,
187+
previousSource,
188+
nextSource,
189+
clearedFields
190+
)
191+
}
192+
193+
function rebindKnowledgeDocumentSecretProvenanceFields(
194+
provenance: DurableSecretProvenance,
195+
previousSource: KnowledgeDocumentSourceValue,
196+
nextSource: KnowledgeDocumentSourceValue,
197+
excludedFields: ReadonlySet<KnowledgeDocumentMetadataField>
169198
): DurableSecretProvenance {
170199
if (provenance.status === 'unknown') return provenance
171200
if (provenance.entries.some((entry) => !entry.sourceValueHash)) return { status: 'unknown' }
172-
const rebound = KNOWLEDGE_DOCUMENT_METADATA_FIELDS.map((field) =>
201+
const rebound = KNOWLEDGE_DOCUMENT_METADATA_FIELDS.filter(
202+
(field) => !excludedFields.has(field)
203+
).map((field) =>
173204
bindKnowledgeDocumentFieldSecretProvenance(
174205
filterDurableSecretProvenanceBySourceValues(provenance, [
175206
createKnowledgeDocumentFieldBinding(field, previousSource[field]),

apps/sim/lib/knowledge/tags/service.ts

Lines changed: 98 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,22 @@
11
import { db } from '@sim/db'
2-
import { document, embedding, knowledgeBaseTagDefinitions } from '@sim/db/schema'
2+
import {
3+
document,
4+
documentSecretProvenance,
5+
embedding,
6+
knowledgeBaseTagDefinitions,
7+
} from '@sim/db/schema'
38
import { createLogger } from '@sim/logger'
49
import { generateId } from '@sim/utils/id'
5-
import { and, eq, isNotNull, isNull, sql } from 'drizzle-orm'
6-
import type { DbOrTx } from '@/lib/db/types'
10+
import { and, eq, inArray, isNotNull, isNull, or, sql } from 'drizzle-orm'
11+
import type { DbOrTx, DbTransaction } from '@/lib/db/types'
712
import { getSlotsForFieldType, SUPPORTED_FIELD_TYPES } from '@/lib/knowledge/constants'
13+
import {
14+
createKnowledgeDocumentSourceValue,
15+
type KnowledgeDocumentMetadataField,
16+
readBoundKnowledgeDocumentSecretProvenance,
17+
rebindKnowledgeDocumentSecretProvenanceAfterMetadataClear,
18+
replaceKnowledgeDocumentSecretProvenanceInTx,
19+
} from '@/lib/knowledge/secret-provenance'
820
import type { BulkTagDefinitionsData, DocumentTagDefinition } from '@/lib/knowledge/tags/types'
921
import type {
1022
CreateTagDefinitionData,
@@ -32,6 +44,7 @@ const VALID_TAG_SLOTS = [
3244
] as const
3345

3446
type ValidTagSlot = (typeof VALID_TAG_SLOTS)[number]
47+
type ClearedTagValues = Partial<Record<ValidTagSlot, null>>
3548

3649
/**
3750
* Validates that a tag slot is a valid slot name
@@ -42,6 +55,88 @@ function validateTagSlot(tagSlot: string): asserts tagSlot is ValidTagSlot {
4255
}
4356
}
4457

58+
async function clearTagSlotsInTx(
59+
tx: DbTransaction,
60+
knowledgeBaseId: string,
61+
tagSlots: readonly ValidTagSlot[]
62+
): Promise<void> {
63+
if (tagSlots.length === 0) return
64+
65+
const clearedTagValues: ClearedTagValues = {}
66+
for (const tagSlot of tagSlots) clearedTagValues[tagSlot] = null
67+
68+
await tx
69+
.update(embedding)
70+
.set(clearedTagValues)
71+
.where(
72+
and(
73+
eq(embedding.knowledgeBaseId, knowledgeBaseId),
74+
or(...tagSlots.map((tagSlot) => isNotNull(embedding[tagSlot])))
75+
)
76+
)
77+
78+
const currentDocuments = await tx
79+
.select()
80+
.from(document)
81+
.where(
82+
and(
83+
eq(document.knowledgeBaseId, knowledgeBaseId),
84+
or(...tagSlots.map((tagSlot) => isNotNull(document[tagSlot])))
85+
)
86+
)
87+
.for('update')
88+
89+
const trackedDocuments = currentDocuments.filter(
90+
(current) => current.secretProvenanceVersion === 1
91+
)
92+
const trackedDocumentIds = trackedDocuments.map((current) => current.id)
93+
const sidecars = trackedDocumentIds.length
94+
? await tx
95+
.select()
96+
.from(documentSecretProvenance)
97+
.where(inArray(documentSecretProvenance.documentId, trackedDocumentIds))
98+
: []
99+
const sidecarByDocumentId = new Map(sidecars.map((sidecar) => [sidecar.documentId, sidecar]))
100+
101+
await tx
102+
.update(document)
103+
.set(clearedTagValues)
104+
.where(
105+
and(
106+
eq(document.knowledgeBaseId, knowledgeBaseId),
107+
or(...tagSlots.map((tagSlot) => isNotNull(document[tagSlot])))
108+
)
109+
)
110+
111+
const clearedFields = new Set<KnowledgeDocumentMetadataField>(tagSlots)
112+
for (const current of trackedDocuments) {
113+
const sidecar = sidecarByDocumentId.get(current.id)
114+
const currentSource = createKnowledgeDocumentSourceValue(current)
115+
const currentProvenance = readBoundKnowledgeDocumentSecretProvenance({
116+
secretProvenanceVersion: current.secretProvenanceVersion,
117+
source: currentSource,
118+
provenanceSourceHash: sidecar?.sourceHash ?? null,
119+
status: sidecar?.status ?? null,
120+
entries: sidecar?.entries,
121+
})
122+
const nextSource = createKnowledgeDocumentSourceValue({
123+
...current,
124+
...clearedTagValues,
125+
})
126+
await replaceKnowledgeDocumentSecretProvenanceInTx(
127+
tx,
128+
current.id,
129+
nextSource,
130+
rebindKnowledgeDocumentSecretProvenanceAfterMetadataClear(
131+
currentProvenance,
132+
currentSource,
133+
nextSource,
134+
clearedFields
135+
)
136+
)
137+
}
138+
}
139+
45140
/**
46141
* Get the field type for a tag slot
47142
*/

0 commit comments

Comments
 (0)