opensource/sim
3
0
Fork 0
mirror of https://github.com/simstudioai/sim.git synced 2026-08-31 01:11:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

improvement(provenance): cleanup secrets boundary (#6374)

Browse Source 
* fix(secrets): preserve raw outputs with durable provenance

* improvement(provenance): cleanup boundary

* fix copy resources

* fix fork copies to work with provenance

* address comments

* fix
This commit is contained in:
Vikhyath Mondreti
2026-08-07 17:12:17 -07:00
committed by GitHub
parent 40c0a571fd
commit 5cf1f9be84
287 changed files with 15274 additions and 5960 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/sim/app/api/function/execute/route.test.ts
+175 -13
View File
@@ -590,7 +590,7 @@ describe('Function Execute API Route', () => {
createMockRequest(
'POST',
{
code: 'return environmentVariables.API_KEY',
code: 'return {{API_KEY}}',
envVars: { API_KEY: 'secret-at-the-end' },
workflowId: 'workflow-1',
workspaceId: 'workspace-1',
@@ -686,7 +686,7 @@ describe('Function Execute API Route', () => {
createMockRequest(
'POST',
{
code: 'print("done")',
code: 'print("{{API_KEY}}")',
language: 'python',
workspaceId: 'workspace-1',
envVars: { API_KEY: 'secret-value' },
@@ -821,6 +821,54 @@ describe('Function Execute API Route', () => {
expect(mockExecuteInSandbox).not.toHaveBeenCalled()
})
it('runs with authenticated incomplete mount provenance and marks exported bytes unknown', async () => {
envFlagsMock.isRemoteSandboxEnabled = true
mockExecuteInSandbox.mockResolvedValueOnce({
result: 'raw result',
stdout: '',
sandboxId: 'sandbox-123',
exportedFiles: { '/home/user/output.txt': 'raw output' },
})
const response = await POST(
createMockRequest(
'POST',
{
code: 'print("done")',
language: 'python',
workspaceId: 'workspace-1',
outputs: {
files: [
{
path: 'files/output.txt',
sandboxPath: '/home/user/output.txt',
mimeType: 'text/plain',
},
],
},
[PRIVATE_SECRET_PROVENANCE_FIELD]: {
version: 1,
complete: false,
selections: [],
},
},
{ [PRIVATE_SECRET_PROVENANCE_HEADER]: PRIVATE_SECRET_PROVENANCE_BUNDLE_V1 }
)
)
expect(response.status).toBe(200)
expect((await response.json()).output.result).toEqual(
expect.objectContaining({ fileId: 'wf_output_txt', vfsPath: 'files/output.txt' })
)
expect(mockExecuteInSandbox).toHaveBeenCalledOnce()
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledWith(
expect.objectContaining({
buffer: Buffer.from('raw output'),
secretProvenance: { status: 'unknown' },
})
)
})
it('does not rewrite a static export path that happens to equal a resolved secret', async () => {
envFlagsMock.isRemoteSandboxEnabled = true
mockExecuteInSandbox.mockResolvedValueOnce({
@@ -853,6 +901,7 @@ describe('Function Execute API Route', () => {
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledWith(
expect.objectContaining({
target: expect.objectContaining({ path: 'files/report-secret-value.txt' }),
secretProvenance: { status: 'exact', entries: [] },
})
)
expect(JSON.stringify(data)).toContain('files/report-secret-value.txt')
@@ -890,7 +939,7 @@ describe('Function Execute API Route', () => {
)
})
it('keeps a binary export unknown when files were mounted without a provenance envelope', async () => {
it('classifies a binary export exact-empty when ordinary files were mounted without secret provenance', async () => {
envFlagsMock.isRemoteSandboxEnabled = true
mockExecuteInSandbox.mockResolvedValueOnce({
result: 'done',
@@ -919,7 +968,7 @@ describe('Function Execute API Route', () => {
expect(response.status).toBe(200)
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledWith(
expect.objectContaining({ secretProvenance: { status: 'unknown' } })
expect.objectContaining({ secretProvenance: { status: 'exact', entries: [] } })
)
})
@@ -987,7 +1036,7 @@ describe('Function Execute API Route', () => {
const response = await POST(
createMockRequest('POST', {
code: 'print("done")',
code: 'print("{{API_KEY}}")',
language: 'python',
workspaceId: 'workspace-1',
envVars: { API_KEY: 'secret-value' },
@@ -2002,6 +2051,117 @@ describe('Function Execute API Route', () => {
expect(Object.values(request.contextVariables)).not.toContain('must-not-bind')
})
it('does not infer provenance from an unused low-entropy environment value', async () => {
mockExecuteInIsolatedVM.mockResolvedValueOnce({ result: 'Box eSign', stdout: '' })
const response = await POST(
createMockRequest(
'POST',
{
code: 'return "Box eSign"',
envVars: { SERVICENOW_PASSWORD: 'x' },
},
{ 'x-sim-request-private-tool-metadata': 'resolved-secret-names-v1' }
)
)
const data = await response.json()
expect(response.status).toBe(200)
expect(data.output.result).toBe('Box eSign')
expect(data.__resolvedSecretNames).toEqual([])
})
it('does not build provenance matchers for unused oversized environment values', async () => {
mockExecuteInIsolatedVM.mockResolvedValueOnce({ result: 'safe', stdout: '' })
const response = await POST(
createMockRequest(
'POST',
{
code: 'return "safe"',
envVars: { UNUSED: 'x'.repeat(65 * 1024) },
},
{ 'x-sim-request-private-tool-metadata': 'resolved-secret-names-v1' }
)
)
expect(response.status).toBe(200)
expect((await response.json()).__resolvedSecretNames).toEqual([])
})
it('conservatively reports only compiled secrets when bounded output classification is exceeded', async () => {
const result = Array.from({ length: 100_001 }, () => 'ordinary')
mockExecuteInIsolatedVM.mockResolvedValueOnce({ result, stdout: '' })
const response = await POST(
createMockRequest(
'POST',
{
code: 'const key = {{API_KEY}}; return params.items',
params: { items: result },
envVars: { API_KEY: 'secret-value', UNUSED: 'x' },
},
{ 'x-sim-request-private-tool-metadata': 'resolved-secret-names-v1' }
)
)
const data = await response.json()
expect(response.status).toBe(200)
expect(response.headers.get('x-sim-private-tool-metadata')).toBe('resolved-secret-names-v1')
expect(data.output.result).toHaveLength(100_001)
expect(data.output.result[0]).toBe('ordinary')
expect(data.__resolvedSecretNames).toEqual(['API_KEY'])
})
it('conservatively reports a compiled secret whose value exceeds matcher capacity', async () => {
mockExecuteInIsolatedVM.mockResolvedValueOnce({ result: 'ordinary', stdout: '' })
const response = await POST(
createMockRequest(
'POST',
{
code: 'const key = {{OVERSIZED_SECRET}}; return "ordinary"',
envVars: { OVERSIZED_SECRET: 's'.repeat(64 * 1024 + 1), UNUSED: 'x' },
},
{ 'x-sim-request-private-tool-metadata': 'resolved-secret-names-v1' }
)
)
const data = await response.json()
expect(response.status).toBe(200)
expect(data.output.result).toBe('ordinary')
expect(data.__resolvedSecretNames).toEqual(['OVERSIZED_SECRET'])
})
it('tracks only compiled names when configured secrets share the same value', async () => {
mockExecuteInIsolatedVM.mockResolvedValueOnce({ result: 'true', stdout: '' })
const oneResponse = await POST(
createMockRequest(
'POST',
{
code: 'return {{SECOND}}',
envVars: { FIRST: 'true', SECOND: 'true' },
},
{ 'x-sim-request-private-tool-metadata': 'resolved-secret-names-v1' }
)
)
mockExecuteInIsolatedVM.mockResolvedValueOnce({ result: 'true', stdout: '' })
const bothResponse = await POST(
createMockRequest(
'POST',
{
code: 'const first = {{FIRST}}; return {{SECOND}}',
envVars: { FIRST: 'true', SECOND: 'true' },
},
{ 'x-sim-request-private-tool-metadata': 'resolved-secret-names-v1' }
)
)
expect((await oneResponse.json()).__resolvedSecretNames).toEqual(['SECOND'])
expect((await bothResponse.json()).__resolvedSecretNames).toEqual(['FIRST', 'SECOND'])
})
it('lowers missing shell placeholders while preserving comments and heredoc delimiters', async () => {
envFlagsMock.isRemoteSandboxEnabled = true
const response = await POST(
@@ -2134,7 +2294,7 @@ describe('Function Execute API Route', () => {
expect(mockExecuteInSandbox).not.toHaveBeenCalled()
})
it('reports exact secret values returned through placeholders and the environment map', async () => {
it('reports exact secret values returned through placeholders without inferring direct environment reads', async () => {
mockExecuteInIsolatedVM.mockResolvedValueOnce({
result: 'secret-valueother-secret',
stdout: '',
@@ -2171,14 +2331,15 @@ describe('Function Execute API Route', () => {
const directData = await directResponse.json()
expect(envData.__resolvedSecretNames).toEqual(['ENV_ONLY', 'SHARED'])
expect(directData.__resolvedSecretNames).toEqual(['API_KEY'])
expect(directData.output.result).toBe('secret-value')
expect(directData.__resolvedSecretNames).toEqual([])
})
it.each([
{ name: 'numeric', secret: '123', result: 123 },
{ name: 'boolean', secret: 'true', result: true },
])(
'records provenance for a typed $name secret returned through direct environment access',
'preserves a typed $name value returned through legacy direct environment access without inferred provenance',
async ({ secret, result }) => {
mockExecuteInIsolatedVM.mockResolvedValueOnce({ result, stdout: '' })
@@ -2197,11 +2358,11 @@ describe('Function Execute API Route', () => {
const data = await response.json()
expect(data.output.result).toBe(result)
expect(data.__resolvedSecretNames).toEqual(['API_KEY'])
expect(data.__resolvedSecretNames).toEqual([])
}
)
it('reports shell substitutions and exact secret output from direct environment access', async () => {
it('reports placeholder output without inferring provenance from legacy shell environment access', async () => {
envFlagsMock.isRemoteSandboxEnabled = true
mockExecuteShellInSandbox.mockResolvedValueOnce({
result: null,
@@ -2245,7 +2406,8 @@ describe('Function Execute API Route', () => {
const directData = await directResponse.json()
expect(referencedData.__resolvedSecretNames).toEqual(['API_KEY'])
expect(directData.__resolvedSecretNames).toEqual(['API_KEY'])
expect(directData.output.stdout).toBe('secret-value')
expect(directData.__resolvedSecretNames).toEqual([])
})
it('returns nonzero shell stderr as a visible 422 error and diagnostic output', async () => {
@@ -2289,8 +2451,8 @@ describe('Function Execute API Route', () => {
)
expect(response.status).toBe(200)
expect((await response.json()).__resolvedSecretNames).toBeUndefined()
expect(response.headers.get('x-sim-private-tool-metadata')).toBeNull()
expect((await response.json()).__resolvedSecretNames).toEqual([])
expect(response.headers.get('x-sim-private-tool-metadata')).toBe('resolved-secret-names-v1')
expect(mockExecuteInIsolatedVM).toHaveBeenCalled()
})
apps/sim/app/api/function/execute/route.ts
+66 -70
View File
@@ -113,8 +113,6 @@ const TAG_PATTERN = createReferencePattern()
const E2B_JS_WRAPPER_LINES = 3
const E2B_PYTHON_WRAPPER_LINES = 1
const MAX_SANDBOX_OUTPUT_FILES = 20
const MAX_PRIVATE_RESOLVED_SECRET_NAMES = 10_000
const MAX_PRIVATE_RESOLVED_SECRET_NAMES_BYTES = 1024 * 1024
const MAX_PRIVATE_FILE_SECRET_MATCH_EVENTS = 1_000_000
const SANDBOX_RUNTIME_PAYLOAD_PATH_ENV = '__SIM_RUNTIME_PAYLOAD_PATH'
@@ -984,12 +982,10 @@ interface FunctionRouteExecutionContext {
resolvedSecretNames: Set<string>
includePrivateResolvedSecretNames: boolean
privateResolvedSecretNamesMetadataType?: ResolvedSecretNamesMetadataType
outputProvenanceComplete: boolean
outputSecretMatcher?: ResolvedSecretMatcher
outputSecretNamesByScanLiteral: Map<string, string[]>
outputSecretPlaintextsByName: Map<string, string>
mountedFileSecretProvenanceScanner?: MountedFileSecretProvenanceScanner
hasMountedSandboxFiles: boolean
}
type ResolvedSecretNamesMetadataType =
@@ -1007,10 +1003,16 @@ function inspectMountedWorkspaceFileProvenance(
): MountedWorkspaceFileProvenanceInspection {
const inspection = inspectPrivateSecretProvenanceRequest(headers, body)
if (inspection.status === 'unsupported') return { status: 'none' }
if (inspection.status !== 'verified' || !isPrivateSecretProvenanceBundleV1(inspection.value)) {
return { status: 'invalid' }
}
if (!inspection.value.complete) {
return {
status: 'verified',
provenance: { version: 1, complete: false, entries: [] },
}
}
if (
inspection.status !== 'verified' ||
!isPrivateSecretProvenanceBundleV1(inspection.value) ||
!inspection.value.complete ||
inspection.value.selections.length !== 1 ||
inspection.value.selections[0]?.key !== MOUNTED_WORKSPACE_FILES_PROVENANCE_KEY
) {
@@ -1188,7 +1190,10 @@ function activateOutputSecretProvenance(
body: unknown,
context: FunctionRouteExecutionContext
): void {
if (!context.outputSecretMatcher) return
if (!context.outputSecretMatcher) {
activateCompiledSecretProvenance(context)
return
}
const matchedPlaintexts = new Set<string>()
const projection = projectResolvedSecretContent(
@@ -1200,7 +1205,7 @@ function activateOutputSecretProvenance(
}
)
if (!projection.safe) {
context.outputProvenanceComplete = false
activateCompiledSecretProvenance(context)
return
}
for (const plaintext of matchedPlaintexts) {
@@ -1211,19 +1216,24 @@ function activateOutputSecretProvenance(
}
/**
* True when any secret material was in scope for this execution — a mounted environment secret, or
* a secret carried by a mounted input file. When false, nothing secret ever reached the sandbox, so
* no export of any kind can carry one.
*
* Mounted bytes are classified from the caller's provenance envelope. Files mounted *without* one
* are unclassifiable rather than clean: absence of an envelope is absence of evidence, not evidence
* the mount carried nothing. Those fail closed here so the classification can never be stronger
* than what the caller actually attested to.
* Conservatively activates only secrets whose placeholders were compiled for this invocation.
* This fallback is used when the bounded output classifier cannot inspect a result; it never
* considers configured-but-unused environment values and never mutates the functional result.
*/
function activateCompiledSecretProvenance(context: FunctionRouteExecutionContext): void {
for (const name of context.outputSecretPlaintextsByName.keys()) {
context.resolvedSecretNames.add(name)
}
}
/**
* True when this execution compiled a secret placeholder or received a mounted file with verified
* secret provenance. Ordinary mounts without a provenance envelope are user data, not evidence that
* a Sim secret was resolved in this call.
*/
function hasSecretMaterialInScope(context: FunctionRouteExecutionContext): boolean {
if (context.outputSecretPlaintextsByName.size > 0) return true
const scanner = context.mountedFileSecretProvenanceScanner
return scanner ? scanner.hasSecrets : context.hasMountedSandboxFiles
return context.mountedFileSecretProvenanceScanner?.hasSecrets ?? false
}
/**
@@ -1268,7 +1278,6 @@ async function getOutputFileSecretProvenance(
MAX_PRIVATE_FILE_SECRET_MATCH_EVENTS
)
} catch {
context.outputProvenanceComplete = false
return { status: 'unknown' }
}
@@ -1293,17 +1302,8 @@ async function getOutputFileSecretProvenance(
}
}
function getPrivateResolvedSecretNames(context: FunctionRouteExecutionContext): string[] | null {
if (!context.outputProvenanceComplete) return null
if (context.resolvedSecretNames.size > MAX_PRIVATE_RESOLVED_SECRET_NAMES) return null
const names = Array.from(context.resolvedSecretNames).sort()
let bytes = 0
for (const name of names) {
bytes += Buffer.byteLength(name, 'utf8')
if (bytes > MAX_PRIVATE_RESOLVED_SECRET_NAMES_BYTES) return null
}
return names
function getPrivateResolvedSecretNames(context: FunctionRouteExecutionContext): string[] {
return Array.from(context.resolvedSecretNames).sort()
}
async function appendResolvedSecretNames(
@@ -1327,21 +1327,17 @@ async function appendPrivateResolvedSecretNames(
): Promise<NextResponse> {
if (!names || !metadataType) return response
try {
const body = (await response.clone().json()) as Record<string, unknown>
const headers = new Headers(response.headers)
headers.delete('content-length')
headers.set(PRIVATE_TOOL_METADATA_RESPONSE_HEADER, metadataType)
return NextResponse.json(
{
...body,
[RESOLVED_SECRET_NAMES_FIELD]: names,
},
{ status: response.status, statusText: response.statusText, headers }
)
} catch {
return response
}
const body = (await response.json()) as Record<string, unknown>
const headers = new Headers(response.headers)
headers.delete('content-length')
headers.set(PRIVATE_TOOL_METADATA_RESPONSE_HEADER, metadataType)
return NextResponse.json(
{
...body,
[RESOLVED_SECRET_NAMES_FIELD]: names,
},
{ status: response.status, statusText: response.statusText, headers }
)
}
/**
@@ -2026,33 +2022,9 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
resolvedSecretNames: new Set<string>(),
includePrivateResolvedSecretNames,
privateResolvedSecretNamesMetadataType,
outputProvenanceComplete: true,
outputSecretNamesByScanLiteral: new Map(),
outputSecretPlaintextsByName: new Map(),
mountedFileSecretProvenanceScanner,
hasMountedSandboxFiles: (_sandboxFiles?.length ?? 0) > 0,
}
for (const [name, plaintext] of Object.entries(envVars)) {
if (!plaintext) continue
routeContext.outputSecretPlaintextsByName.set(name, plaintext)
const scanLiterals = new Set([plaintext, JSON.stringify(plaintext).slice(1, -1)])
for (const scanLiteral of scanLiterals) {
const names = routeContext.outputSecretNamesByScanLiteral.get(scanLiteral) ?? []
names.push(name)
routeContext.outputSecretNamesByScanLiteral.set(scanLiteral, names)
}
}
if (routeContext.outputSecretNamesByScanLiteral.size > 0) {
try {
routeContext.outputSecretMatcher = createResolvedSecretMatcher(
[...routeContext.outputSecretNamesByScanLiteral].map(([plaintext, names]) => ({
plaintext,
replacement: `{{${names[0]}}}`,
}))
)
} catch {
routeContext.outputProvenanceComplete = false
}
}
const lang = isValidCodeLanguage(language) ? language : DEFAULT_CODE_LANGUAGE
@@ -2081,6 +2053,30 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
environmentVariables: envVars,
reservedNames: Object.keys(contextVariables),
})
for (const name of compilation.resolvedSecretNames) {
if (!Object.hasOwn(envVars, name)) continue
const plaintext = envVars[name]
if (!plaintext) continue
routeContext.outputSecretPlaintextsByName.set(name, plaintext)
const scanLiterals = new Set([plaintext, JSON.stringify(plaintext).slice(1, -1)])
for (const scanLiteral of scanLiterals) {
const names = routeContext.outputSecretNamesByScanLiteral.get(scanLiteral) ?? []
names.push(name)
routeContext.outputSecretNamesByScanLiteral.set(scanLiteral, names)
}
}
if (routeContext.outputSecretNamesByScanLiteral.size > 0) {
try {
routeContext.outputSecretMatcher = createResolvedSecretMatcher(
[...routeContext.outputSecretNamesByScanLiteral].map(([plaintext, names]) => ({
plaintext,
replacement: `{{${[...names].sort()[0]}}}`,
}))
)
} catch {
activateCompiledSecretProvenance(routeContext)
}
}
resolvedCode = compilation.code
compilerInternalIdentifiers = [...compilation.internalIdentifiers]
compilerPrivateInputs = [...compilation.privateInputs]
apps/sim/app/api/guardrails/validate/route.test.ts
+6 -6
View File
@@ -115,11 +115,11 @@ describe('POST /api/guardrails/validate', () => {
},
})
mockValidateHallucination.mockResolvedValue({ passed: true, score: 8 })
mockImportProvenance.mockResolvedValue(true)
mockImportProvenance.mockResolvedValue({ success: true, matched: true })
mockRegistryIsComplete.mockReturnValue(true)
mockPrepareCopilotEnvironmentContext.mockResolvedValue({
resolvedSecretTraceRegistry: {
importProvenanceForValue: mockImportProvenance,
importProvenanceForValueAtInputPath: mockImportProvenance,
isComplete: mockRegistryIsComplete,
},
})
@@ -242,7 +242,7 @@ describe('POST /api/guardrails/validate', () => {
)
expect(res.status).toBe(200)
expect(mockImportProvenance).toHaveBeenCalledWith(provenance, 'secret value', {
expect(mockImportProvenance).toHaveBeenCalledWith(provenance, 'secret value', ['input'], {
trusted: true,
})
})
@@ -371,7 +371,7 @@ describe('POST /api/guardrails/validate', () => {
)
})
it('rejects a headerless internal hallucination check before model execution', async () => {
it('preserves a headerless legacy internal hallucination check', async () => {
hybridAuthMockFns.mockCheckSessionOrInternalAuth.mockResolvedValue({
success: true,
userId: 'user-1',
@@ -388,9 +388,9 @@ describe('POST /api/guardrails/validate', () => {
})
)
expect(res.status).toBe(400)
expect(res.status).toBe(200)
expect(mockImportProvenance).not.toHaveBeenCalled()
expect(mockValidateHallucination).not.toHaveBeenCalled()
expect(mockValidateHallucination).toHaveBeenCalled()
})
it('rejects invalid internal billing attribution as a protocol error', async () => {
apps/sim/app/api/guardrails/validate/route.ts
+8 -14
View File
@@ -257,25 +257,19 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
if (provenanceInspection.status === 'invalid') {
return NextResponse.json({ error: 'Invalid model input provenance' }, { status: 400 })
}
if (
provenanceInspection.status === 'unsupported' &&
auth.authType === AuthType.INTERNAL_JWT
) {
return NextResponse.json(
{ error: 'Model input provenance is unavailable' },
{ status: 400 }
)
}
if (provenanceInspection.status === 'verified' && auth.authType !== AuthType.INTERNAL_JWT) {
return NextResponse.json({ error: 'Invalid model input provenance' }, { status: 400 })
}
const provenanceReady =
provenanceInspection.status === 'verified'
? await resolvedSecretTraceRegistry.importProvenanceForValue(
provenanceInspection.value,
inputStr,
{ trusted: true }
)
? (
await resolvedSecretTraceRegistry.importProvenanceForValueAtInputPath(
provenanceInspection.value,
inputStr,
['input'],
{ trusted: true }
)
).success
: true
if (!provenanceReady || !resolvedSecretTraceRegistry.isComplete()) {
return NextResponse.json(
apps/sim/app/api/knowledge/secret-provenance.test.ts
+107 -3
View File
@@ -230,7 +230,7 @@ describe('knowledge write secret provenance', () => {
if (!result.success) expect(result.response.status).toBe(400)
})
it('rejects an unavailable verified selection before a write can start', () => {
it('persists authenticated unavailable selection lineage as unknown', () => {
const bundle = {
version: 1 as const,
complete: true,
@@ -257,7 +257,111 @@ describe('knowledge write secret provenance', () => {
selectionKeys: ['document-source:0'],
})
expect(result.success).toBe(false)
if (!result.success) expect(result.response.status).toBe(400)
expect(result).toEqual({ success: true, provenances: [{ status: 'unknown' }] })
})
it('persists an authenticated incomplete document bundle as unknown', () => {
const payload = {
[PRIVATE_SECRET_PROVENANCE_FIELD]: {
version: 1 as const,
complete: false,
selections: [],
},
}
const result = resolveKnowledgeDocumentWriteSecretProvenance({
request: createRequest(payload),
payload,
authType: AuthType.INTERNAL_JWT,
userId: 'user-1',
workspaceId: 'workspace-1',
documents: [{ documentTagsData: JSON.stringify([{ tagName: 'region', value: 'west' }]) }],
})
expect(result).toEqual({
success: true,
provenances: [
{
filename: { status: 'unknown' },
content: { status: 'unknown' },
tags: [{ tagName: 'region', provenance: { status: 'unknown' } }],
},
],
})
})
it('keeps persisted tag names raw while retaining tag-value provenance', () => {
const documentTagsData = JSON.stringify([{ tagName: 'private-name', value: 'west' }])
const payload = {
documents: [{ filename: 'doc.md', documentTagsData }],
[PRIVATE_SECRET_PROVENANCE_FIELD]: {
version: 1 as const,
complete: true,
selections: [
{
key: 'document-filename:0',
provenance: {
version: 1 as const,
complete: true,
entries: [],
scope: PRIVATE_PROVENANCE_SCOPE,
},
},
{
key: 'document-content:0',
provenance: {
version: 1 as const,
complete: true,
entries: [],
scope: PRIVATE_PROVENANCE_SCOPE,
},
},
{
key: 'document-tag-value:0:0',
provenance: {
version: 1 as const,
complete: true,
entries: [{ name: 'TAG_VALUE', encryptedValue: 'encrypted-tag-value' }],
scope: PRIVATE_PROVENANCE_SCOPE,
},
},
],
},
}
const result = resolveKnowledgeDocumentWriteSecretProvenance({
request: createRequest(payload),
payload,
authType: AuthType.INTERNAL_JWT,
userId: 'user-1',
workspaceId: 'workspace-1',
documents: [{ documentTagsData }],
})
expect(result).toEqual({
success: true,
provenances: [
{
filename: { status: 'exact', entries: [] },
content: { status: 'exact', entries: [] },
tags: [
{
tagName: 'private-name',
provenance: {
status: 'exact',
entries: [
{
name: 'TAG_VALUE',
encryptedValue: 'encrypted-tag-value',
sourceUserId: 'user-1',
sourceWorkspaceId: 'workspace-1',
},
],
},
},
],
},
],
})
})
})
apps/sim/app/api/knowledge/secret-provenance.ts
+15 -18
View File
@@ -23,7 +23,6 @@ import {
import {
knowledgeDocumentContentSelectionKey,
knowledgeDocumentFilenameSelectionKey,
knowledgeDocumentTagNameSelectionKey,
knowledgeDocumentTagValueSelectionKey,
parseKnowledgeDocumentTagProvenanceTargets,
} from '@/lib/knowledge/secret-provenance-selection'
@@ -59,11 +58,16 @@ export function resolveKnowledgeWriteSecretProvenance(options: {
if (inspection.status !== 'verified' || options.authType !== AuthType.INTERNAL_JWT) {
return { success: false, response: invalidKnowledgeProvenanceResponse() }
}
if (
!isPrivateSecretProvenanceBundleV1(inspection.value) ||
!inspection.value.complete ||
inspection.value.selections.length !== options.selectionKeys.length
) {
if (!isPrivateSecretProvenanceBundleV1(inspection.value)) {
return { success: false, response: invalidKnowledgeProvenanceResponse() }
}
if (!inspection.value.complete) {
return {
success: true,
provenances: options.selectionKeys.map(() => ({ status: 'unknown' })),
}
}
if (inspection.value.selections.length !== options.selectionKeys.length) {
return { success: false, response: invalidKnowledgeProvenanceResponse() }
}
const provenances = options.selectionKeys.map((selectionKey) =>
@@ -72,9 +76,7 @@ export function resolveKnowledgeWriteSecretProvenance(options: {
...(options.workspaceId ? { workspaceId: options.workspaceId } : {}),
})
)
if (
provenances.some((provenance) => provenance === undefined || provenance.status === 'unknown')
) {
if (provenances.some((provenance) => provenance === undefined)) {
return { success: false, response: invalidKnowledgeProvenanceResponse() }
}
return { success: true, provenances: provenances as DurableSecretProvenance[] }
@@ -84,7 +86,7 @@ type KnowledgeDocumentWriteProvenanceResolution =
| { success: true; provenances?: KnowledgeDocumentWriteSecretProvenance[] }
| { success: false; response: NextResponse }
/** Resolves field-separated document input provenance and rejects dynamic secret tag names. */
/** Resolves provenance for durable document fields; persisted tag names remain raw and untracked. */
export function resolveKnowledgeDocumentWriteSecretProvenance(options: {
request: NextRequest
payload: unknown
@@ -99,10 +101,9 @@ export function resolveKnowledgeDocumentWriteSecretProvenance(options: {
const selectionKeys = options.documents.flatMap((_document, documentIndex) => [
knowledgeDocumentFilenameSelectionKey(documentIndex),
knowledgeDocumentContentSelectionKey(documentIndex),
...tagTargets[documentIndex].flatMap((_tag, tagIndex) => [
knowledgeDocumentTagNameSelectionKey(documentIndex, tagIndex),
knowledgeDocumentTagValueSelectionKey(documentIndex, tagIndex),
]),
...tagTargets[documentIndex].map((_tag, tagIndex) =>
knowledgeDocumentTagValueSelectionKey(documentIndex, tagIndex)
),
])
const resolved = resolveKnowledgeWriteSecretProvenance({
request: options.request,
@@ -122,11 +123,7 @@ export function resolveKnowledgeDocumentWriteSecretProvenance(options: {
const content = resolved.provenances[provenanceIndex++]
const tagProvenances: KnowledgeDocumentWriteSecretProvenance['tags'][number][] = []
for (const tag of tags) {
const tagName = resolved.provenances[provenanceIndex++]
const tagValue = resolved.provenances[provenanceIndex++]
if (tagName.status !== 'exact' || tagName.entries.length > 0) {
return { success: false, response: invalidKnowledgeProvenanceResponse() }
}
tagProvenances.push({ tagName: tag.tagName, provenance: tagValue })
}
provenances.push({ filename, content, tags: tagProvenances })
apps/sim/app/api/mcp/serve/[serverId]/route.test.ts
+3 -4
View File
@@ -1008,7 +1008,7 @@ describe('MCP Serve Route', () => {
expect(JSON.stringify(body)).not.toContain(RESOLVED_SECRET_PROVENANCE_FIELD)
})
it('fails closed when a successful workflow MCP response omits private provenance', async () => {
it('preserves a successful legacy workflow MCP response without private provenance', async () => {
dbChainMockFns.limit
.mockResolvedValueOnce([
{
@@ -1041,9 +1041,8 @@ describe('MCP Serve Route', () => {
const response = await POST(req, { params: Promise.resolve({ serverId: 'server-1' }) })
const body = await response.json()
expect(response.status).toBe(500)
expect(body.error.message).toBe('Tool execution failed')
expect(JSON.stringify(body)).not.toContain('secret-value')
expect(response.status).toBe(200)
expect(body.result.content[0].text).toBe('"secret-value"')
})
it.each([
apps/sim/app/api/mcp/serve/[serverId]/route.ts
+9 -19
View File
@@ -256,8 +256,7 @@ interface WorkflowExecutionProvenance {
async function consumeWorkflowExecutionProvenance(
response: Response,
value: unknown,
scope: { userId: string; workspaceId: string }
value: unknown
): Promise<WorkflowExecutionProvenance> {
const inspection = inspectPrivateToolMetadataEnvelope(
response.headers,
@@ -265,8 +264,7 @@ async function consumeWorkflowExecutionProvenance(
RESOLVED_SECRET_PROVENANCE_METADATA_V1
)
if (inspection.status === 'unsupported') {
if (!response.ok) return { value, hasPrivateProvenance: false }
throw new Error('MCP workflow execution provenance is unavailable')
return { value, hasPrivateProvenance: false }
}
if (inspection.status === 'invalid' || !isJsonObject(value)) {
throw new Error('MCP workflow execution provenance is invalid')
@@ -912,10 +910,7 @@ async function handleToolsCall(
})
const rawExecuteResult = await readWorkflowExecutionResult(response, abortSignal.signal)
const provenance = await consumeWorkflowExecutionProvenance(response, rawExecuteResult, {
userId: actorUserId,
workspaceId: wf.workspaceId,
})
const provenance = await consumeWorkflowExecutionProvenance(response, rawExecuteResult)
const executeResult = provenance.value
const executeResultObject = isJsonObject(executeResult) ? executeResult : null
@@ -959,17 +954,12 @@ async function handleToolsCall(
: executeResultObject && hasResponseField(executeResultObject, 'output')
? executeResultObject.output
: executeResult
if (!provenance.hasPrivateProvenance) {
throw new Error('MCP workflow execution provenance is unavailable')
}
const projectedToolOutput = await projectWorkflowMcpModelContent(
toolOutput,
provenance.privateProvenance,
{
userId: actorUserId,
workspaceId: wf.workspaceId,
}
)
const projectedToolOutput = provenance.hasPrivateProvenance
? await projectWorkflowMcpModelContent(toolOutput, provenance.privateProvenance, {
userId: actorUserId,
workspaceId: wf.workspaceId,
})
: toolOutput
const result: CallToolResult = {
content: [{ type: 'text', text: serializeToolText(projectedToolOutput) }],
isError: executeResultObject?.success === false,
apps/sim/app/api/mcp/tools/execute/route.test.ts
+59 -40
View File
@@ -9,17 +9,11 @@ const {
mockDiscoverServerTools,
mockExecuteTool,
mockGetExecutionTimeout,
mockReadResponseToBufferWithLimit,
} = vi.hoisted(() => ({
mockCapExecutionTimeoutMs: vi.fn((_policy: number, requested?: number) => requested ?? 0),
mockDiscoverServerTools: vi.fn(),
mockExecuteTool: vi.fn(),
mockGetExecutionTimeout: vi.fn(() => 0),
mockReadResponseToBufferWithLimit: vi.fn(),
}))
vi.mock('@/lib/core/utils/stream-limits', () => ({
readResponseToBufferWithLimit: mockReadResponseToBufferWithLimit,
}))
vi.mock('@/lib/mcp/middleware', () => ({
@@ -102,47 +96,26 @@ describe('MCP tool execution private secret provenance', () => {
vi.clearAllMocks()
mockDiscoverServerTools.mockResolvedValue([{ name: 'example_tool', inputSchema: {} }])
mockExecuteTool.mockResolvedValue({ content: [{ type: 'text', text: 'ok' }] })
mockReadResponseToBufferWithLimit.mockImplementation(async (response: Response) =>
Buffer.from(await response.arrayBuffer())
)
})
it('returns fail-closed scoped provenance only to an authenticated internal caller', async () => {
it('returns provenance activated by this MCP transport call', async () => {
mockDiscoverServerTools.mockImplementationOnce(
async (
_userId: string,
_serverId: string,
_workspaceId: string,
_forceRefresh: boolean,
report: (value: unknown) => void
recordProvenance?: (provenance: unknown) => void
) => {
report({
recordProvenance?.({
version: 1,
complete: false,
entries: [],
complete: true,
entries: [{ name: 'MCP_TOKEN', encryptedValue: 'encrypted-mcp-token' }],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
})
return [{ name: 'example_tool', inputSchema: {} }]
}
)
mockExecuteTool.mockImplementationOnce(
async (
_userId: string,
_serverId: string,
_toolCall: unknown,
_workspaceId: string,
_headers: unknown,
report: (value: unknown) => void
) => {
report({
version: 1,
complete: true,
entries: [{ name: 'NEW_TOKEN', encryptedValue: 'encrypted-v2' }],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
})
return { content: [{ type: 'text', text: 'ok' }] }
}
)
const request = createRequest({
'x-sim-request-private-tool-metadata': 'resolved-secret-provenance-v1',
})
@@ -155,10 +128,12 @@ describe('MCP tool execution private secret provenance', () => {
)
expect(body.__resolvedSecretTraceProvenance).toEqual({
version: 1,
complete: false,
entries: [],
complete: true,
entries: [{ name: 'MCP_TOKEN', encryptedValue: 'encrypted-mcp-token' }],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
})
expect(mockDiscoverServerTools.mock.calls[0]?.[4]).toEqual(expect.any(Function))
expect(mockExecuteTool.mock.calls[0]?.[5]).toEqual(expect.any(Function))
})
it('does not expose private provenance metadata to a session caller', async () => {
@@ -176,10 +151,10 @@ describe('MCP tool execution private secret provenance', () => {
expect(mockExecuteTool.mock.calls[0]?.[5]).toBeUndefined()
})
it('preserves the functional response when private provenance cannot be attached', async () => {
mockReadResponseToBufferWithLimit.mockRejectedValueOnce(new Error('Response exceeds limit'))
it('preserves MCP error status and message when attaching private provenance', async () => {
mockExecuteTool.mockResolvedValueOnce({
content: [{ type: 'text', text: 'unchanged' }],
isError: true,
content: [{ type: 'text', text: 'Provider rejected the request' }],
})
const request = createRequest({
'x-sim-request-private-tool-metadata': 'resolved-secret-provenance-v1',
@@ -188,10 +163,54 @@ describe('MCP tool execution private secret provenance', () => {
const response = await POST(request, {})
const body = (await response.json()) as Record<string, unknown>
expect(response.status).toBe(400)
expect(response.headers.get('x-sim-private-tool-metadata')).toBe(
'resolved-secret-provenance-v1'
)
expect(body).toMatchObject({
success: false,
error: 'Provider rejected the request',
__resolvedSecretTraceProvenance: {
version: 1,
complete: true,
entries: [],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
},
})
})
it('attaches private provenance without imposing a second functional response limit', async () => {
const largeText = 'x'.repeat(10 * 1024 * 1024 + 1)
mockExecuteTool.mockResolvedValueOnce({
content: [{ type: 'text', text: largeText }],
})
const request = createRequest({
'x-sim-request-private-tool-metadata': 'resolved-secret-provenance-v1',
})
const response = await (async () => {
const responseJsonSpy = vi.spyOn(Response.prototype, 'json')
try {
const result = await POST(request, {})
expect(responseJsonSpy).not.toHaveBeenCalled()
return result
} finally {
responseJsonSpy.mockRestore()
}
})()
const body = (await response.json()) as Record<string, unknown>
expect(response.status).toBe(200)
expect(response.ok).toBe(true)
expect(response.headers.has('x-sim-private-tool-metadata')).toBe(false)
expect(body).not.toHaveProperty('__resolvedSecretTraceProvenance')
expect(response.headers.get('x-sim-private-tool-metadata')).toBe(
'resolved-secret-provenance-v1'
)
expect(body.__resolvedSecretTraceProvenance).toEqual({
version: 1,
complete: true,
entries: [],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
})
expect(body).toMatchObject({
success: true,
data: {
@@ -201,7 +220,7 @@ describe('MCP tool execution private secret provenance', () => {
})
expect(
(body.data as { output: { content: Array<{ text?: unknown }> } }).output.content[0]?.text
).toBe('unchanged')
).toBe(largeText)
})
it('uses the remaining workflow deadline for trusted internal tool calls', async () => {
apps/sim/app/api/mcp/tools/execute/route.ts
+34 -55
View File
@@ -11,15 +11,13 @@ import {
} from '@/lib/billing/core/billing-attribution'
import { capExecutionTimeoutMs, getExecutionTimeout } from '@/lib/core/execution-limits'
import type { SubscriptionPlan } from '@/lib/core/rate-limiter/types'
import { readResponseToBufferWithLimit } from '@/lib/core/utils/stream-limits'
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
import { SIM_VIA_HEADER } from '@/lib/execution/call-chain'
import { parseRemainingExecutionDeadlineMs } from '@/lib/execution/execution-deadline-header'
import {
PRIVATE_TOOL_METADATA_RESPONSE_HEADER,
RESOLVED_SECRET_PROVENANCE_FIELD,
RESOLVED_SECRET_PROVENANCE_METADATA_V1,
requestsPrivateToolMetadata,
serializePrivateToolMetadataResponseEnvelope,
} from '@/lib/execution/private-tool-metadata'
import {
mcpBodyReadErrorResponse,
@@ -34,7 +32,7 @@ import {
type McpToolCall,
type McpToolResult,
} from '@/lib/mcp/types'
import { categorizeError, createMcpErrorResponse, createMcpSuccessResponse } from '@/lib/mcp/utils'
import { categorizeError } from '@/lib/mcp/utils'
import {
assertPermissionsAllowed,
McpToolsNotAllowedError,
@@ -45,7 +43,6 @@ import {
} from '@/executor/utils/resolved-secret-trace-registry'
const logger = createLogger('McpToolExecutionAPI')
const MAX_PRIVATE_MCP_RESPONSE_BYTES = 10 * 1024 * 1024
export const dynamic = 'force-dynamic'
@@ -68,33 +65,21 @@ function hasType(prop: unknown): prop is SchemaProperty {
return typeof prop === 'object' && prop !== null && 'type' in prop
}
async function attachPrivateProvenance(
response: NextResponse,
provenance: ResolvedSecretTraceProvenanceAccumulator
): Promise<NextResponse> {
let payload: Record<string, unknown>
try {
const body = await readResponseToBufferWithLimit(response.clone(), {
maxBytes: MAX_PRIVATE_MCP_RESPONSE_BYTES,
label: 'MCP private metadata response',
allowNoBodyFallback: true,
})
const parsed: unknown = JSON.parse(body.toString('utf8'))
if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) {
throw new Error('MCP response is not a JSON object')
}
payload = parsed as Record<string, unknown>
} catch {
return response
function createToolExecutionResponse(
body: Record<string, unknown>,
status: number,
provenance: ResolvedSecretTraceProvenanceAccumulator | undefined
): NextResponse {
if (!provenance) {
return NextResponse.json(body, { status })
}
const headers = new Headers(response.headers)
headers.delete('content-length')
headers.set(PRIVATE_TOOL_METADATA_RESPONSE_HEADER, RESOLVED_SECRET_PROVENANCE_METADATA_V1)
return NextResponse.json(
{ ...payload, [RESOLVED_SECRET_PROVENANCE_FIELD]: provenance.exportProvenance() },
{ status: response.status, headers }
const envelope = serializePrivateToolMetadataResponseEnvelope(
body,
RESOLVED_SECRET_PROVENANCE_METADATA_V1,
provenance.exportProvenance()
)
return NextResponse.json(envelope.body, { status, headers: envelope.headers })
}
/**
@@ -115,13 +100,22 @@ export const POST = withRouteHandler(
resolvedSecretTraceProvenance.record(provenance)
}
: undefined
const response = await (async (): Promise<NextResponse> => {
const errorResponse = (message: string, status: number): NextResponse =>
createToolExecutionResponse(
{ success: false, error: message },
status,
resolvedSecretTraceProvenance
)
const successResponse = <T>(data: T, status = 200): NextResponse =>
createToolExecutionResponse({ success: true, data }, status, resolvedSecretTraceProvenance)
return (async (): Promise<NextResponse> => {
try {
const rawBody = await readMcpJsonBodyWithLimit(request)
const parsedBody = mcpToolExecutionBodySchema.safeParse(rawBody)
if (!parsedBody.success) {
return createMcpErrorResponse(parsedBody.error, 'Invalid request format', 400)
return errorResponse('Invalid request format', 400)
}
const body = parsedBody.data
@@ -148,7 +142,7 @@ export const POST = withRouteHandler(
})
} catch (err) {
if (err instanceof McpToolsNotAllowedError) {
return createMcpErrorResponse(err, err.message, 403)
return errorResponse(err.message, 403)
}
throw err
}
@@ -172,11 +166,7 @@ export const POST = withRouteHandler(
logger.warn(`[${requestId}] Tool ${toolName} not found on server ${serverId}`, {
availableTools: tools.map((t) => t.name),
})
return createMcpErrorResponse(
new Error('Tool not found'),
'Tool not found on the specified server',
404
)
return errorResponse('Tool not found on the specified server', 404)
}
if (tool.inputSchema?.properties) {
@@ -241,11 +231,7 @@ export const POST = withRouteHandler(
const validationError = validateToolArguments(tool, args)
if (validationError) {
logger.warn(`[${requestId}] Tool validation failed: ${validationError}`)
return createMcpErrorResponse(
new Error(`Invalid arguments for tool ${toolName}: ${validationError}`),
'Invalid tool arguments',
400
)
return errorResponse('Invalid tool arguments', 400)
}
}
@@ -317,11 +303,7 @@ export const POST = withRouteHandler(
logger.warn(
`[${requestId}] Tool execution returned error for ${toolName} on ${serverId}`
)
return createMcpErrorResponse(
transformedResult,
transformedResult.error || 'Tool execution failed',
400
)
return errorResponse(transformedResult.error || 'Tool execution failed', 400)
}
logger.info(`[${requestId}] Successfully executed tool ${toolName} on server ${serverId}`)
@@ -342,7 +324,7 @@ export const POST = withRouteHandler(
})
}
return createMcpSuccessResponse(transformedResult)
return successResponse(transformedResult)
} catch (error) {
if (getErrorMessage(error) === 'Tool execution timeout') {
resolvedSecretTraceProvenance?.markIncomplete()
@@ -359,27 +341,24 @@ export const POST = withRouteHandler(
logger.warn(`[${requestId}] OAuth re-authorization required for MCP tool execution`, {
serverId: errorServerId,
})
return NextResponse.json(
return createToolExecutionResponse(
{
success: false,
error: 'OAuth re-authorization required',
code: 'reauth_required',
serverId: errorServerId,
},
{ status: 401 }
401,
resolvedSecretTraceProvenance
)
}
logger.error(`[${requestId}] Error executing MCP tool:`, error)
const { message, status } = categorizeError(error)
return createMcpErrorResponse(new Error(message), message, status)
return errorResponse(message, status)
}
})()
return resolvedSecretTraceProvenance
? attachPrivateProvenance(response, resolvedSecretTraceProvenance)
: response
}
)
)
apps/sim/app/api/memory/secret-provenance.test.ts
+33 -4
View File
@@ -77,14 +77,19 @@ describe('memory write secret provenance', () => {
expect(result).toEqual({ success: true })
})
it('rejects an unavailable verified selection before persistence', () => {
it('persists authenticated unavailable selection lineage as unknown', () => {
const bundle = {
version: 1 as const,
complete: true,
selections: [
{
key: 'data',
provenance: { version: 1 as const, complete: false, entries: [] },
provenance: {
version: 1 as const,
complete: false,
entries: [],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
},
},
],
}
@@ -103,8 +108,32 @@ describe('memory write secret provenance', () => {
workspaceId: 'workspace-1',
})
expect(result.success).toBe(false)
if (!result.success) expect(result.response.status).toBe(400)
expect(result).toEqual({ success: true, provenance: { status: 'unknown' } })
})
it('persists an authenticated incomplete bundle as unknown', () => {
const payload = {
[PRIVATE_SECRET_PROVENANCE_FIELD]: {
version: 1 as const,
complete: false,
selections: [],
},
}
const request = new NextRequest('http://localhost/api/memory', {
method: 'POST',
headers: { [PRIVATE_SECRET_PROVENANCE_HEADER]: PRIVATE_SECRET_PROVENANCE_BUNDLE_V1 },
body: JSON.stringify(payload),
})
expect(
resolveMemoryWriteSecretProvenance({
request,
payload,
authType: AuthType.INTERNAL_JWT,
userId: 'user-1',
workspaceId: 'workspace-1',
})
).toEqual({ success: true, provenance: { status: 'unknown' } })
})
it('accepts exact-empty provenance from the workflow owner in the actor workspace', () => {
apps/sim/app/api/memory/secret-provenance.ts
+8 -6
View File
@@ -57,18 +57,20 @@ export function resolveMemoryWriteSecretProvenance(options: {
if (inspection.status !== 'verified' || options.authType !== AuthType.INTERNAL_JWT) {
return { success: false, response: invalidMemoryProvenanceResponse() }
}
if (
!isPrivateSecretProvenanceBundleV1(inspection.value) ||
!inspection.value.complete ||
inspection.value.selections.length !== 1
) {
if (!isPrivateSecretProvenanceBundleV1(inspection.value)) {
return { success: false, response: invalidMemoryProvenanceResponse() }
}
if (!inspection.value.complete) {
return { success: true, provenance: { status: 'unknown' } }
}
if (inspection.value.selections.length !== 1) {
return { success: false, response: invalidMemoryProvenanceResponse() }
}
const provenance = durableSecretProvenanceFromPrivateBundle(inspection.value, 'data', {
userId: options.userId,
workspaceId: options.workspaceId,
})
return provenance?.status === 'exact'
return provenance
? { success: true, provenance }
: { success: false, response: invalidMemoryProvenanceResponse() }
}
apps/sim/app/api/mothership/execute/route.test.ts
+22 -93
View File
@@ -275,12 +275,7 @@ describe('mothership private trace provenance transport', () => {
)
expect(response.status).toBe(200)
expect(mockBuildTaggedMcpToolSchemas).toHaveBeenCalledWith(
'user-1',
'workspace-1',
['123'],
expect.any(Function)
)
expect(mockBuildTaggedMcpToolSchemas).toHaveBeenCalledWith('user-1', 'workspace-1', ['123'])
expect(mockProcessContextsServer).toHaveBeenCalledWith(
[
{
@@ -415,7 +410,7 @@ describe('mothership private trace provenance transport', () => {
})
})
it('returns encrypted provenance on a marker-gated successful request', async () => {
it('returns exact-empty output provenance on a marker-gated successful request', async () => {
mockRunHeadlessCopilotLifecycle.mockImplementation(
async (_payload: Record<string, unknown>, options: CopilotLifecycleOptions) => {
expect(options.environmentContext).not.toHaveProperty('decryptedEnvVars')
@@ -447,48 +442,31 @@ describe('mothership private trace provenance transport', () => {
expect(body.__resolvedSecretTraceProvenance).toEqual({
version: 1,
complete: true,
entries: [{ name: 'API_KEY', encryptedValue: 'encrypted-secret' }],
entries: [],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
})
expect(JSON.stringify(body.__resolvedSecretTraceProvenance)).not.toContain('secret-value')
expect(mockGetPersonalAndWorkspaceEnv).toHaveBeenCalledTimes(1)
})
it('imports only MCP provenance present in the discovered schemas', async () => {
const provenance = {
version: 1,
complete: true,
entries: [
{ name: 'API_KEY', encryptedValue: 'encrypted-secret' },
{ name: 'UNRELATED', encryptedValue: 'encrypted-unrelated' },
],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
}
mockDecryptSecret.mockImplementation(async (encryptedValue: string) => ({
decrypted: encryptedValue === 'encrypted-secret' ? 'secret-value' : 'unrelated-value',
}))
mockBuildTaggedMcpToolSchemas.mockImplementationOnce(
async (
_userId: string,
_workspaceId: string,
_serverIds: string[],
report: (value: unknown) => void
) => {
report(provenance)
return [{ name: 'mcp-docs', description: 'Uses secret-value' }]
}
)
it('keeps discovered MCP schemas raw without activating matching configured secrets', async () => {
mockBuildTaggedMcpToolSchemas.mockResolvedValueOnce([
{ name: 'mcp-docs', description: 'Uses secret-value' },
])
mockRunHeadlessCopilotLifecycle.mockImplementation(
async (payload: Record<string, unknown>, options: CopilotLifecycleOptions) => {
const registry =
options.environmentContext?.resolvedSecretTraceRegistry ??
options.resolvedSecretTraceRegistry
expect(registry?.exportProvenance()).toEqual({
...provenance,
entries: [{ name: 'API_KEY', encryptedValue: 'encrypted-secret' }],
version: 1,
complete: true,
entries: [],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
})
expect(JSON.stringify(payload)).not.toContain('encrypted-secret')
expect(JSON.stringify(payload)).not.toContain('__resolvedSecretTraceProvenance')
expect(payload.mothershipTools).toEqual([
{ name: 'mcp-docs', description: 'Uses secret-value' },
])
return successResult()
}
)
@@ -514,62 +492,15 @@ describe('mothership private trace provenance transport', () => {
expect({ status: response.status, provenance: body.__resolvedSecretTraceProvenance }).toEqual({
status: 200,
provenance: {
...provenance,
entries: [{ name: 'API_KEY', encryptedValue: 'encrypted-secret' }],
version: 1,
complete: true,
entries: [],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
},
})
})
it('omits MCP tools with malformed discovery provenance without poisoning the lifecycle', async () => {
mockBuildTaggedMcpToolSchemas.mockImplementationOnce(
async (
_userId: string,
_workspaceId: string,
_serverIds: string[],
report: (value: unknown) => void
) => {
report({ version: 1, complete: true, entries: 'invalid' })
return []
}
)
mockRunHeadlessCopilotLifecycle.mockImplementation(
async (payload: Record<string, unknown>, options: CopilotLifecycleOptions) => {
const registry =
options.environmentContext?.resolvedSecretTraceRegistry ??
options.resolvedSecretTraceRegistry
expect(registry?.isComplete()).toBe(true)
expect(payload).not.toHaveProperty('mothershipTools')
return successResult()
}
)
const response = await POST(
createMockRequest(
'POST',
{
...requestBody,
contexts: [{ kind: 'mcp', label: 'Docs', serverId: 'server-1' }],
},
{
Authorization: 'Bearer internal',
'x-sim-billing-attribution': 'billing',
'x-sim-request-private-tool-metadata': 'resolved-secret-provenance-v1',
},
'http://localhost:3000/api/mothership/execute'
)
)
const body = await response.json()
expect(response.status).toBe(200)
expect(body.__resolvedSecretTraceProvenance).toEqual({
version: 1,
complete: true,
entries: [],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
})
})
it('returns encrypted provenance with marker-gated failures', async () => {
it('returns exact-empty provenance for an already projected marker-gated failure', async () => {
mockRunHeadlessCopilotLifecycle.mockImplementation(
async (_payload: Record<string, unknown>, options: CopilotLifecycleOptions) => {
activateSecret(options)
@@ -601,12 +532,10 @@ describe('mothership private trace provenance transport', () => {
'resolved-secret-provenance-v1'
)
expect(body.content).toBe('secret-value')
expect(body.__resolvedSecretTraceProvenance.entries).toEqual([
{ name: 'API_KEY', encryptedValue: 'encrypted-secret' },
])
expect(body.__resolvedSecretTraceProvenance.entries).toEqual([])
})
it('places encrypted provenance only on the terminal streamed event', async () => {
it('places exact-empty provenance only on the terminal streamed event', async () => {
mockRunHeadlessCopilotLifecycle.mockImplementation(
async (_payload: Record<string, unknown>, options: CopilotLifecycleOptions) => {
activateSecret(options)
@@ -642,7 +571,7 @@ describe('mothership private trace provenance transport', () => {
data: {
content: 'secret-value',
__resolvedSecretTraceProvenance: {
entries: [{ name: 'API_KEY', encryptedValue: 'encrypted-secret' }],
entries: [],
},
},
})
apps/sim/app/api/mothership/execute/route.ts
+7 -38
View File
@@ -38,7 +38,6 @@ import {
} from '@/lib/workspaces/permissions/utils'
import {
createIncompleteResolvedSecretTraceRegistry,
ResolvedSecretTraceProvenanceAccumulator,
type ResolvedSecretTraceRegistry,
} from '@/executor/utils/resolved-secret-trace-registry'
import type { ChatContext } from '@/stores/panel'
@@ -60,7 +59,9 @@ function withPrivateProvenance<T extends Record<string, unknown>>(
return {
...payload,
...(include && registry
? { [RESOLVED_SECRET_PROVENANCE_FIELD]: registry.exportProvenance() }
? {
[RESOLVED_SECRET_PROVENANCE_FIELD]: registry.exportCommittedProvenanceForInputPaths([]),
}
: {}),
}
}
@@ -201,14 +202,6 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
activeResolvedSecretTraceRegistry = createIncompleteResolvedSecretTraceRegistry(scope)
}
resolvedSecretTraceRegistry = activeResolvedSecretTraceRegistry
const mcpDiscoveryProvenance = new ResolvedSecretTraceProvenanceAccumulator({
userId,
workspaceId,
})
const recordMcpDiscoveryProvenance = (provenance: unknown): void => {
mcpDiscoveryProvenance.record(provenance)
}
const effectiveChatId = chatId || generateId()
messageId = providedMessageId || generateId()
requestId = providedRequestId || generateId()
@@ -227,39 +220,15 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
const nonMcpAgentMentions = agentMentions?.filter((context) => context.kind !== 'mcp')
const userPermission = workspaceAccess.permission
const mothershipToolsPromise = Promise.allSettled([
buildSelectedMcpToolSchemas(
userId,
workspaceId,
mcpTools ?? [],
recordMcpDiscoveryProvenance
),
buildTaggedMcpToolSchemas(
userId,
workspaceId,
taggedMcpServerIds,
recordMcpDiscoveryProvenance
),
]).then(async (results) => {
buildSelectedMcpToolSchemas(userId, workspaceId, mcpTools ?? []),
buildTaggedMcpToolSchemas(userId, workspaceId, taggedMcpServerIds),
]).then((results) => {
const groups = results.map((result) => {
if (result.status === 'rejected') throw result.reason
return result.value
})
const byName = new Map(groups.flat().map((tool) => [tool.name, tool]))
const tools = [...byName.values()]
if (activeResolvedSecretTraceRegistry) {
const discoveryRegistry = activeResolvedSecretTraceRegistry.forkForToolInput(tools)
const imported = await discoveryRegistry.importProvenanceForValue(
mcpDiscoveryProvenance.exportProvenance(),
tools,
{ trusted: true }
)
if (!imported || !discoveryRegistry.isComplete()) {
reqLogger.warn('Omitting MCP tools with unverifiable secret provenance')
return []
}
activeResolvedSecretTraceRegistry.mergeToolCallRegistry(discoveryRegistry)
}
return tools
return [...byName.values()]
})
const [workspaceContext, integrationTools, mothershipTools, entitlements, agentContexts] =
await Promise.all([
apps/sim/app/api/providers/route.test.ts
+101 -11
View File
@@ -3,6 +3,10 @@
*/
import { createMockRequest, hybridAuthMockFns } from '@sim/testing'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import {
PRIVATE_MODEL_INPUT_STATE_HEADER,
PROJECTED_MODEL_INPUT_PATHS_V1,
} from '@/lib/execution/model-input-provenance'
const {
mockExecuteProviderRequest,
@@ -10,18 +14,18 @@ const {
mockCheckWorkspaceAccess,
mockAuthorizeCredentialUse,
mockPrepareCopilotEnvironmentContext,
mockCollectProviderModelInputProvenanceValues,
mockImportProvenance,
mockRegistryIsComplete,
mockProjectResolvedSecretModelContent,
} = vi.hoisted(() => ({
mockExecuteProviderRequest: vi.fn(),
mockRequireBillingAttributionHeader: vi.fn(),
mockCheckWorkspaceAccess: vi.fn(),
mockAuthorizeCredentialUse: vi.fn(),
mockPrepareCopilotEnvironmentContext: vi.fn(),
mockCollectProviderModelInputProvenanceValues: vi.fn(),
mockImportProvenance: vi.fn(),
mockRegistryIsComplete: vi.fn(),
mockProjectResolvedSecretModelContent: vi.fn(),
}))
vi.mock('@/providers', () => ({
@@ -45,8 +49,8 @@ vi.mock('@/lib/copilot/environment-context', () => ({
prepareCopilotEnvironmentContext: mockPrepareCopilotEnvironmentContext,
}))
vi.mock('@/providers/model-input-provenance', () => ({
collectProviderModelInputProvenanceValues: mockCollectProviderModelInputProvenanceValues,
vi.mock('@/executor/utils/resolved-secret-content-projection', () => ({
projectResolvedSecretModelContent: mockProjectResolvedSecretModelContent,
}))
vi.mock('@/app/api/auth/oauth/utils', () => ({
@@ -89,6 +93,7 @@ function createProviderRequest(
},
{
'x-sim-private-model-input-provenance': 'resolved-secret-provenance-v1',
[PRIVATE_MODEL_INPUT_STATE_HEADER]: PROJECTED_MODEL_INPUT_PATHS_V1,
...headers,
}
)
@@ -109,12 +114,15 @@ describe('POST /api/providers', () => {
model: 'gpt-4o',
tokens: { input: 1, output: 1, total: 2 },
})
mockCollectProviderModelInputProvenanceValues.mockReturnValue(['selected-model-input'])
mockImportProvenance.mockResolvedValue(true)
mockRegistryIsComplete.mockReturnValue(true)
mockProjectResolvedSecretModelContent.mockImplementation((value) => ({
safe: true,
value,
}))
mockPrepareCopilotEnvironmentContext.mockResolvedValue({
resolvedSecretTraceRegistry: {
importProvenanceForValue: mockImportProvenance,
importProvenance: mockImportProvenance,
isComplete: mockRegistryIsComplete,
},
})
@@ -200,9 +208,91 @@ describe('POST /api/providers', () => {
)
expect(res.status).toBe(200)
expect(mockImportProvenance).toHaveBeenCalledWith(provenance, expect.any(Array), {
trusted: true,
expect(mockImportProvenance).toHaveBeenCalledWith(provenance, { trusted: true })
})
it('projects legacy private prompt provenance on the provider-facing copy', async () => {
mockProjectResolvedSecretModelContent.mockReturnValue({
safe: true,
value: {
systemPrompt: 'Use {{TOKEN}} safely',
context: '[{"role":"user","content":"{{TOKEN}}"}]',
},
})
const res = await POST(
createMockRequest(
'POST',
{
provider: 'openai',
model: 'gpt-4o',
workspaceId: 'ws-1',
systemPrompt: 'Use secret-value safely',
context: '[{"role":"user","content":"secret-value"}]',
__resolvedSecretTraceProvenance: { version: 1, complete: true, entries: [] },
},
{ 'x-sim-private-model-input-provenance': 'resolved-secret-provenance-v1' }
)
)
expect(res.status).toBe(200)
expect(mockExecuteProviderRequest).toHaveBeenCalledWith(
'openai',
expect.objectContaining({
systemPrompt: 'Use {{TOKEN}} safely',
context: '[{"role":"user","content":"{{TOKEN}}"}]',
}),
expect.anything()
)
})
it('does not re-project an explicitly projected private request', async () => {
mockProjectResolvedSecretModelContent.mockReturnValue({
safe: true,
value: { systemPrompt: 'Bo{{TOKEN}}', context: undefined },
})
const res = await POST(
createProviderRequest({
provider: 'openai',
model: 'gpt-4o',
workspaceId: 'ws-1',
systemPrompt: 'Box',
})
)
expect(res.status).toBe(200)
expect(mockProjectResolvedSecretModelContent).not.toHaveBeenCalled()
expect(mockExecuteProviderRequest).toHaveBeenCalledWith(
'openai',
expect.objectContaining({ systemPrompt: 'Box' }),
expect.anything()
)
})
it('rejects a projected marker without a private provenance envelope', async () => {
const res = await POST(
createMockRequest(
'POST',
{ provider: 'openai', model: 'gpt-4o', workspaceId: 'ws-1' },
{ [PRIVATE_MODEL_INPUT_STATE_HEADER]: PROJECTED_MODEL_INPUT_PATHS_V1 }
)
)
expect(res.status).toBe(400)
expect(mockExecuteProviderRequest).not.toHaveBeenCalled()
})
it('rejects an unknown private projection marker', async () => {
const res = await POST(
createProviderRequest(
{ provider: 'openai', model: 'gpt-4o', workspaceId: 'ws-1' },
{ [PRIVATE_MODEL_INPUT_STATE_HEADER]: 'unknown-projection' }
)
)
expect(res.status).toBe(400)
expect(mockExecuteProviderRequest).not.toHaveBeenCalled()
})
it('rejects a partial private provenance envelope', async () => {
@@ -218,7 +308,7 @@ describe('POST /api/providers', () => {
expect(mockExecuteProviderRequest).not.toHaveBeenCalled()
})
it('rejects an internal request without the private provenance envelope', async () => {
it('preserves legacy internal requests without the private provenance envelope', async () => {
const res = await POST(
createMockRequest('POST', {
provider: 'openai',
@@ -227,9 +317,9 @@ describe('POST /api/providers', () => {
})
)
expect(res.status).toBe(400)
expect(res.status).toBe(200)
expect(mockImportProvenance).not.toHaveBeenCalled()
expect(mockExecuteProviderRequest).not.toHaveBeenCalled()
expect(mockExecuteProviderRequest).toHaveBeenCalled()
})
it('omits provisional stream output from the execution header', async () => {
apps/sim/app/api/providers/route.ts
+51 -15
View File
@@ -2,6 +2,7 @@ import { db } from '@sim/db'
import { account } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { getErrorMessage, toError } from '@sim/utils/errors'
import { isPlainRecord } from '@sim/utils/object'
import { eq } from 'drizzle-orm'
import { type NextRequest, NextResponse } from 'next/server'
import { executeProviderContract } from '@/lib/api/contracts/providers'
@@ -16,7 +17,10 @@ import {
import { prepareCopilotEnvironmentContext } from '@/lib/copilot/environment-context'
import { generateRequestId } from '@/lib/core/utils/request'
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
import { inspectModelInputProvenanceRequest } from '@/lib/execution/model-input-provenance'
import {
inspectModelInputProjectionState,
inspectModelInputProvenanceRequest,
} from '@/lib/execution/model-input-provenance'
import { checkWorkspaceAccess } from '@/lib/workspaces/permissions/utils'
import {
getServiceAccountToken,
@@ -30,8 +34,8 @@ import {
ProviderNotAllowedError,
} from '@/ee/access-control/utils/permission-check'
import type { StreamingExecution } from '@/executor/types'
import { projectResolvedSecretModelContent } from '@/executor/utils/resolved-secret-content-projection'
import { executeProviderRequest } from '@/providers'
import { collectProviderModelInputProvenanceValues } from '@/providers/model-input-provenance'
import { projectStreamingExecutionToByteStream } from '@/providers/stream-pump'
import type { ProviderRequest } from '@/providers/types'
@@ -224,7 +228,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
hasBillingAttribution: !!billingAttribution,
})
const providerRequest: ProviderRequest = {
let providerRequest: ProviderRequest = {
model,
systemPrompt,
context,
@@ -254,22 +258,54 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
verbosity,
}
const provenanceInspection = inspectModelInputProvenanceRequest(request.headers, body)
if (provenanceInspection.status === 'unsupported') {
return NextResponse.json({ error: 'Model input provenance is unavailable' }, { status: 400 })
}
if (provenanceInspection.status === 'invalid') {
const projectionState = inspectModelInputProjectionState(request.headers)
if (
provenanceInspection.status === 'invalid' ||
projectionState === 'invalid' ||
(projectionState === 'projected' && provenanceInspection.status !== 'verified')
) {
return NextResponse.json({ error: 'Invalid model input provenance' }, { status: 400 })
}
const providerRuntimeContext = await prepareCopilotEnvironmentContext(auth.userId, workspaceId)
const provenanceReady =
await providerRuntimeContext.resolvedSecretTraceRegistry.importProvenanceForValue(
provenanceInspection.value,
collectProviderModelInputProvenanceValues(providerRequest, provider),
{ trusted: true }
)
if (!provenanceReady || !providerRuntimeContext.resolvedSecretTraceRegistry.isComplete()) {
return NextResponse.json({ error: 'Model input provenance is unavailable' }, { status: 400 })
if (provenanceInspection.status === 'verified') {
const provenanceReady =
await providerRuntimeContext.resolvedSecretTraceRegistry.importProvenance(
provenanceInspection.value,
{ trusted: true }
)
if (!provenanceReady || !providerRuntimeContext.resolvedSecretTraceRegistry.isComplete()) {
return NextResponse.json(
{ error: 'Model input provenance is unavailable' },
{ status: 400 }
)
}
if (projectionState === 'unmarked') {
const projection = projectResolvedSecretModelContent(
{ systemPrompt: providerRequest.systemPrompt, context: providerRequest.context },
providerRuntimeContext.resolvedSecretTraceRegistry
)
if (!projection.safe || !isPlainRecord(projection.value)) {
return NextResponse.json(
{ error: 'Model input provenance is unavailable' },
{ status: 400 }
)
}
const projectedSystemPrompt = projection.value.systemPrompt
const projectedContext = projection.value.context
if (
(projectedSystemPrompt !== undefined && typeof projectedSystemPrompt !== 'string') ||
(projectedContext !== undefined && typeof projectedContext !== 'string')
) {
return NextResponse.json({ error: 'Invalid model input provenance' }, { status: 400 })
}
providerRequest = {
...providerRequest,
systemPrompt: projectedSystemPrompt,
context: projectedContext,
}
}
}
const response = await executeProviderRequest(provider, providerRequest, providerRuntimeContext)
apps/sim/app/api/tools/file/manage/route.test.ts
+47 -4
View File
@@ -347,6 +347,39 @@ describe('POST /api/tools/file/manage content provenance', () => {
)
})
it('persists an authenticated file write with unavailable lineage as unknown', async () => {
const response = await POST(
createMockRequest(
'POST',
{
operation: 'write',
workspaceId: 'workspace-1',
fileName: 'new.txt',
content: 'possibly secret',
__privateSecretProvenance: {
version: 1,
complete: false,
selections: [],
},
},
PRIVATE_SECRET_PROVENANCE_HEADER
)
)
expect(response.status).toBe(200)
expect(mockUploadWorkspaceFile).toHaveBeenCalledWith(
'workspace-1',
'user-1',
Buffer.from('possibly secret'),
'new.txt',
'text/plain',
{
folderId: null,
secretProvenance: { status: 'unknown' },
}
)
})
it('atomically binds append provenance to the exact predecessor version', async () => {
const existing = workspaceFile('file-1')
mockResolveWorkspaceFileReference.mockResolvedValue(existing)
@@ -487,7 +520,7 @@ describe('POST /api/tools/file/manage content provenance', () => {
)
})
it('rejects a secret-bearing archive before downloading or extracting it', async () => {
it('extracts a secret-bearing archive with unknown output provenance', async () => {
const zip = new JSZip()
zip.file('child.txt', 'secret-value')
mockDownloadFileFromStorage.mockResolvedValue(
@@ -511,9 +544,19 @@ describe('POST /api/tools/file/manage content provenance', () => {
})
)
expect(response.status).toBe(422)
expect(mockDownloadFileFromStorage).not.toHaveBeenCalled()
expect(mockUploadWorkspaceFile).not.toHaveBeenCalled()
expect(response.status).toBe(200)
expect(mockDownloadFileFromStorage).toHaveBeenCalledTimes(1)
expect(mockUploadWorkspaceFile).toHaveBeenCalledWith(
'workspace-1',
'user-1',
Buffer.from('secret-value'),
'child.txt',
'text/plain',
{
folderId: null,
secretProvenance: { status: 'unknown' },
}
)
})
it('omits source scope when canonical files have different owners', async () => {
apps/sim/app/api/tools/file/manage/route.ts
+21 -20
View File
@@ -330,26 +330,37 @@ function resolveFileMutationSecretProvenance(options: {
if (
inspection.status !== 'verified' ||
options.authType !== AuthType.INTERNAL_JWT ||
!isPrivateSecretProvenanceBundleV1(inspection.value) ||
!inspection.value.complete ||
inspection.value.selections.length !== options.selectionKeys.length
!isPrivateSecretProvenanceBundleV1(inspection.value)
) {
return { success: false, error: 'Invalid file secret provenance' }
}
const destinationScope = { userId: options.userId, workspaceId: options.workspaceId }
const provenanceBySelection = new Map<string, WorkspaceFileSecretProvenance>()
if (!inspection.value.complete) {
for (const selectionKey of options.selectionKeys) {
provenanceBySelection.set(selectionKey, { status: 'unknown' })
}
return { success: true, provenanceBySelection }
}
if (inspection.value.selections.length !== options.selectionKeys.length) {
return { success: false, error: 'Invalid file secret provenance' }
}
const destinationScope = { userId: options.userId, workspaceId: options.workspaceId }
for (const selectionKey of options.selectionKeys) {
const provenance = durableSecretProvenanceFromPrivateBundle(
inspection.value,
selectionKey,
destinationScope
)
if (
!provenance ||
provenance.status === 'unknown' ||
provenance.entries.some((entry) => !entry.name || !entry.sourceUserId)
) {
if (!provenance) {
return { success: false, error: 'Invalid file secret provenance' }
}
if (provenance.status === 'unknown') {
provenanceBySelection.set(selectionKey, provenance)
continue
}
if (provenance.entries.some((entry) => !entry.name || !entry.sourceUserId)) {
return { success: false, error: 'Invalid file secret provenance' }
}
provenanceBySelection.set(selectionKey, {
@@ -383,7 +394,7 @@ function resolveFileWriteSecretProvenance(options: {
})
if (!resolution.success || !resolution.provenanceBySelection) return resolution
const content = resolution.provenanceBySelection.get('content')
if (!content || content.status !== 'exact') {
if (!content) {
return { success: false, error: 'Invalid file secret provenance' }
}
return { success: true, contentProvenance: content }
@@ -1128,16 +1139,6 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
targetOwnerUserId: userId,
sources: canonicalArchiveSource.concat(selectedArchiveSource),
})
if (archiveProvenance.status === 'unknown' || archiveProvenance.entries.length > 0) {
return NextResponse.json(
{
success: false,
error:
'Archive cannot be decompressed because its secret provenance is not exact-empty',
},
{ status: 422 }
)
}
const archiveBuffer = await downloadFileFromStorage(archive, requestId, logger, {
maxBytes: MAX_ARCHIVE_BYTES,
apps/sim/app/api/tools/fireflies/upload-audio/route.ts
-1
View File
@@ -93,7 +93,6 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
headers: request.headers,
payload: body,
isInternalRequest: true,
allowLegacyWithoutEnvelope: true,
})
if (!modelInputProvenance.success) {
return errorResponse(modelInputProvenance.error, modelInputProvenance.status)
apps/sim/app/api/workflows/[id]/execute/route.async.test.ts
+81 -6
View File
@@ -360,6 +360,7 @@ function createInternalProvenanceRequest(
useDraftState?: boolean
provenance?: typeof WORKFLOW_INPUT_PROVENANCE
selectionKey?: string
bundleComplete?: boolean
includeHeader?: boolean
includeField?: boolean
} = {}
@@ -371,6 +372,7 @@ function createInternalProvenanceRequest(
useDraftState,
provenance = WORKFLOW_INPUT_PROVENANCE,
selectionKey = 'input',
bundleComplete = true,
includeHeader = true,
includeField = true,
} = options
@@ -387,8 +389,8 @@ function createInternalProvenanceRequest(
? {
[PRIVATE_SECRET_PROVENANCE_FIELD]: {
version: 1,
complete: true,
selections: [{ key: selectionKey, provenance }],
complete: bundleComplete,
selections: bundleComplete ? [{ key: selectionKey, provenance }] : [],
},
}
: {}),
@@ -496,6 +498,9 @@ describe('workflow execute async route', () => {
close: vi.fn().mockResolvedValue(undefined),
})
loggingSessionMockFns.mockWaitForPostExecution.mockReset().mockResolvedValue(undefined)
loggingSessionMockFns.mockExportResolvedSecretTraceProvenanceForValue
.mockReset()
.mockReturnValue({ version: 1, complete: false, entries: [] })
mockExecuteWorkflowJob.mockReset().mockResolvedValue({ success: true })
encryptionMockFns.mockDecryptSecret.mockReset().mockImplementation(async (value: string) => ({
decrypted: value === 'encrypted-token' ? 'secret-value' : 'other-secret',
@@ -532,6 +537,26 @@ describe('workflow execute async route', () => {
expect(executionOptions.snapshot.input).toEqual({ hello: 'world' })
})
it('runs authenticated incomplete workflow input with incomplete downstream lineage', async () => {
configureExecutionCaller(EXECUTION_CALLERS[4])
const response = await POST(createInternalProvenanceRequest({ bundleComplete: false }), {
params: Promise.resolve({ id: 'workflow-1' }),
})
expect(response.status).toBe(200)
const executionOptions = mockExecuteWorkflowCore.mock.calls[0]?.[0]
expect(executionOptions).toMatchObject({
trustedInitialResolvedSecretTraceProvenance: {
version: 1,
complete: false,
entries: [],
},
})
expect(executionOptions.snapshot.input).toEqual({ input: { token: 'secret-value' } })
expect(executionOptions.snapshot.input).not.toHaveProperty(PRIVATE_SECRET_PROVENANCE_FIELD)
})
it.each([
{ name: 'standard stream', useDraftState: false },
{ name: 'manual event stream', useDraftState: true },
@@ -1955,19 +1980,29 @@ describe('workflow execute async route', () => {
expect(runFromBlock?.sourceSnapshot).not.toHaveProperty('resolvedSecretTraceProvenance')
})
it('returns encrypted resolution provenance only to an authenticated internal tool caller', async () => {
it('exports exact provenance for the final response body to an authenticated internal caller', async () => {
const caller = EXECUTION_CALLERS[4]
configureExecutionCaller(caller)
const provenance = {
const runProvenance = {
version: 1,
complete: true,
entries: [{ name: 'UNRELATED_SECRET', encryptedValue: 'encrypted-unrelated-secret' }],
}
const responseProvenance = {
version: 1,
complete: true,
entries: [{ name: 'CHILD_SECRET', encryptedValue: 'encrypted-child-secret' }],
}
loggingSessionMockFns.mockExportResolvedSecretTraceProvenanceForValue.mockReturnValueOnce(
responseProvenance
)
mockExecuteWorkflowCore.mockResolvedValueOnce({
success: true,
status: 'completed',
output: { ok: true },
executionState: { resolvedSecretTraceProvenance: provenance },
executionState: {
resolvedSecretTraceProvenance: runProvenance,
},
metadata: {
duration: 100,
startTime: '2026-01-01T00:00:00Z',
@@ -1985,8 +2020,48 @@ describe('workflow execute async route', () => {
)
await expect(response.json()).resolves.toMatchObject({
output: { ok: true },
__resolvedSecretTraceProvenance: provenance,
__resolvedSecretTraceProvenance: responseProvenance,
})
expect(
loggingSessionMockFns.mockExportResolvedSecretTraceProvenanceForValue
).toHaveBeenCalledWith(
expect.objectContaining({
success: true,
executionId: 'execution-123',
output: { ok: true },
})
)
})
it('includes a thrown execution error in the exact response-provenance boundary', async () => {
const caller = EXECUTION_CALLERS[4]
configureExecutionCaller(caller)
const responseProvenance = {
version: 1,
complete: true,
entries: [{ name: 'ERROR_SECRET', encryptedValue: 'encrypted-error-secret' }],
}
loggingSessionMockFns.mockExportResolvedSecretTraceProvenanceForValue.mockReturnValueOnce(
responseProvenance
)
mockExecuteWorkflowCore.mockRejectedValueOnce(new Error('resolved error value'))
const request = createCallerExecutionRequest(caller, undefined, 'sync')
request.headers.set('x-sim-request-private-tool-metadata', 'resolved-secret-provenance-v1')
const response = await POST(request, { params: Promise.resolve({ id: 'workflow-1' }) })
const body = await response.json()
expect(response.status).toBe(500)
expect(body).toMatchObject({
success: false,
error: 'resolved error value',
__resolvedSecretTraceProvenance: responseProvenance,
})
expect(
loggingSessionMockFns.mockExportResolvedSecretTraceProvenanceForValue
).toHaveBeenCalledWith(
expect.objectContaining({ success: false, error: 'resolved error value' })
)
})
it('does not expose private provenance metadata to non-internal callers', async () => {
apps/sim/app/api/workflows/[id]/execute/route.ts
+11 -15
View File
@@ -157,12 +157,7 @@ import type {
IterationContext,
SerializableExecutionState,
} from '@/executor/execution/types'
import type {
BlockLog,
ExecutionResult,
NormalizedBlockOutput,
StreamingExecution,
} from '@/executor/types'
import type { BlockLog, NormalizedBlockOutput, StreamingExecution } from '@/executor/types'
import { getExecutionErrorStatus, hasExecutionResult } from '@/executor/utils/errors'
import type { ResolvedSecretTraceProvenanceV1 } from '@/executor/utils/resolved-secret-trace-registry'
import { Serializer } from '@/serializer'
@@ -202,7 +197,7 @@ function createExecutionJsonResponse(
body: Record<string, unknown>,
init: ResponseInit | undefined,
includePrivateProvenance: boolean,
result?: ExecutionResult
loggingSession?: LoggingSession
): NextResponse {
if (!includePrivateProvenance) {
return NextResponse.json(body, init)
@@ -213,11 +208,12 @@ function createExecutionJsonResponse(
return NextResponse.json(
{
...body,
[RESOLVED_SECRET_PROVENANCE_FIELD]: result?.executionState?.resolvedSecretTraceProvenance ?? {
version: 1,
complete: false,
entries: [],
},
[RESOLVED_SECRET_PROVENANCE_FIELD]:
loggingSession?.exportResolvedSecretTraceProvenanceForValue(body) ?? {
version: 1,
complete: false,
entries: [],
},
},
{ ...init, headers }
)
@@ -1583,7 +1579,7 @@ async function handleExecutePost(
},
{ status: 408 },
includePrivateTraceProvenance,
result
loggingSession
)
}
@@ -1654,7 +1650,7 @@ async function handleExecutePost(
filteredResult,
undefined,
includePrivateTraceProvenance,
result
loggingSession
)
} catch (error: unknown) {
const executionTimedOut = didExecutionTimeOut(error)
@@ -1716,7 +1712,7 @@ async function handleExecutePost(
},
{ status },
includePrivateTraceProvenance,
executionResult
loggingSession
)
} finally {
requestAbort.cleanup()
apps/sim/blocks/blocks/knowledge.ts
+15
View File
@@ -1,3 +1,4 @@
import { isPlainRecord } from '@sim/utils/object'
import { PackageSearchIcon } from '@/components/icons'
import { DEFAULT_RERANKER_MODEL, SUPPORTED_RERANKER_MODELS } from '@/lib/knowledge/reranker-models'
import type { BlockConfig } from '@/blocks/types'
@@ -377,6 +378,7 @@ export const KnowledgeBlock: BlockConfig = {
}
},
params: (params) => {
params = { ...params }
const knowledgeBaseId = params.knowledgeBaseId ? String(params.knowledgeBaseId).trim() : ''
if (!knowledgeBaseId) {
throw new Error('Knowledge base ID is required')
@@ -428,6 +430,19 @@ export const KnowledgeBlock: BlockConfig = {
params.documentId = String(params.upsertDocumentId).trim()
}
if (
(params.operation === 'create_document' || params.operation === 'upsert_document') &&
typeof params.documentTags === 'string' &&
params.documentTags.trim().length > 0
) {
try {
const documentTags: unknown = JSON.parse(params.documentTags)
if (Array.isArray(documentTags) || isPlainRecord(documentTags)) {
params.documentTags = documentTags
}
} catch {}
}
// Convert enabled dropdown string to boolean for update_chunk
if (params.operation === 'update_chunk' && typeof params.enabled === 'string') {
params.enabled = params.enabled === 'true'
apps/sim/ee/workspace-forking/lib/copy/cleanup-failed.test.ts
+122 -1
View File
@@ -1,7 +1,13 @@
/**
* @vitest-environment node
*/
import { knowledgeBase, workflow, workflowBlocks, workflowDeploymentVersion } from '@sim/db/schema'
import {
document,
knowledgeBase,
workflow,
workflowBlocks,
workflowDeploymentVersion,
} from '@sim/db/schema'
import { dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing'
import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
@@ -269,6 +275,31 @@ describe('cleanup-failed', () => {
expect(updates()).toHaveLength(0)
expect(mockInvalidateDeployedStateCache).not.toHaveBeenCalled()
})
it('attempts every workflow, then reports a workflow-scoped cleanup failure', async () => {
queueTableRows(workflowDeploymentVersion, [
{ id: 'dv-failed', version: 5, state: versionState('failed-kb') },
])
queueTableRows(workflowDeploymentVersion, [
{ id: 'dv-cleaned', version: 5, state: versionState('failed-kb') },
])
dbChainMockFns.set.mockImplementationOnce(() => {
throw new Error('first workflow update failed')
})
await expect(
clearFailedReferencesInDeploymentVersions(
new Set(['wf-failed', 'wf-cleaned']),
failedByKind(),
'test'
)
).rejects.toThrow('Failed to clear deployment-version references for 1 workflow(s)')
// The second workflow is still processed after the first workflow's update fails.
expect(dbChainMockFns.update).toHaveBeenCalledTimes(2)
expect(mockInvalidateDeployedStateCache).toHaveBeenCalledTimes(1)
expect(mockInvalidateDeployedStateCache).toHaveBeenCalledWith('dv-cleaned')
})
})
describe('clearFailedForkResourceReferences', () => {
@@ -316,6 +347,75 @@ describe('cleanup-failed', () => {
expect(deletes()[0].table).toBe(knowledgeBase)
})
it('keeps a failed copied knowledge base when it contains a non-fork document', async () => {
queueTableRows(knowledgeBase, [{ id: 'failed-kb' }])
const cleaned = await clearFailedForkResourceReferences({
childWorkspaceId: 'child-ws',
failures: [{ kind: 'knowledge-base', childId: 'failed-kb', documentChildIds: [] }],
requestId: 'test',
})
expect(cleaned).toEqual({ cleared: 0, clearingFailed: false })
expect(dbChainMockFns.update).not.toHaveBeenCalled()
expect(dbChainMockFns.delete).not.toHaveBeenCalled()
expect(mockInvalidateDeployedStateCache).not.toHaveBeenCalled()
})
it('clears only failed fork-document references when a user document keeps the KB alive', async () => {
queueTableRows(knowledgeBase, [{ id: 'failed-kb' }])
queueTableRows(workflow, [{ id: 'wf-1' }])
queueTableRows(workflowBlocks, [
{
...draftBlockRow('failed-kb'),
subBlocks: {
...draftBlockRow('failed-kb').subBlocks,
documentId: {
id: 'documentId',
type: 'document-selector',
value: 'fork_document_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa',
},
},
},
])
const cleaned = await clearFailedForkResourceReferences({
childWorkspaceId: 'child-ws',
failures: [
{
kind: 'knowledge-base',
childId: 'failed-kb',
documentChildIds: ['fork_document_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa'],
},
],
requestId: 'test',
})
expect(cleaned).toEqual({ cleared: 1, clearingFailed: false })
const cleared = updates()[0].values.subBlocks as Record<string, { value: unknown }>
expect(cleared.knowledgeBaseId.value).toBe('failed-kb')
expect(cleared.documentId.value).toBe('')
expect(deletes().map(({ table }) => table)).toEqual([document])
})
it('guards the final KB delete against a non-fork document inserted during cleanup', async () => {
queueTableRows(workflow, [])
await clearFailedForkResourceReferences({
childWorkspaceId: 'child-ws',
failures: [{ kind: 'knowledge-base', childId: 'failed-kb', documentChildIds: [] }],
requestId: 'test',
})
const deletePredicate = dbChainMockFns.where.mock.calls.at(-1)?.[0]
expect(deletePredicate).toEqual(
expect.objectContaining({
type: 'and',
conditions: expect.arrayContaining([expect.objectContaining({ type: 'notExists' })]),
})
)
})
it('sweeps a deployed target version even when no draft referenced the failed id', async () => {
// Draft is clean (other-kb), but a deployed target version still points at the dropped
// placeholder - the deployed-target scope (not draft divergence) catches it.
@@ -376,5 +476,26 @@ describe('cleanup-failed', () => {
// The drop is skipped, so the placeholder row survives (no delete issued).
expect(dbChainMockFns.delete).not.toHaveBeenCalled()
})
it('keeps placeholders when a deployed-version cleanup fails after draft cleanup succeeds', async () => {
queueTableRows(workflow, [{ id: 'wf-1' }])
queueTableRows(workflowBlocks, [draftBlockRow('other-kb')])
queueTableRows(workflowDeploymentVersion, [
{ id: 'dv-failed', version: 5, state: versionState('failed-kb') },
])
dbChainMockFns.set.mockImplementationOnce(() => {
throw new Error('deployment update failed')
})
const cleaned = await clearFailedForkResourceReferences({
childWorkspaceId: 'child-ws',
failures: [{ kind: 'knowledge-base', childId: 'failed-kb', documentChildIds: [] }],
deployedTargetWorkflowIds: ['wf-deployed'],
requestId: 'test',
})
expect(cleaned).toEqual({ cleared: 0, clearingFailed: true })
expect(dbChainMockFns.delete).not.toHaveBeenCalled()
})
})
})
apps/sim/ee/workspace-forking/lib/copy/cleanup-failed.ts
+73 -10
View File
@@ -9,10 +9,13 @@ import {
} from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { getErrorMessage } from '@sim/utils/errors'
import { and, asc, eq, gt, inArray } from 'drizzle-orm'
import { and, asc, eq, exists, gt, inArray, isNull, notExists, sql } from 'drizzle-orm'
import { isRecord, type SubBlockRecord } from '@/lib/workflows/persistence/remap-internal-ids'
import { invalidateDeployedStateCache } from '@/lib/workflows/persistence/utils'
import type { ForkFailedResource } from '@/ee/workspace-forking/lib/copy/copy-resources'
import {
FORK_DOCUMENT_ID_PATTERN,
type ForkFailedResource,
} from '@/ee/workspace-forking/lib/copy/copy-resources'
import type { ForkCopyResolver } from '@/ee/workspace-forking/lib/remap/fork-bootstrap'
import {
clearDependentsOnRemap,
@@ -28,6 +31,30 @@ const WORKFLOW_PAGE = 200
/** Deployment versions loaded per page so a workflow with many versions never loads all at once. */
const DEPLOYMENT_VERSION_PAGE = 100
async function findKnowledgeBasesWithNonForkDocuments(ids: string[]): Promise<Set<string>> {
if (ids.length === 0) return new Set()
const rows = await db
.select({ id: knowledgeBase.id })
.from(knowledgeBase)
.where(and(inArray(knowledgeBase.id, ids), exists(liveNonForkDocumentQuery())))
.limit(ids.length)
return new Set(rows.map(({ id }) => id))
}
function liveNonForkDocumentQuery() {
return db
.select({ id: document.id })
.from(document)
.where(
and(
eq(document.knowledgeBaseId, knowledgeBase.id),
sql`${document.id} !~ ${FORK_DOCUMENT_ID_PATTERN}`,
isNull(document.deletedAt),
isNull(document.archivedAt)
)
)
}
/** Identity-or-clear resolver: a failed id resolves to null (cleared), any other id to itself. */
function buildFailedResolver(failedByKind: Map<ForkRemapKind, Set<string>>): ForkCopyResolver {
return (kind, id) => (failedByKind.get(kind)?.has(id) ? null : id)
@@ -77,12 +104,9 @@ function clearFailedSubBlockReferences(
* is the count of failed resources whose references were cleared.
*
* Storage accounting: this cleanup never decrements storage usage because it never removes
* anything that was counted. Copied file blobs are the only counted copies (incremented in
* `executeForkFileBlobCopies` only after the blob lands), and a failed file's blob never
* landed - its metadata row is intentionally left re-uploadable, and nothing was charged. The
* dropped table/KB/document placeholders are DB rows the upload path never counts, and any KB
* blobs copied before their KB failed are left in storage (rows only are dropped here) but
* uncounted - mirroring the KB upload path, which never counts KB blobs.
* anything that remains counted. A failed file copy is not charged and leaves its metadata row
* re-uploadable. A failed KB copy reverses its usage and retires its active file-ownership rows
* before reaching this cleanup; deterministic blobs remain available for a safe retry.
*/
export async function clearFailedForkResourceReferences(params: {
childWorkspaceId: string
@@ -94,6 +118,12 @@ export async function clearFailedForkResourceReferences(params: {
const { childWorkspaceId, failures, requestId = 'unknown' } = params
if (failures.length === 0) return { cleared: 0, clearingFailed: false }
const failedKnowledgeBaseIds = failures.flatMap((failure) =>
failure.kind === 'knowledge-base' ? [failure.childId] : []
)
const retainedKnowledgeBaseIds =
await findKnowledgeBasesWithNonForkDocuments(failedKnowledgeBaseIds)
const failedByKind = new Map<ForkRemapKind, Set<string>>()
const markFailed = (kind: ForkRemapKind, id: string) => {
const set = failedByKind.get(kind)
@@ -105,24 +135,43 @@ export async function clearFailedForkResourceReferences(params: {
// Standalone documents copied into an already-existing target KB (the doc-into-mapped-KB sync
// path) - dropped individually, since their KB is not ours to remove.
const docIds: string[] = []
let cleanupCount = 0
for (const failure of failures) {
if (failure.kind === 'table') {
markFailed('table', failure.childId)
tableIds.push(failure.childId)
cleanupCount += 1
} else if (failure.kind === 'knowledge-document') {
markFailed('knowledge-document', failure.childId)
docIds.push(failure.childId)
cleanupCount += 1
} else if (failure.kind === 'file') {
// A failed file blob: clear `file-upload` references to its copied storage key. No row to
// drop - the metadata row is left in place so the user can re-upload the missing blob.
markFailed('file', failure.childKey)
cleanupCount += 1
} else {
if (retainedKnowledgeBaseIds.has(failure.childId)) {
for (const docId of failure.documentChildIds) {
markFailed('knowledge-document', docId)
docIds.push(docId)
}
if (failure.documentChildIds.length > 0) cleanupCount += 1
logger.warn(
`[${requestId}] Keeping a failed copied knowledge base that contains non-fork documents`,
{ childWorkspaceId, childKnowledgeBaseId: failure.childId }
)
continue
}
markFailed('knowledge-base', failure.childId)
for (const docId of failure.documentChildIds) markFailed('knowledge-document', docId)
kbIds.push(failure.childId)
cleanupCount += 1
}
}
if (failedByKind.size === 0) return { cleared: 0, clearingFailed: false }
// Whether BOTH reference-clear phases completed without throwing. The placeholder drop below is
// gated on this: if clearing threw, a workflow (draft or deployed version) may still reference
// the failed id, so dropping its placeholder would create a dangling reference to a deleted row.
@@ -185,7 +234,9 @@ export async function clearFailedForkResourceReferences(params: {
await db.delete(userTableDefinitions).where(inArray(userTableDefinitions.id, tableIds))
}
if (kbIds.length > 0) {
await db.delete(knowledgeBase).where(inArray(knowledgeBase.id, kbIds))
await db
.delete(knowledgeBase)
.where(and(inArray(knowledgeBase.id, kbIds), notExists(liveNonForkDocumentQuery())))
}
if (docIds.length > 0) {
await db.delete(document).where(inArray(document.id, docIds))
@@ -197,7 +248,7 @@ export async function clearFailedForkResourceReferences(params: {
})
}
return { cleared: failures.length, clearingFailed: false }
return { cleared: cleanupCount, clearingFailed: false }
}
/**
@@ -304,6 +355,9 @@ export function rewriteDeploymentVersionState(
* no-op. After a version is rewritten its cached deployed state is evicted so execute/serve rebuilds
* from the cleaned snapshot. Bounded work (no long transaction): per-version short UPDATEs, versions
* keyset-paginated, and a per-workflow failure is logged without aborting the other workflows.
* After every workflow has been attempted, any failures are reported to the caller so it can keep
* the failed resource placeholders in place rather than deleting rows a deployed version may still
* reference.
*/
export async function clearFailedReferencesInDeploymentVersions(
workflowIds: ReadonlySet<string>,
@@ -312,6 +366,7 @@ export async function clearFailedReferencesInDeploymentVersions(
): Promise<void> {
if (workflowIds.size === 0) return
const resolve = buildFailedResolver(failedByKind)
const failures: unknown[] = []
for (const workflowId of workflowIds) {
try {
@@ -355,10 +410,18 @@ export async function clearFailedReferencesInDeploymentVersions(
afterVersion = versions[versions.length - 1].version
}
} catch (error) {
failures.push(error)
logger.error(`[${requestId}] Failed to clear references in deployment versions`, {
workflowId,
error: getErrorMessage(error),
})
}
}
if (failures.length > 0) {
throw new AggregateError(
failures,
`Failed to clear deployment-version references for ${failures.length} workflow(s)`
)
}
}
apps/sim/ee/workspace-forking/lib/copy/copy-resources.test.ts
+813 -57
View File
File diff suppressed because it is too large Load Diff
apps/sim/ee/workspace-forking/lib/copy/copy-resources.ts
+468 -161
View File
@@ -14,13 +14,26 @@ import {
userTableRowSecretProvenance,
userTableRows,
workflowMcpServer,
workspaceFiles,
} from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { sha256Hex } from '@sim/security/hash'
import { getErrorMessage } from '@sim/utils/errors'
import { generateId } from '@sim/utils/id'
import { omit } from '@sim/utils/object'
import { and, asc, eq, gt, inArray, isNotNull, isNull, type SQL, sql } from 'drizzle-orm'
import {
and,
asc,
eq,
exists,
gt,
inArray,
isNotNull,
isNull,
or,
type SQL,
sql,
} from 'drizzle-orm'
import {
decrementStorageUsageForBillingContextInTx,
incrementStorageUsageForBillingContextInTx,
@@ -29,7 +42,10 @@ import {
} from '@/lib/billing/storage'
import { mapWithConcurrency } from '@/lib/core/utils/concurrency'
import type { DbOrTx } from '@/lib/db/types'
import type { DurableSecretProvenance } from '@/lib/execution/durable-secret-provenance'
import {
type DurableSecretProvenance,
hashDurableSecretProvenanceValue,
} from '@/lib/execution/durable-secret-provenance'
import {
createKnowledgeDocumentSourceValue,
type KnowledgeDocumentSourceValue,
@@ -50,11 +66,17 @@ import {
headObject,
uploadFile,
} from '@/lib/uploads/core/storage-service'
import {
type KnowledgeBaseFileOwnership,
recordKnowledgeBaseFileOwnership,
} from '@/lib/uploads/server/metadata'
import { MAX_FILE_SIZE } from '@/lib/uploads/utils/validation'
import { isRecord } from '@/lib/workflows/persistence/remap-internal-ids'
import type {
ForkMappingUpsert,
ForkResourceType,
import {
deleteCopiedResourceMappingsByTargets,
type ForkMappingUpsert,
type ForkResourceType,
persistCopiedResourceMappings,
} from '@/ee/workspace-forking/lib/mapping/mapping-store'
import type { ForkBlockIdResolver } from '@/ee/workspace-forking/lib/remap/block-identity'
import {
@@ -97,6 +119,7 @@ function isForkProvenancePageWithinBudget(sidecars: readonly { entries: unknown
* processes one page at a time, so peak concurrency stays at this cap regardless of KB size.
*/
const KB_DOCUMENT_COPY_CONCURRENCY = 5
export const FORK_DOCUMENT_ID_PATTERN = '^fork_document_[0-9a-f]{40}$'
function deriveCopyIdentity(
kind: 'document' | 'embedding',
@@ -107,9 +130,50 @@ function deriveCopyIdentity(
return `fork_${kind}_${digest}`
}
/** Stable object key so a replay overwrites or reuses the same copied KB blob. */
function deriveKbDocumentStorageKey(childDocumentId: string): string {
return `kb/fork-${childDocumentId}`
/**
* Stable legacy key prefix for a copied KB document. New blobs append their content digest so
* retries of different source snapshots cannot overwrite one another; the unsuffixed form remains
* valid for copies finalized by an older worker during a rolling deployment.
*/
function deriveKbDocumentStorageKey(childDocumentId: string, contentHash?: string): string {
const prefix = `kb/fork-${childDocumentId}`
return contentHash ? `${prefix}-${contentHash}` : prefix
}
function isKbDocumentStorageKey(key: string, childDocumentId: string): boolean {
const prefix = deriveKbDocumentStorageKey(childDocumentId)
if (key === prefix) return true
if (!key.startsWith(`${prefix}-`)) return false
return /^[0-9a-f]{64}$/.test(key.slice(prefix.length + 1))
}
interface TargetDocumentExpectation {
childDocumentId: string
childKnowledgeBaseId: string
}
interface TargetDocumentState {
id: string
knowledgeBaseId: string
storageKey: string | null
archivedAt: Date | null
deletedAt: Date | null
}
function validateTargetDocumentState(
row: TargetDocumentState,
expected: TargetDocumentExpectation
): 'active' | 'archived' {
if (row.id !== expected.childDocumentId) {
throw new Error(`Copied document ${row.id} has an unexpected identity`)
}
if (row.knowledgeBaseId !== expected.childKnowledgeBaseId || row.deletedAt) {
throw new Error(`Copied document ${row.id} has conflicting storage identity`)
}
if (row.storageKey !== null && !isKbDocumentStorageKey(row.storageKey, row.id)) {
throw new Error(`Copied document ${row.id} has conflicting storage`)
}
return row.archivedAt ? 'archived' : 'active'
}
/**
@@ -157,6 +221,13 @@ export interface CopyResourcesParams {
* omits it, defaulting to the deterministic derive (a fresh child has no pairs).
*/
resolveBlockId?: ForkBlockIdResolver
/** Canonical fork-edge orientation for document identities completed by the background copy. */
documentMappingContext: ForkDocumentMappingContext
}
export interface ForkDocumentMappingContext {
edgeChildWorkspaceId: string
sourceIsParent: boolean
}
export interface ForkContentPlanEntry {
@@ -198,7 +269,10 @@ export interface ForkContentDocumentEntry {
sourceDocId: string
childDocId: string
childKnowledgeBaseId: string
/** Source blob fields captured at placeholder time, for the post-commit blob re-key. */
/**
* Source blob fields retained in the serialized payload for rolling-deploy and queued-job
* compatibility. Current workers re-read the live source row before copying it.
*/
storageKey: string | null
fileUrl: string
fileSize: number
@@ -217,6 +291,11 @@ export interface ForkContentPlan {
skills: ForkContentSkillEntry[]
/** Documents copied into an already-existing target KB (sync-only; empty at fork create). */
documents: ForkContentDocumentEntry[]
/**
* Optional only so workers deployed during a rollout can still consume already-queued payloads.
* Every newly planned fork/sync includes it.
*/
documentMappingContext?: ForkDocumentMappingContext
}
/**
@@ -301,6 +380,7 @@ export async function copyForkResourceContainers(
knowledgeBases: [],
skills: [],
documents: [],
documentMappingContext: params.documentMappingContext,
}
const names: ForkCopiedResourceNames = {
tables: [],
@@ -739,12 +819,13 @@ export async function planForkMappedKbDocumentCopies(params: {
resolver: ForkReferenceResolver
referencedDocumentIds: string[]
alreadyCopiedSourceDocIds: Set<string>
now: Date
}): Promise<{
documents: ForkContentDocumentEntry[]
docIdMap: Map<string, string>
mappingEntries: ForkMappingUpsert[]
}> {
const { tx, resolver, referencedDocumentIds, alreadyCopiedSourceDocIds } = params
const { tx, resolver, referencedDocumentIds, alreadyCopiedSourceDocIds, now } = params
const documents: ForkContentDocumentEntry[] = []
const docIdMap = new Map<string, string>()
const mappingEntries: ForkMappingUpsert[] = []
@@ -768,41 +849,74 @@ export async function planForkMappedKbDocumentCopies(params: {
)
)
const planned = docs.flatMap((doc) => {
const targetKbId = resolver('knowledge-base', doc.knowledgeBaseId)
if (targetKbId == null) return []
return [{ doc, targetKbId, childDocId: deriveCopyIdentity('document', targetKbId, doc.id) }]
})
const existingTargets =
planned.length === 0
? []
: await tx
.select({
id: document.id,
knowledgeBaseId: document.knowledgeBaseId,
storageKey: document.storageKey,
archivedAt: document.archivedAt,
deletedAt: document.deletedAt,
})
.from(document)
.where(
inArray(
document.id,
planned.map(({ childDocId }) => childDocId)
)
)
const existingTargetById = new Map(existingTargets.map((target) => [target.id, target]))
const inserts: (typeof document.$inferInsert)[] = []
for (const doc of docs) {
for (const { doc, targetKbId, childDocId } of planned) {
// The parent KB must already exist in the target. The resolver returns a target KB id only
// for a mapped, still-existing KB (validTargetIdsByKind), so this is FK-safe; a doc whose KB
// isn't mapped resolves null here and is left for its reference to be cleared.
const targetKbId = resolver('knowledge-base', doc.knowledgeBaseId)
if (targetKbId == null) continue
const childDocId = deriveCopyIdentity('document', targetKbId, doc.id)
inserts.push({
...doc,
id: childDocId,
knowledgeBaseId: targetKbId,
connectorId: null,
storageKey: null,
fileUrl: '',
fileSize: 0,
deletedAt: null,
archivedAt: new Date(),
})
const existingTarget = existingTargetById.get(childDocId)
const expectedTarget = {
childDocumentId: childDocId,
childKnowledgeBaseId: targetKbId,
}
const existingTargetState = existingTarget
? validateTargetDocumentState(existingTarget, expectedTarget)
: null
if (!existingTarget) {
inserts.push({
...doc,
id: childDocId,
knowledgeBaseId: targetKbId,
connectorId: null,
storageKey: null,
fileUrl: '',
fileSize: 0,
deletedAt: null,
archivedAt: now,
})
}
docIdMap.set(doc.id, childDocId)
mappingEntries.push({
resourceType: 'knowledge_document',
parentResourceId: doc.id,
childResourceId: childDocId,
})
documents.push({
sourceDocId: doc.id,
childDocId,
childKnowledgeBaseId: targetKbId,
storageKey: doc.storageKey,
fileUrl: doc.fileUrl,
fileSize: doc.fileSize,
filename: doc.filename,
mimeType: doc.mimeType,
})
if (!existingTarget || existingTargetState === 'archived') {
documents.push({
sourceDocId: doc.id,
childDocId,
childKnowledgeBaseId: targetKbId,
storageKey: doc.storageKey,
fileUrl: doc.fileUrl,
fileSize: doc.fileSize,
filename: doc.filename,
mimeType: doc.mimeType,
})
}
}
if (inserts.length > 0) await tx.insert(document).values(inserts)
return { documents, docIdMap, mappingEntries }
@@ -1002,15 +1116,18 @@ export async function copyForkResourceContent(params: {
kb.documentIdMap[source.id] ?? deriveCopyIdentity('document', kb.childId, source.id),
}))
const activeTargetDocumentIds = await getActiveTargetDocumentIds(
documentCopies.map(({ childDocumentId }) => childDocumentId)
documentCopies.map(({ childDocumentId }) => ({
childDocumentId,
childKnowledgeBaseId: kb.childId,
}))
)
const documentsToCopy = documentCopies.filter(
({ childDocumentId }) => !activeTargetDocumentIds.has(childDocumentId)
)
// Copy the page's documents with bounded concurrency. The mapper never rejects
// (it captures its error), so all in-flight work settles before this resolves - no
// orphaned writes survive a failure - and a captured error is rethrown after to keep
// the KB ALL-OR-NOTHING (any failed doc fails the whole KB -> cleanup below).
// (it captures its error), so all in-flight work settles before this resolves and a
// captured error is rethrown after to keep the KB ALL-OR-NOTHING (any failed doc fails
// the whole KB -> cleanup below).
if (documentsToCopy.length > 0) {
const resolvedBillingContext = await getBillingContext()
const docErrors = await mapWithConcurrency(
@@ -1035,6 +1152,22 @@ export async function copyForkResourceContent(params: {
const docError = docErrors.find((error) => error != null)
if (docError) throw docError
}
const mappingContext = contentPlan.documentMappingContext
if (mappingContext) {
await db.transaction(async (tx) => {
await persistCopiedResourceMappings({
executor: tx,
edgeChildWorkspaceId: mappingContext.edgeChildWorkspaceId,
userId,
sourceIsParent: mappingContext.sourceIsParent,
entries: documentCopies.map(({ source, childDocumentId }) => ({
resourceType: 'knowledge_document',
parentResourceId: source.id,
childResourceId: childDocumentId,
})),
})
})
}
afterDocId = docs[docs.length - 1].id
if (docs.length < CONTENT_PAGE) break
}
@@ -1052,6 +1185,19 @@ export async function copyForkResourceContent(params: {
{ cause: rollbackError }
)
}
if (contentPlan.documentMappingContext) {
try {
await deleteFailedKnowledgeBaseDocumentMappings(
kb.childId,
contentPlan.documentMappingContext
)
} catch (mappingCleanupError) {
logger.error(`[${requestId}] Failed to clean mappings for a failed copied KB`, {
childKnowledgeBaseId: kb.childId,
error: getErrorMessage(mappingCleanupError),
})
}
}
failedResources += 1
failures.push({
kind: 'knowledge-base',
@@ -1072,50 +1218,54 @@ export async function copyForkResourceContent(params: {
// own documents are never touched.
for (const docEntry of contentPlan.documents) {
try {
const active = await isActiveTargetDocument(docEntry.childDocId)
const active = await isActiveTargetDocument({
childDocumentId: docEntry.childDocId,
childKnowledgeBaseId: docEntry.childKnowledgeBaseId,
})
if (active) {
copiedResources += 1
continue
}
const [source] = await db
.select()
.from(document)
.where(
and(
eq(document.id, docEntry.sourceDocId),
isNull(document.deletedAt),
isNull(document.archivedAt)
)
)
.limit(1)
if (!source) {
throw new Error(`Source document ${docEntry.sourceDocId} is missing`)
}
const resolvedBillingContext = await getBillingContext()
const blob = await copyKbDocumentBlob(
{
storageKey: docEntry.storageKey,
filename: docEntry.filename,
mimeType: docEntry.mimeType,
},
await copyKbDocument({
source,
childDocumentId: docEntry.childDocId,
childKnowledgeBaseId: docEntry.childKnowledgeBaseId,
childWorkspaceId,
userId,
docEntry.childDocId
)
try {
await copyDocumentEmbeddings(
docEntry.sourceDocId,
docEntry.childDocId,
docEntry.childKnowledgeBaseId
)
await finalizeKbDocument({
childDocumentId: docEntry.childDocId,
childKnowledgeBaseId: docEntry.childKnowledgeBaseId,
billingContext: resolvedBillingContext,
bytes: blob ? docEntry.fileSize : 0,
values: {
knowledgeBaseId: docEntry.childKnowledgeBaseId,
connectorId: null,
storageKey: blob?.storageKey ?? null,
fileUrl: blob?.fileUrl ?? docEntry.fileUrl,
fileSize: docEntry.fileSize,
archivedAt: null,
deletedAt: null,
uploadedBy: userId,
},
})
} catch (error) {
if (blob) await cleanupCopiedKbBlob(blob.storageKey)
throw error
}
billingContext: resolvedBillingContext,
})
copiedResources += 1
} catch (error) {
if (contentPlan.documentMappingContext) {
try {
await deleteCopiedResourceMappingsByTargets({
executor: db,
edgeChildWorkspaceId: contentPlan.documentMappingContext.edgeChildWorkspaceId,
sourceIsParent: contentPlan.documentMappingContext.sourceIsParent,
targets: [{ resourceType: 'knowledge_document', resourceId: docEntry.childDocId }],
})
} catch (mappingCleanupError) {
logger.error(`[${requestId}] Failed to clean mapping for a failed copied document`, {
childDocumentId: docEntry.childDocId,
error: getErrorMessage(mappingCleanupError),
})
}
}
failedResources += 1
failures.push({ kind: 'knowledge-document', childId: docEntry.childDocId })
logger.warn(`[${requestId}] Failed to copy document into mapped KB during sync`, {
@@ -1172,24 +1322,38 @@ export async function copyForkResourceContent(params: {
return { copied: copiedResources, failed: failedResources, failures }
}
async function getActiveTargetDocumentIds(childDocumentIds: string[]): Promise<Set<string>> {
if (childDocumentIds.length === 0) return new Set()
const active = await db
.select({ id: document.id })
async function getActiveTargetDocumentIds(
expectations: TargetDocumentExpectation[]
): Promise<Set<string>> {
if (expectations.length === 0) return new Set()
const expectedById = new Map(expectations.map((expected) => [expected.childDocumentId, expected]))
const existing = await db
.select({
id: document.id,
knowledgeBaseId: document.knowledgeBaseId,
storageKey: document.storageKey,
archivedAt: document.archivedAt,
deletedAt: document.deletedAt,
})
.from(document)
.where(
and(
inArray(document.id, childDocumentIds),
isNull(document.deletedAt),
isNull(document.archivedAt)
inArray(
document.id,
expectations.map(({ childDocumentId }) => childDocumentId)
)
)
.limit(childDocumentIds.length)
return new Set(active.map((row) => row.id))
.limit(expectations.length)
const activeIds = new Set<string>()
for (const row of existing) {
const expected = expectedById.get(row.id)
if (!expected) throw new Error(`Copied document ${row.id} was not requested`)
if (validateTargetDocumentState(row, expected) === 'active') activeIds.add(row.id)
}
return activeIds
}
async function isActiveTargetDocument(childDocumentId: string): Promise<boolean> {
return (await getActiveTargetDocumentIds([childDocumentId])).has(childDocumentId)
async function isActiveTargetDocument(expectation: TargetDocumentExpectation): Promise<boolean> {
return (await getActiveTargetDocumentIds([expectation])).has(expectation.childDocumentId)
}
/**
@@ -1231,19 +1395,21 @@ async function finalizeKbDocument(params: {
billingContext: StorageBillingContext
bytes: number
values: Partial<typeof document.$inferInsert>
fileOwnership?: KnowledgeBaseFileOwnership
secretProvenance?: DurableSecretProvenance
provenanceSource?: KnowledgeDocumentSourceValue
}): Promise<void> {
}): Promise<string | null> {
const {
childDocumentId,
childKnowledgeBaseId,
billingContext,
bytes,
values,
fileOwnership,
secretProvenance,
provenanceSource,
} = params
await db.transaction(async (tx) => {
return db.transaction(async (tx) => {
const [lockedKnowledgeBase] = await tx
.select({ workspaceId: knowledgeBase.workspaceId })
.from(knowledgeBase)
@@ -1257,6 +1423,11 @@ async function finalizeKbDocument(params: {
`Copied document knowledge base ${childKnowledgeBaseId} moved from workspace ${billingContext.workspaceId}; refusing stale storage charge`
)
}
if (fileOwnership && fileOwnership.workspaceId !== lockedKnowledgeBase.workspaceId) {
throw new Error(
`Copied document ${childDocumentId} ownership does not match its knowledge base workspace`
)
}
const [activated] = await tx
.update(document)
@@ -1264,6 +1435,7 @@ async function finalizeKbDocument(params: {
.where(
and(
eq(document.id, childDocumentId),
eq(document.knowledgeBaseId, childKnowledgeBaseId),
isNull(document.deletedAt),
isNotNull(document.archivedAt)
)
@@ -1272,7 +1444,15 @@ async function finalizeKbDocument(params: {
if (!activated) {
const [active] = await tx
.select({ id: document.id })
.select({
id: document.id,
knowledgeBaseId: document.knowledgeBaseId,
storageKey: document.storageKey,
filename: document.filename,
mimeType: document.mimeType,
fileSize: document.fileSize,
uploadedBy: document.uploadedBy,
})
.from(document)
.where(
and(
@@ -1282,8 +1462,32 @@ async function finalizeKbDocument(params: {
)
)
.limit(1)
if (active) return
throw new Error(`Copied document placeholder ${childDocumentId} is missing`)
if (!active) throw new Error(`Copied document placeholder ${childDocumentId} is missing`)
if (active.knowledgeBaseId !== childKnowledgeBaseId) {
throw new Error(`Copied document ${childDocumentId} has conflicting active storage`)
}
if (fileOwnership) {
const activeStorageKey = active.storageKey
if (!activeStorageKey || !isKbDocumentStorageKey(activeStorageKey, childDocumentId)) {
throw new Error(`Copied document ${childDocumentId} has conflicting active storage`)
}
await recordKnowledgeBaseFileOwnership(
{
key: activeStorageKey,
userId: active.uploadedBy ?? fileOwnership.userId,
workspaceId: fileOwnership.workspaceId,
originalName: active.filename,
contentType: active.mimeType,
size: active.fileSize,
},
tx
)
}
return active.storageKey
}
if (fileOwnership) {
await recordKnowledgeBaseFileOwnership(fileOwnership, tx)
}
if (secretProvenance && provenanceSource) {
@@ -1296,6 +1500,7 @@ async function finalizeKbDocument(params: {
}
await incrementStorageUsageForBillingContextInTx(tx, billingContext, bytes)
return fileOwnership?.key ?? null
})
}
@@ -1320,51 +1525,77 @@ async function copyKbDocument(params: {
userId,
billingContext,
} = params
await ensureKbDocumentPlaceholder(source, childDocumentId, childKnowledgeBaseId, userId)
const sourceSecretContext = await loadKnowledgeDocumentDurableSecretProvenance(source.id)
const sourceSnapshotHash = hashDurableSecretProvenanceValue(
createKnowledgeDocumentSourceValue(source)
)
const provenanceSnapshotHash = hashDurableSecretProvenanceValue(sourceSecretContext.source)
if (!sourceSnapshotHash || sourceSnapshotHash !== provenanceSnapshotHash) {
throw new Error(`Knowledge document ${source.id} changed while preparing its fork copy`)
}
await ensureKbDocumentPlaceholder(source, childDocumentId, childKnowledgeBaseId, userId)
const blob = await copyKbDocumentBlob(source, childWorkspaceId, userId, childDocumentId)
try {
await copyDocumentEmbeddings(source.id, childDocumentId, childKnowledgeBaseId)
const copiedValues = {
...omit(source, ['id', 'knowledgeBaseId']),
knowledgeBaseId: childKnowledgeBaseId,
connectorId: null,
storageKey: blob?.storageKey ?? null,
fileUrl: blob?.fileUrl ?? source.fileUrl,
archivedAt: null,
deletedAt: null,
uploadedBy: userId,
secretProvenanceVersion: sourceSecretContext.tracked ? 1 : null,
await copyDocumentEmbeddings(source.id, childDocumentId, childKnowledgeBaseId)
const copiedValues = {
...omit(source, ['id', 'knowledgeBaseId']),
knowledgeBaseId: childKnowledgeBaseId,
connectorId: null,
storageKey: blob?.storageKey ?? null,
fileUrl: blob?.fileUrl ?? source.fileUrl,
archivedAt: null,
deletedAt: null,
uploadedBy: userId,
secretProvenanceVersion: sourceSecretContext.tracked ? 1 : null,
}
const copiedSource = createKnowledgeDocumentSourceValue(copiedValues)
const finalizedStorageKey = await finalizeKbDocument({
childDocumentId,
childKnowledgeBaseId,
billingContext,
bytes: blob ? source.fileSize : 0,
values: copiedValues,
...(blob
? {
fileOwnership: {
key: blob.storageKey,
userId,
workspaceId: childWorkspaceId,
originalName: source.filename,
contentType: source.mimeType,
size: source.fileSize,
},
}
: {}),
...(sourceSecretContext.tracked
? {
secretProvenance: rebindKnowledgeDocumentSecretProvenance(
sourceSecretContext.provenance,
sourceSecretContext.source,
copiedSource
),
provenanceSource: copiedSource,
}
: {}),
})
if (blob && finalizedStorageKey !== blob.storageKey) {
try {
await deleteFile({ key: blob.storageKey, context: 'knowledge-base' })
} catch (error) {
logger.warn(`Failed to remove an unreferenced losing fork document blob`, {
childDocumentId,
storageKey: blob.storageKey,
error: getErrorMessage(error),
})
}
const copiedSource = createKnowledgeDocumentSourceValue(copiedValues)
await finalizeKbDocument({
childDocumentId,
childKnowledgeBaseId,
billingContext,
bytes: blob ? source.fileSize : 0,
values: copiedValues,
...(sourceSecretContext.tracked
? {
secretProvenance: rebindKnowledgeDocumentSecretProvenance(
sourceSecretContext.provenance,
sourceSecretContext.source,
copiedSource
),
provenanceSource: copiedSource,
}
: {}),
})
} catch (error) {
if (blob) await cleanupCopiedKbBlob(blob.storageKey)
throw error
}
}
/**
* Reverse any documents already activated for a KB when a later document fails,
* preserving the existing all-or-nothing KB failure semantics without a long
* parent transaction. The aggregate keeps memory bounded regardless of KB size.
* parent transaction. Accounting reversal and archival are limited to reserved
* deterministic fork identities rather than every document in the target KB.
*/
async function rollbackCopiedKbDocuments(
childKnowledgeBaseId: string,
@@ -1389,6 +1620,7 @@ async function rollbackCopiedKbDocuments(
.where(
and(
eq(document.knowledgeBaseId, childKnowledgeBaseId),
sql`${document.id} ~ ${FORK_DOCUMENT_ID_PATTERN}`,
isNull(document.deletedAt),
isNull(document.archivedAt),
isNotNull(document.storageKey)
@@ -1396,12 +1628,41 @@ async function rollbackCopiedKbDocuments(
)
const bytes = Number(usage?.total ?? 0)
await decrementStorageUsageForBillingContextInTx(tx, billingContext, bytes)
await tx
.update(workspaceFiles)
.set({ deletedAt: new Date() })
.where(
and(
eq(workspaceFiles.workspaceId, childWorkspaceId),
eq(workspaceFiles.context, 'knowledge-base'),
isNull(workspaceFiles.deletedAt),
exists(
tx
.select({ id: document.id })
.from(document)
.where(
and(
eq(document.knowledgeBaseId, childKnowledgeBaseId),
sql`${document.id} ~ ${FORK_DOCUMENT_ID_PATTERN}`,
isNull(document.deletedAt),
isNull(document.archivedAt),
eq(document.storageKey, workspaceFiles.key),
or(
eq(workspaceFiles.key, sql<string>`'kb/fork-' || ${document.id}`),
sql`${workspaceFiles.key} ~ ('^kb/fork-' || ${document.id} || '-[0-9a-f]{64}$')`
)
)
)
)
)
)
await tx
.update(document)
.set({ archivedAt: new Date() })
.where(
and(
eq(document.knowledgeBaseId, childKnowledgeBaseId),
sql`${document.id} ~ ${FORK_DOCUMENT_ID_PATTERN}`,
isNull(document.deletedAt),
isNull(document.archivedAt)
)
@@ -1409,6 +1670,50 @@ async function rollbackCopiedKbDocuments(
})
}
/**
* Remove the identity rows completed for earlier pages of a KB whose later page failed. Target
* document ids are keyset-paged from the failed KB, so cleanup never retains the whole KB in app
* memory. The target side is selected from the same serialized edge orientation used to write the
* mappings.
*/
async function deleteFailedKnowledgeBaseDocumentMappings(
childKnowledgeBaseId: string,
mappingContext: ForkDocumentMappingContext
): Promise<void> {
let afterId: string | null = null
for (;;) {
const rows = await db
.select({ id: document.id })
.from(document)
.where(
afterId == null
? and(
eq(document.knowledgeBaseId, childKnowledgeBaseId),
sql`${document.id} ~ ${FORK_DOCUMENT_ID_PATTERN}`
)
: and(
eq(document.knowledgeBaseId, childKnowledgeBaseId),
sql`${document.id} ~ ${FORK_DOCUMENT_ID_PATTERN}`,
gt(document.id, afterId)
)
)
.orderBy(asc(document.id))
.limit(CONTENT_PAGE)
if (rows.length === 0) break
await deleteCopiedResourceMappingsByTargets({
executor: db,
edgeChildWorkspaceId: mappingContext.edgeChildWorkspaceId,
sourceIsParent: mappingContext.sourceIsParent,
targets: rows.map(({ id }) => ({
resourceType: 'knowledge_document' as const,
resourceId: id,
})),
})
if (rows.length < CONTENT_PAGE) break
afterId = rows[rows.length - 1].id
}
}
async function copyDocumentEmbeddings(
sourceDocumentId: string,
childDocumentId: string,
@@ -1488,48 +1793,50 @@ async function copyDocumentEmbeddings(
* `verifyKBFileAccess` grants a child-workspace member - without it the copied object is
* download-denied (no binding = deny). Returns the new `storageKey` + serve `fileUrl`, or null
* when there is no internal blob to copy (external/`data:` docs have a null `storageKey`) or the
* copy fails. A stored source blob is required to copy successfully; callers
* keep the target placeholder archived and report the existing resource failure.
* copy fails. A stored source blob is required to copy successfully; callers keep the target
* placeholder archived and report the existing resource failure. The content digest in the key
* makes reuse safe for identical retries and prevents a later source snapshot from adopting or
* overwriting bytes left by an earlier failed attempt. Ownership is recorded before storage I/O,
* matching the presigned-upload lifecycle: successful finalization reuses the immutable binding,
* while the existing orphan-binding sweep eventually reclaims an abandoned object or reservation.
*/
async function copyKbDocumentBlob(
doc: { storageKey: string | null; filename: string; mimeType: string },
doc: { storageKey: string | null; filename: string; mimeType: string; fileSize: number },
childWorkspaceId: string,
userId: string,
childDocumentId: string
): Promise<{ storageKey: string; fileUrl: string } | null> {
if (!doc.storageKey) return null
const targetKey = deriveKbDocumentStorageKey(childDocumentId)
try {
const existing = await headObject(targetKey, 'knowledge-base')
if (!existing) {
const buffer = await downloadFile({
key: doc.storageKey,
context: 'knowledge-base',
maxBytes: MAX_FILE_SIZE,
})
await uploadFile({
file: buffer,
fileName: doc.filename,
contentType: doc.mimeType,
context: 'knowledge-base',
customKey: targetKey,
preserveKey: true,
persistMetadata: false,
metadata: {
userId,
workspaceId: childWorkspaceId,
originalName: doc.filename,
},
})
}
} catch (error) {
await cleanupCopiedKbBlob(targetKey)
throw error
const buffer = await downloadFile({
key: doc.storageKey,
context: 'knowledge-base',
maxBytes: MAX_FILE_SIZE,
})
const targetKey = deriveKbDocumentStorageKey(childDocumentId, sha256Hex(buffer))
await recordKnowledgeBaseFileOwnership({
key: targetKey,
userId,
workspaceId: childWorkspaceId,
originalName: doc.filename,
contentType: doc.mimeType,
size: doc.fileSize,
})
const existing = await headObject(targetKey, 'knowledge-base')
if (!existing) {
await uploadFile({
file: buffer,
fileName: doc.filename,
contentType: doc.mimeType,
context: 'knowledge-base',
customKey: targetKey,
preserveKey: true,
persistMetadata: false,
metadata: {
userId,
workspaceId: childWorkspaceId,
originalName: doc.filename,
},
})
}
return { storageKey: targetKey, fileUrl: `/api/files/serve/${encodeURIComponent(targetKey)}` }
}
/** Best-effort orphan cleanup after DB finalization or embedding copy fails. */
async function cleanupCopiedKbBlob(storageKey: string): Promise<void> {
await deleteFile({ key: storageKey, context: 'knowledge-base' }).catch(() => {})
}
apps/sim/ee/workspace-forking/lib/create-fork.test.ts
+8
View File
@@ -200,6 +200,14 @@ describe('createFork storage headroom gate', () => {
bytes: 500,
})
expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(1)
expect(mockCopyForkResourceContainers).toHaveBeenCalledWith(
expect.objectContaining({
documentMappingContext: {
edgeChildWorkspaceId: result.workspace.id,
sourceIsParent: true,
},
})
)
})
it('seeds identity mappings for copied FILES by storage key (a later sync must not re-offer them)', async () => {
apps/sim/ee/workspace-forking/lib/create-fork.ts
+4
View File
@@ -257,6 +257,10 @@ export async function createFork(params: CreateForkParams): Promise<CreateForkRe
},
workflowIdMap,
referencedDocumentIds: Array.from(referencedDocumentIds),
documentMappingContext: {
edgeChildWorkspaceId: childWorkspaceId,
sourceIsParent: true,
},
})
forkedResourceNames = resourceResult.names
apps/sim/ee/workspace-forking/lib/mapping/mapping-store.test.ts
+36
View File
@@ -5,6 +5,8 @@ import { describe, expect, it } from 'vitest'
import {
buildForkResolver,
type ForkMappingRow,
type ForkMappingUpsert,
orientCopiedResourceMappings,
} from '@/ee/workspace-forking/lib/mapping/mapping-store'
const credentialRow: ForkMappingRow = {
@@ -15,6 +17,40 @@ const credentialRow: ForkMappingRow = {
childResourceId: 'cred-child',
}
const copiedEntry: ForkMappingUpsert = {
resourceType: 'knowledge_document',
parentResourceId: 'runtime-source-doc',
childResourceId: 'runtime-target-doc',
}
describe('orientCopiedResourceMappings', () => {
it('keeps fork/pull source-parent mappings in canonical orientation', () => {
expect(orientCopiedResourceMappings(true, [copiedEntry])).toEqual({
entries: [copiedEntry],
deleteKeys: [],
})
})
it('swaps push source-child mappings and removes the prior row keyed by that child', () => {
expect(orientCopiedResourceMappings(false, [copiedEntry])).toEqual({
entries: [
{
resourceType: 'knowledge_document',
parentResourceId: 'runtime-target-doc',
childResourceId: 'runtime-source-doc',
},
],
deleteKeys: [{ resourceType: 'knowledge_document', childResourceId: 'runtime-source-doc' }],
})
})
it('does not produce a push mapping for an unmapped null target', () => {
expect(
orientCopiedResourceMappings(false, [{ ...copiedEntry, childResourceId: null }])
).toEqual({ entries: [], deleteKeys: [] })
})
})
describe('buildForkResolver', () => {
it('resolves source->target for a pull (source is parent)', () => {
const resolve = buildForkResolver([credentialRow], { sourceIsParent: true })
apps/sim/ee/workspace-forking/lib/mapping/mapping-store.ts
+106 -7
View File
@@ -29,6 +29,21 @@ export interface ForkMappingUpsert {
childResourceId: string | null
}
export interface PersistCopiedResourceMappingsParams {
executor: DbOrTx
edgeChildWorkspaceId: string
userId: string
/** Whether the runtime copy source is the canonical parent side of the fork edge. */
sourceIsParent: boolean
/** Runtime source -> runtime target identities produced by the copy operation. */
entries: ForkMappingUpsert[]
}
export interface OrientedCopiedResourceMappings {
entries: ForkMappingUpsert[]
deleteKeys: Array<{ resourceType: ForkResourceType; childResourceId: string }>
}
const RESOURCE_TYPE_TO_FORK_KIND: Record<ForkResourceType, ForkRemapKind | null> = {
workflow: null,
oauth_credential: 'credential',
@@ -216,6 +231,56 @@ export async function upsertEdgeMappings(
}
}
/**
* Orient runtime source -> target copy identities into the edge's canonical parent -> child
* storage shape. Pull/fork copies already have that orientation. Push copies run child -> parent,
* so their pairs are swapped and the old row keyed by the source-child identity is removed before
* the replacement is upserted.
*/
export function orientCopiedResourceMappings(
sourceIsParent: boolean,
entries: ForkMappingUpsert[]
): OrientedCopiedResourceMappings {
if (sourceIsParent) return { entries, deleteKeys: [] }
const oriented: ForkMappingUpsert[] = []
const deleteKeys: OrientedCopiedResourceMappings['deleteKeys'] = []
for (const entry of entries) {
if (entry.childResourceId == null) continue
oriented.push({
resourceType: entry.resourceType,
parentResourceId: entry.childResourceId,
childResourceId: entry.parentResourceId,
})
deleteKeys.push({
resourceType: entry.resourceType,
childResourceId: entry.parentResourceId,
})
}
return { entries: oriented, deleteKeys }
}
/**
* Persist identities created by a copy operation through one shared orientation boundary. Used by
* both the promote transaction and the post-commit full-KB document copier so those paths cannot
* disagree about parent/child direction.
*/
export async function persistCopiedResourceMappings({
executor,
edgeChildWorkspaceId,
userId,
sourceIsParent,
entries,
}: PersistCopiedResourceMappingsParams): Promise<void> {
if (entries.length === 0) return
const oriented = orientCopiedResourceMappings(sourceIsParent, entries)
if (oriented.entries.length === 0) return
if (oriented.deleteKeys.length > 0) {
await deleteEdgeMappingsByChildResources(executor, edgeChildWorkspaceId, oriented.deleteKeys)
}
await upsertEdgeMappings(executor, edgeChildWorkspaceId, userId, oriented.entries)
}
/**
* Remove mapping rows matched by their child-side (source) resource id, grouped by
* resource type into a single OR-of-INs - one query for the whole push save (the
@@ -226,19 +291,53 @@ export async function deleteEdgeMappingsByChildResources(
tx: DbOrTx,
childWorkspaceId: string,
pairs: Array<{ resourceType: ForkResourceType; childResourceId: string }>
): Promise<void> {
await deleteEdgeMappingsByResourceIds(
tx,
childWorkspaceId,
'child',
pairs.map(({ resourceType, childResourceId }) => ({
resourceType,
resourceId: childResourceId,
}))
)
}
/** Remove copy mappings whose runtime target resources are being discarded after a failed fill. */
export async function deleteCopiedResourceMappingsByTargets(params: {
executor: DbOrTx
edgeChildWorkspaceId: string
sourceIsParent: boolean
targets: Array<{ resourceType: ForkResourceType; resourceId: string }>
}): Promise<void> {
const { executor, edgeChildWorkspaceId, sourceIsParent, targets } = params
await deleteEdgeMappingsByResourceIds(
executor,
edgeChildWorkspaceId,
sourceIsParent ? 'child' : 'parent',
targets
)
}
async function deleteEdgeMappingsByResourceIds(
tx: DbOrTx,
childWorkspaceId: string,
side: 'parent' | 'child',
pairs: Array<{ resourceType: ForkResourceType; resourceId: string }>
): Promise<void> {
if (pairs.length === 0) return
const idsByType = new Map<ForkResourceType, string[]>()
for (const { resourceType, childResourceId } of pairs) {
for (const { resourceType, resourceId } of pairs) {
const list = idsByType.get(resourceType)
if (list) list.push(childResourceId)
else idsByType.set(resourceType, [childResourceId])
if (list) list.push(resourceId)
else idsByType.set(resourceType, [resourceId])
}
const resourceColumn =
side === 'parent'
? workspaceForkResourceMap.parentResourceId
: workspaceForkResourceMap.childResourceId
const conditions = Array.from(idsByType, ([resourceType, ids]) =>
and(
eq(workspaceForkResourceMap.resourceType, resourceType),
inArray(workspaceForkResourceMap.childResourceId, ids)
)
and(eq(workspaceForkResourceMap.resourceType, resourceType), inArray(resourceColumn, ids))
)
await tx
.delete(workspaceForkResourceMap)
apps/sim/ee/workspace-forking/lib/promote/copy-unmapped.test.ts
+75 -64
View File
@@ -9,22 +9,19 @@ import {
import type { DbOrTx } from '@/lib/db/types'
const {
mockUpsertEdgeMappings,
mockDeleteEdgeMappingsByChildResources,
mockPersistCopiedResourceMappings,
mockCopyForkResourceContainers,
mockPlanForkMappedKbDocumentCopies,
mockPlanForkFileCopies,
} = vi.hoisted(() => ({
mockUpsertEdgeMappings: vi.fn(),
mockDeleteEdgeMappingsByChildResources: vi.fn(),
mockPersistCopiedResourceMappings: vi.fn(),
mockCopyForkResourceContainers: vi.fn(),
mockPlanForkMappedKbDocumentCopies: vi.fn(),
mockPlanForkFileCopies: vi.fn(),
}))
vi.mock('@/ee/workspace-forking/lib/mapping/mapping-store', () => ({
upsertEdgeMappings: mockUpsertEdgeMappings,
deleteEdgeMappingsByChildResources: mockDeleteEdgeMappingsByChildResources,
persistCopiedResourceMappings: mockPersistCopiedResourceMappings,
resourceTypeToForkKind: vi.fn(),
}))
@@ -40,14 +37,12 @@ vi.mock('@/ee/workspace-forking/lib/copy/copy-files', () => ({
}))
import type { ForkEdge } from '@/ee/workspace-forking/lib/lineage/lineage'
import type { ForkMappingUpsert } from '@/ee/workspace-forking/lib/mapping/mapping-store'
import {
augmentForkResolver,
buildPromoteCopySelection,
copyPromoteUnmappedResources,
FORK_COPYABLE_KIND_TO_SELECTION_KEY,
hasPromoteCopySelection,
persistPromoteCopiedMappings,
} from '@/ee/workspace-forking/lib/promote/copy-unmapped'
import { isForkCopyableKind } from '@/ee/workspace-forking/lib/promote/promote-plan'
import type { ForkRemapKind } from '@/ee/workspace-forking/lib/remap/remap-references'
@@ -237,51 +232,6 @@ describe('augmentForkResolver', () => {
})
})
describe('persistPromoteCopiedMappings', () => {
const tx = {} as DbOrTx
const entry: ForkMappingUpsert = {
resourceType: 'knowledge_base',
parentResourceId: 'src-kb',
childResourceId: 'dst-kb',
}
beforeEach(() => {
vi.clearAllMocks()
})
it('pull keeps the source(parent)->target(child) orientation as-is', async () => {
await persistPromoteCopiedMappings(tx, 'edge-child', 'user-1', 'pull', [entry])
expect(mockUpsertEdgeMappings).toHaveBeenCalledWith(tx, 'edge-child', 'user-1', [entry])
expect(mockDeleteEdgeMappingsByChildResources).not.toHaveBeenCalled()
})
it('push swaps to target(parent)->source(child) and deletes the prior row keyed on the source child', async () => {
await persistPromoteCopiedMappings(tx, 'edge-child', 'user-1', 'push', [entry])
// Delete keys on the source child resource (the swapped child id = the original parent id).
expect(mockDeleteEdgeMappingsByChildResources).toHaveBeenCalledWith(tx, 'edge-child', [
{ resourceType: 'knowledge_base', childResourceId: 'src-kb' },
])
// The swap flips parent/child: the new copy (dst) becomes the parent side on push.
expect(mockUpsertEdgeMappings).toHaveBeenCalledWith(tx, 'edge-child', 'user-1', [
{ resourceType: 'knowledge_base', parentResourceId: 'dst-kb', childResourceId: 'src-kb' },
])
})
it('push skips an entry with a null child id (the narrowing guard, no bogus mapping)', async () => {
await persistPromoteCopiedMappings(tx, 'edge-child', 'user-1', 'push', [
{ resourceType: 'knowledge_base', parentResourceId: 'src-kb', childResourceId: null },
])
expect(mockDeleteEdgeMappingsByChildResources).not.toHaveBeenCalled()
expect(mockUpsertEdgeMappings).not.toHaveBeenCalled()
})
it('returns without writing when there are no entries', async () => {
await persistPromoteCopiedMappings(tx, 'edge-child', 'user-1', 'push', [])
expect(mockDeleteEdgeMappingsByChildResources).not.toHaveBeenCalled()
expect(mockUpsertEdgeMappings).not.toHaveBeenCalled()
})
})
describe('copyPromoteUnmappedResources - files + folder content-refs', () => {
const tx = {} as DbOrTx
// Only edge.childWorkspaceId is read by the copy path.
@@ -320,6 +270,46 @@ describe('copyPromoteUnmappedResources - files + folder content-refs', () => {
})
})
it('threads push orientation through the shared container and mapping boundaries', async () => {
await copyPromoteUnmappedResources({
tx,
edge,
sourceWorkspaceId: 'child-source-ws',
targetWorkspaceId: 'parent-target-ws',
direction: 'push',
userId: 'user-1',
now: new Date(),
selection: {
customTools: [],
skills: [],
tables: [],
knowledgeBases: [],
files: [],
mcpServers: [],
},
workflowIdMap: new Map(),
folderIdMap: new Map(),
resolver: () => null,
resolveBlockId,
referencedDocumentIds: [],
})
expect(mockCopyForkResourceContainers).toHaveBeenCalledWith(
expect.objectContaining({
documentMappingContext: {
edgeChildWorkspaceId: 'edge-child',
sourceIsParent: false,
},
})
)
expect(mockPersistCopiedResourceMappings).toHaveBeenCalledWith(
expect.objectContaining({
edgeChildWorkspaceId: 'edge-child',
sourceIsParent: false,
})
)
})
it('copies selected files (keyMap + blobTasks), persists the file mapping, and threads file + folder content-ref maps', async () => {
mockPlanForkFileCopies.mockResolvedValue({
keyMap: new Map([['workspace/SRC/a.png', 'workspace/DST/a.png']]),
@@ -351,6 +341,7 @@ describe('copyPromoteUnmappedResources - files + folder content-refs', () => {
tables: [],
knowledgeBases: [],
files: ['workspace/SRC/a.png'],
mcpServers: [],
},
workflowIdMap: new Map(),
folderIdMap: new Map([['fld-src', 'fld-dst']]),
@@ -371,13 +362,19 @@ describe('copyPromoteUnmappedResources - files + folder content-refs', () => {
)
// The file mapping is persisted (pull keeps source(parent)->target(child) orientation) so a
// re-sync resolves the copy instead of re-copying it.
expect(mockUpsertEdgeMappings).toHaveBeenCalledWith(tx, 'edge-child', 'user-1', [
{
resourceType: 'file',
parentResourceId: 'workspace/SRC/a.png',
childResourceId: 'workspace/DST/a.png',
},
])
expect(mockPersistCopiedResourceMappings).toHaveBeenCalledWith({
executor: tx,
edgeChildWorkspaceId: 'edge-child',
userId: 'user-1',
sourceIsParent: true,
entries: [
{
resourceType: 'file',
parentResourceId: 'workspace/SRC/a.png',
childResourceId: 'workspace/DST/a.png',
},
],
})
// The folder map AND the file key/id maps reach the in-content rewriter.
expect(result.contentRefMaps.folders).toEqual({ 'fld-src': 'fld-dst' })
expect(result.contentRefMaps.fileKeys).toEqual({ 'workspace/SRC/a.png': 'workspace/DST/a.png' })
@@ -430,6 +427,7 @@ describe('copyPromoteUnmappedResources - files + folder content-refs', () => {
tables: ['tbl-unref'],
knowledgeBases: [],
files: [],
mcpServers: [],
},
workflowIdMap: new Map(),
folderIdMap: new Map(),
@@ -438,9 +436,15 @@ describe('copyPromoteUnmappedResources - files + folder content-refs', () => {
referencedDocumentIds: [],
})
expect(mockUpsertEdgeMappings).toHaveBeenCalledWith(tx, 'edge-child', 'user-1', [
{ resourceType: 'table', parentResourceId: 'tbl-unref', childResourceId: 'tbl-copy' },
])
expect(mockPersistCopiedResourceMappings).toHaveBeenCalledWith({
executor: tx,
edgeChildWorkspaceId: 'edge-child',
userId: 'user-1',
sourceIsParent: true,
entries: [
{ resourceType: 'table', parentResourceId: 'tbl-unref', childResourceId: 'tbl-copy' },
],
})
})
it('threads the plan-provided referencedDocumentIds into both doc-copy paths (no in-tx re-scan)', async () => {
@@ -478,10 +482,17 @@ describe('copyPromoteUnmappedResources - files + folder content-refs', () => {
// The promote-built block-id resolver reaches the table remap unchanged, so copied
// tables' workflow-group outputs use the persisted-pair ids, not the derive.
resolveBlockId,
documentMappingContext: {
edgeChildWorkspaceId: 'edge-child',
sourceIsParent: true,
},
})
)
expect(mockPlanForkMappedKbDocumentCopies).toHaveBeenCalledWith(
expect.objectContaining({ referencedDocumentIds: ['doc-1', 'doc-2'] })
expect.objectContaining({
referencedDocumentIds: ['doc-1', 'doc-2'],
now: expect.any(Date),
})
)
})
})
apps/sim/ee/workspace-forking/lib/promote/copy-unmapped.ts
+13 -48
View File
@@ -16,10 +16,9 @@ import {
} from '@/ee/workspace-forking/lib/copy/copy-resources'
import type { ForkEdge } from '@/ee/workspace-forking/lib/lineage/lineage'
import {
deleteEdgeMappingsByChildResources,
type ForkMappingUpsert,
persistCopiedResourceMappings,
resourceTypeToForkKind,
upsertEdgeMappings,
} from '@/ee/workspace-forking/lib/mapping/mapping-store'
import type { ForkBlockIdResolver } from '@/ee/workspace-forking/lib/remap/block-identity'
import type {
@@ -230,6 +229,10 @@ export async function copyPromoteUnmappedResources(params: {
// rewritten through the same plan resolver that remaps subblock-value env refs.
resolveEnvName: (key) => resolver('env-var', key),
resolveBlockId,
documentMappingContext: {
edgeChildWorkspaceId: edge.childWorkspaceId,
sourceIsParent: direction === 'pull',
},
})
// Copy the selected workspace files (keyed by storage key) - metadata inserts in the tx, blob
@@ -259,6 +262,7 @@ export async function copyPromoteUnmappedResources(params: {
resolver,
referencedDocumentIds,
alreadyCopiedSourceDocIds: new Set(containerDocMap.keys()),
now,
})
result.contentPlan.documents.push(...mappedKbDocs.documents)
@@ -272,11 +276,13 @@ export async function copyPromoteUnmappedResources(params: {
childResourceId: child,
})
)
await persistPromoteCopiedMappings(tx, edge.childWorkspaceId, userId, direction, [
...result.mappingEntries,
...fileMappingEntries,
...mappedKbDocs.mappingEntries,
])
await persistCopiedResourceMappings({
executor: tx,
edgeChildWorkspaceId: edge.childWorkspaceId,
userId,
sourceIsParent: direction === 'pull',
entries: [...result.mappingEntries, ...fileMappingEntries, ...mappedKbDocs.mappingEntries],
})
const copyIdMapByKind = new Map<ForkRemapKind, Map<string, string>>()
for (const [resourceType, sourceToTarget] of result.idMap) {
@@ -313,44 +319,3 @@ export async function copyPromoteUnmappedResources(params: {
blobTasks: fileResult.blobTasks,
}
}
/**
* Persist the copied resources' id mappings for the edge. The copy returns entries oriented
* source(parent)->target(child); a pull matches that orientation directly (fill-null upsert), a
* push swaps it (the parent side is the new TARGET) and first drops any prior row keyed on the
* source child resource so a changed target can't leak a second mapping.
*/
export async function persistPromoteCopiedMappings(
tx: DbOrTx,
childWorkspaceId: string,
userId: string,
direction: 'push' | 'pull',
entries: ForkMappingUpsert[]
): Promise<void> {
if (entries.length === 0) return
if (direction === 'pull') {
await upsertEdgeMappings(tx, childWorkspaceId, userId, entries)
return
}
// Push: re-key on the source child resource. Skip any entry with a null child id (copy entries
// always carry one; the guard narrows the type so neither the swap nor the delete needs a cast).
// After the swap every childResourceId is the original (non-null) parent id, keyed for the
// delete-then-insert that prevents a changed target from leaking a second mapping.
const swapped: ForkMappingUpsert[] = []
const deleteKeys: Array<{
resourceType: ForkMappingUpsert['resourceType']
childResourceId: string
}> = []
for (const entry of entries) {
if (entry.childResourceId == null) continue
swapped.push({
resourceType: entry.resourceType,
parentResourceId: entry.childResourceId,
childResourceId: entry.parentResourceId,
})
deleteKeys.push({ resourceType: entry.resourceType, childResourceId: entry.parentResourceId })
}
if (swapped.length === 0) return
await deleteEdgeMappingsByChildResources(tx, childWorkspaceId, deleteKeys)
await upsertEdgeMappings(tx, childWorkspaceId, userId, swapped)
}
apps/sim/executor/execution/block-executor.test.ts
+84 -8
View File
@@ -5,6 +5,7 @@ import { loggerMock } from '@sim/testing'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import { clearLargeValueCacheForTests } from '@/lib/execution/payloads/cache'
import { isLargeArrayManifest } from '@/lib/execution/payloads/large-array-manifest-metadata'
import { isLargeValueRef } from '@/lib/execution/payloads/large-value-ref'
import { projectTraceSpansForSecrets } from '@/lib/logs/execution/trace-secret-projection'
import { buildTraceSpans } from '@/lib/logs/execution/trace-spans/trace-spans'
import { BlockType, EDGE } from '@/executor/constants'
@@ -146,6 +147,64 @@ describe('BlockExecutor', () => {
})
})
it('carries complete encrypted candidates through large-output compaction', async () => {
const block = createBlock()
const workflow: SerializedWorkflow = {
version: '1',
blocks: [block],
connections: [],
loops: {},
parallels: {},
}
const state = new ExecutionState()
const resolver = new VariableResolver(workflow, {}, state)
const onBlockComplete = vi.fn(async () => {})
const registry = new ResolvedSecretTraceRegistry(
[{ name: 'API_KEY', plaintext: 'secret-value', encryptedValue: 'encrypted-secret' }],
{ userId: 'user-1', workspaceId: 'workspace-1' }
)
const handler: BlockHandler = {
canHandle: () => true,
execute: async (blockContext) => {
blockContext.resolvedSecretTraceRegistry?.recordResolved('API_KEY', 'secret-value')
return {
result: {
huge: 'p'.repeat(9 * 1024 * 1024),
public: 'ok',
secret: 'secret-value',
},
}
},
}
const executor = new BlockExecutor([handler], resolver, { onBlockComplete }, state)
const ctx = createContext(state)
ctx.resolvedSecretTraceRegistry = registry
await executor.execute(ctx, createNode(block), block)
await vi.waitFor(() => expect(onBlockComplete).toHaveBeenCalledOnce())
const storedOutput = state.getBlockOutput(block.id)
const storedResult = storedOutput?.result as Record<string, unknown>
expect(isLargeValueRef(storedResult.huge)).toBe(true)
expect(storedResult.public).toBe('ok')
expect(storedResult.secret).toBe('secret-value')
const expectedProvenance = {
version: 1,
complete: true,
entries: [{ name: 'API_KEY', encryptedValue: 'encrypted-secret' }],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
}
expect(state.getBlockState(block.id)?.resolvedSecretTraceProvenance).toEqual(expectedProvenance)
expect(onBlockComplete.mock.calls[0]?.[3]?.resolvedSecretTraceProvenance).toEqual(
expectedProvenance
)
expect(onBlockComplete.mock.calls[0]?.[3]?.displayResolvedSecretTraceProvenance).toEqual(
expectedProvenance
)
expect(JSON.stringify(expectedProvenance)).not.toContain('secret-value')
})
it('persists stable outer-branch aliases for completed parallel branch outputs', async () => {
const block = createBlock()
const workflow: SerializedWorkflow = {
@@ -327,7 +386,7 @@ describe('BlockExecutor', () => {
it('projects lifecycle callback diagnostics without changing block execution', async () => {
const secret = 'lifecycle-secret-7f3a91'
const startError = new Error(`start failed ${secret} __var_API_KEY`)
const startError = new Error('start failed __var_API_KEY')
const completionError = new Error(`completion failed ${secret} __sim_code_6_binding_0`)
const block = createBlock()
const workflow: SerializedWorkflow = {
@@ -341,7 +400,15 @@ describe('BlockExecutor', () => {
const resolver = new VariableResolver(workflow, {}, state)
const output = { result: `raw ${secret}` }
const executor = new BlockExecutor(
[{ canHandle: () => true, execute: async () => output }],
[
{
canHandle: () => true,
execute: async (blockContext) => {
blockContext.resolvedSecretTraceRegistry?.recordResolved('API_KEY', secret)
return output
},
},
],
resolver,
{
workspaceId: 'workspace-1',
@@ -382,8 +449,7 @@ describe('BlockExecutor', () => {
expect.objectContaining({
blockId: block.id,
blockType: BlockType.FUNCTION,
errorType: 'error',
hasStack: true,
error: 'completion failed {{API_KEY}} [RUNTIME_BINDING]',
})
)
})
@@ -392,7 +458,7 @@ describe('BlockExecutor', () => {
expect.objectContaining({
blockId: block.id,
blockType: BlockType.FUNCTION,
error: 'start failed {{API_KEY}} {{API_KEY}}',
error: 'start failed [REDACTED_SECRET]',
})
)
const loggerPayload = JSON.stringify(executionLogger?.warn.mock.calls)
@@ -400,7 +466,7 @@ describe('BlockExecutor', () => {
expect(loggerPayload).not.toContain(secret)
expect(loggerPayload).not.toContain('__var_')
expect(loggerPayload).not.toContain('__sim_')
expect(startError.message).toContain(secret)
expect(startError.message).toContain('__var_API_KEY')
expect(completionError.message).toContain(secret)
})
@@ -795,7 +861,8 @@ describe('BlockExecutor', () => {
const resolver = new VariableResolver(workflow, {}, state)
const handler: BlockHandler = {
canHandle: () => true,
execute: async () => {
execute: async (blockContext) => {
blockContext.resolvedSecretTraceRegistry?.recordResolved('API_KEY', secret)
throw new Error(rawError)
},
}
@@ -900,10 +967,17 @@ describe('BlockExecutor streaming pump', () => {
failAfterText?: string
streamError?: Error
onFullContent?: (content: string) => void | Promise<void>
resolvedSecret?: { name: string; value: string }
}): BlockHandler {
return {
canHandle: () => true,
execute: async () => {
execute: async (blockContext) => {
if (options.resolvedSecret) {
blockContext.resolvedSecretTraceRegistry?.recordResolved(
options.resolvedSecret.name,
options.resolvedSecret.value
)
}
const timeSegment: Record<string, unknown> = {
type: 'model',
name: 'claude-test',
@@ -1031,6 +1105,7 @@ describe('BlockExecutor streaming pump', () => {
const handler = createAgentEventsStreamingHandler({
failAfterText: 'partial',
streamError: rawError,
resolvedSecret: { name: 'API_KEY', value: secret },
})
const { executor, block, state } = createExecutor(handler)
const ctx = createContext(state)
@@ -1079,6 +1154,7 @@ describe('BlockExecutor streaming pump', () => {
onFullContent: async () => {
throw callbackError
},
resolvedSecret: { name: 'API_KEY', value: secret },
})
const { executor, block, state } = createExecutor(handler)
block.config.params = { responseFormat: 'json' }
apps/sim/executor/execution/block-executor.ts
+3 -2
View File
@@ -101,8 +101,9 @@ export class BlockExecutor {
}
const parentResolvedSecretTraceRegistry = ctx.resolvedSecretTraceRegistry
const blockResolvedSecretTraceRegistry =
parentResolvedSecretTraceRegistry?.forkForToolInputValues([])
const blockResolvedSecretTraceRegistry = parentResolvedSecretTraceRegistry?.forkForInputPaths(
[]
)
const blockCtx = blockResolvedSecretTraceRegistry
? { ...ctx, resolvedSecretTraceRegistry: blockResolvedSecretTraceRegistry }
: ctx
apps/sim/executor/execution/engine.test.ts
+5 -3
View File
@@ -1009,10 +1009,12 @@ describe('ExecutionEngine', () => {
startNode.outgoingEdges.set('edge1', { target: 'error-node' })
const dag = createMockDAG([startNode, errorNode])
const registry = new ResolvedSecretTraceRegistry([
{ name: 'API_KEY', plaintext: secret, encryptedValue: 'encrypted-api-key' },
])
registry.recordResolved('API_KEY', secret)
const context = createMockContext({
resolvedSecretTraceRegistry: new ResolvedSecretTraceRegistry([
{ name: 'API_KEY', plaintext: secret, encryptedValue: 'encrypted-api-key' },
]),
resolvedSecretTraceRegistry: registry,
})
const edgeManager = createMockEdgeManager((node) => {
if (node.id === 'start') return ['error-node']
apps/sim/executor/execution/types.ts
+4 -4
View File
@@ -106,7 +106,7 @@ export interface SerializableExecutionState {
workflowVariableResolvedSecretTraceProvenance?: Record<string, ResolvedSecretTraceProvenanceV1>
/** Exact-value provenance for the persisted workflow input. Absence means legacy/untracked. */
workflowInputResolvedSecretTraceProvenance?: ResolvedSecretTraceProvenanceV1
/** Exact-value provenance for the persisted terminal output. Absence means legacy/untracked. */
/** Encrypted candidates for the persisted terminal output. Absence means legacy/untracked. */
finalOutputResolvedSecretTraceProvenance?: ResolvedSecretTraceProvenanceV1
/** Presence distinguishes current checkpoints from legacy states that predate provenance. */
resolvedSecretTraceCheckpointVersion?: 1
@@ -170,11 +170,11 @@ export interface BlockCompletionCallbackData {
input?: unknown
output: NormalizedBlockOutput
/**
* Encrypted provenance filtered to this exact block output. Internal durable
* consumers use it when the raw output crosses a storage boundary.
* Encrypted candidates active in this block call. Internal durable consumers
* filter them against the exact value that crosses a storage boundary.
*/
resolvedSecretTraceProvenance?: ResolvedSecretTraceProvenanceV1
/** Internal display-only provenance filtered to this callback's input/output envelope. */
/** Internal encrypted candidates filtered against the display envelope during projection. */
displayResolvedSecretTraceProvenance?: ResolvedSecretTraceProvenanceV1
executionTime: number
startedAt: string
apps/sim/executor/handlers/agent/agent-handler.test.ts
+998
View File
File diff suppressed because it is too large Load Diff
apps/sim/executor/handlers/agent/agent-handler.ts
+835 -75
View File
File diff suppressed because it is too large Load Diff
apps/sim/executor/handlers/agent/memory.test.ts
+51 -14
View File
@@ -266,20 +266,34 @@ describe('Memory', () => {
expect(mockRedactObjectStrings).toHaveBeenCalledOnce()
})
it('does not write memory when projection fails closed', async () => {
it('persists raw memory with unknown lineage when provenance is unavailable', async () => {
const registry = new ResolvedSecretTraceRegistry()
registry.markIncomplete()
const appendMessage = vi
.spyOn(memoryService as any, 'appendMessage')
.mockResolvedValue(undefined)
await expect(
memoryService.appendToMemory(createContext(registry) as never, inputs, {
role: 'user',
content: 'possibly secret',
})
).rejects.toThrow('Memory content could not be safely projected')
expect(appendMessage).not.toHaveBeenCalled()
const message = { role: 'user' as const, content: 'possibly secret' }
await memoryService.appendToMemory(createContext(registry) as never, inputs, message)
expect(appendMessage).toHaveBeenCalledWith('workspace-1', 'conversation-1', message, {
status: 'unknown',
})
})
it('seeds raw memory with unknown lineage when provenance is unavailable', async () => {
const registry = new ResolvedSecretTraceRegistry()
registry.markIncomplete()
const seedMemoryRecord = vi
.spyOn(memoryService as any, 'seedMemoryRecord')
.mockResolvedValue(undefined)
const message = { role: 'assistant' as const, content: 'possibly secret' }
await memoryService.seedMemory(createContext(registry) as never, inputs, [message])
expect(seedMemoryRecord).toHaveBeenCalledWith('workspace-1', 'conversation-1', [message], {
status: 'unknown',
})
})
it('preserves legacy stored messages when no current resolution activated the value', async () => {
@@ -351,7 +365,7 @@ describe('Memory', () => {
}
const projected = (memoryService as any).projectMessageForModel(
createContext(registry),
registry,
message
) as Message
@@ -387,9 +401,16 @@ describe('Memory', () => {
messages: [message],
provenance: {
status: 'exact',
entries: [{ name: 'TOKEN', encryptedValue: 'ciphertext' }],
entries: [
{
name: 'TOKEN',
encryptedValue: 'ciphertext',
sourceValueHash: hashDurableSecretProvenanceValue(message),
},
],
},
})
mockDecryptSecret.mockResolvedValue({ decrypted: secret })
const [fetched] = await memoryService.fetchMemoryMessages(
createContext(registry) as never,
inputs
@@ -402,7 +423,7 @@ describe('Memory', () => {
)
it.each(['name', 'functionName', 'toolCallId', 'toolName'] as const)(
'rejects an active resolved secret in the %s control field',
'does not plaintext-scan the %s control field',
(field) => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'TOKEN', plaintext: 'control-secret', encryptedValue: 'ciphertext' },
@@ -428,12 +449,28 @@ describe('Memory', () => {
],
} as Message
expect(() =>
(memoryService as any).projectMessageForModel(createContext(registry), message)
).toThrow('Memory content could not be safely projected')
expect((memoryService as any).projectMessageForModel(registry, message)).toEqual(message)
}
)
it('does not project unrelated active secrets into legacy memory', async () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'TOKEN', plaintext: 'x', encryptedValue: 'ciphertext' },
])
registry.recordResolved('TOKEN', 'x')
vi.spyOn(memoryService as any, 'fetchMemory').mockResolvedValueOnce({
messages: [{ role: 'assistant', content: 'Box' }],
provenance: { status: 'exact', entries: [] },
})
const messages = await memoryService.fetchMemoryMessages(
createContext(registry) as never,
inputs
)
expect(messages).toEqual([{ role: 'assistant', content: 'Box' }])
})
it('does not activate provenance from a message dropped by the selected window', async () => {
const oldSecretMessage: Message = { role: 'user', content: 'same-value' }
const retainedPublicMessage: Message = { role: 'assistant', content: 'same-value' }
apps/sim/executor/handlers/agent/memory.ts
+19 -78
View File
@@ -24,7 +24,7 @@ import {
projectResolvedSecretModelContent,
projectResolvedSecretModelJsonStrings,
} from '@/executor/utils/resolved-secret-content-projection'
import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
import { PROVIDER_DEFINITIONS } from '@/providers/models'
const logger = createLogger('Memory')
@@ -85,7 +85,21 @@ export class Memory {
throw new Error('Memory content could not be safely projected')
}
return messages.map((message) => this.projectMessageForModel(ctx, message))
return Promise.all(
messages.map(async (message) => {
const messageProvenance = filterDurableSecretProvenanceBySourceValues(selectedProvenance, [
message,
])
const modelRegistry = new ResolvedSecretTraceRegistry(
[],
ctx.resolvedSecretTraceRegistry?.exportProvenance().scope
)
if (!(await importDurableSecretProvenance(modelRegistry, messageProvenance, message))) {
throw new Error('Memory content could not be safely projected')
}
return this.projectMessageForModel(modelRegistry, message)
})
)
}
private captureMessagesProvenance(
@@ -122,9 +136,6 @@ export class Memory {
const provenance = ctx.resolvedSecretTraceRegistry
? this.captureMessagesProvenance(ctx.resolvedSecretTraceRegistry, [message])
: undefined
if (provenance?.status === 'unknown') {
throw new Error('Memory content could not be safely projected')
}
await this.appendMessage(workspaceId, key, message, provenance)
@@ -172,9 +183,6 @@ export class Memory {
const provenance = ctx.resolvedSecretTraceRegistry
? this.captureMessagesProvenance(ctx.resolvedSecretTraceRegistry, messagesToStore)
: undefined
if (provenance?.status === 'unknown') {
throw new Error('Memory content could not be safely projected')
}
await this.seedMemoryRecord(workspaceId, key, messagesToStore, provenance)
logger.debug('Seeded memory', {
@@ -204,23 +212,7 @@ export class Memory {
}
}
private projectMessageForModel(ctx: ExecutionContext, message: Message): Message {
const controlValues = this.readModelControlValues(message)
const controlProjection = projectResolvedSecretModelContent(
controlValues,
ctx.resolvedSecretTraceRegistry
)
if (
!controlProjection.safe ||
!Array.isArray(controlProjection.value) ||
controlProjection.value.length !== controlValues.length ||
controlProjection.value.some(
(value, index) => typeof value !== 'string' || value !== controlValues[index]
)
) {
throw new Error('Memory content could not be safely projected')
}
private projectMessageForModel(registry: ResolvedSecretTraceRegistry, message: Message): Message {
const functionArguments = this.readFunctionCallArguments(message.function_call)
const toolArguments = message.tool_calls?.map((toolCall) => {
if (!isPlainRecord(toolCall)) {
@@ -228,13 +220,10 @@ export class Memory {
}
return this.readFunctionCallArguments(toolCall.function)
})
const contentProjection = projectResolvedSecretModelContent(
message.content,
ctx.resolvedSecretTraceRegistry
)
const contentProjection = projectResolvedSecretModelContent(message.content, registry)
const argumentProjection = projectResolvedSecretModelJsonStrings(
[functionArguments, ...(toolArguments ?? [])],
ctx.resolvedSecretTraceRegistry
registry
)
if (
!contentProjection.safe ||
@@ -306,54 +295,6 @@ export class Memory {
return functionCall.arguments
}
private readModelControlValues(message: Message): string[] {
if (!isPlainRecord(message)) {
throw new Error('Memory content could not be safely projected')
}
const controls: string[] = []
for (const key of ['name', 'tool_call_id'] as const) {
if (!(key in message)) continue
const value = message[key]
if (value !== undefined && value !== null) {
if (typeof value !== 'string') {
throw new Error('Memory content could not be safely projected')
}
controls.push(value)
}
}
if (message.function_call !== undefined && message.function_call !== null) {
if (!isPlainRecord(message.function_call)) {
throw new Error('Memory content could not be safely projected')
}
const name = message.function_call.name
if (name !== undefined && name !== null) {
if (typeof name !== 'string') {
throw new Error('Memory content could not be safely projected')
}
controls.push(name)
}
}
for (const toolCall of message.tool_calls ?? []) {
if (!isPlainRecord(toolCall)) {
throw new Error('Memory content could not be safely projected')
}
for (const value of [
toolCall.id,
isPlainRecord(toolCall.function) ? toolCall.function.name : undefined,
]) {
if (value !== undefined && value !== null) {
if (typeof value !== 'string') {
throw new Error('Memory content could not be safely projected')
}
controls.push(value)
}
}
}
return controls
}
private requireWorkspaceId(ctx: ExecutionContext): string {
if (!ctx.workspaceId) {
throw new Error('workspaceId is required for memory operations')
apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts
+111 -3
View File
@@ -32,7 +32,11 @@ vi.mock('@/lib/model-router/resolve', () => ({
SIM_AUTO_SYSTEM_PREAMBLE: 'Sim auto system preamble',
}))
import { PRIVATE_MODEL_INPUT_PROVENANCE_HEADER } from '@/lib/execution/model-input-provenance'
import {
PRIVATE_MODEL_INPUT_PROVENANCE_HEADER,
PRIVATE_MODEL_INPUT_STATE_HEADER,
PROJECTED_MODEL_INPUT_PATHS_V1,
} from '@/lib/execution/model-input-provenance'
import {
RESOLVED_SECRET_PROVENANCE_FIELD,
RESOLVED_SECRET_PROVENANCE_METADATA_V1,
@@ -223,8 +227,18 @@ describe('EvaluatorBlockHandler', () => {
encryptedValue: 'encrypted-evaluator-credential',
},
])
registry.recordResolved('CONTENT_SECRET', contentSecret)
registry.recordResolved('METRIC_SECRET', metricSecret)
registry.recordResolvedAtInputPath('CONTENT_SECRET', contentSecret, ['content'])
registry.recordResolvedInputProjection(['content'], contentSecret, '{{CONTENT_SECRET}}')
registry.recordResolvedAtInputPath('METRIC_SECRET', metricSecret, [
'metrics',
'0',
'description',
])
registry.recordResolvedInputProjection(
['metrics', '0', 'description'],
metricSecret,
'{{METRIC_SECRET}}'
)
registry.recordResolved('API_KEY', credentialSecret)
mockContext.resolvedSecretTraceRegistry = registry
@@ -246,6 +260,9 @@ describe('EvaluatorBlockHandler', () => {
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBe(
RESOLVED_SECRET_PROVENANCE_METADATA_V1
)
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBe(
PROJECTED_MODEL_INPUT_PATHS_V1
)
expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({
version: 1,
complete: true,
@@ -263,6 +280,96 @@ describe('EvaluatorBlockHandler', () => {
expect(requestBody.apiKey).toBe(credentialSecret)
})
it('projects every model-bound metric leaf and maps the score back to the raw metric name', async () => {
const rawMetric = {
name: 'private-metric-key',
description: 'private metric instructions',
range: { min: 'private minimum', max: 'private maximum' },
}
const projectedMetric = {
name: '{{METRIC_NAME_SECRET}}',
description: '{{METRIC_DESCRIPTION_SECRET}}',
range: { min: '{{METRIC_MIN_SECRET}}', max: '{{METRIC_MAX_SECRET}}' },
}
const secrets = [
{
name: 'METRIC_NAME_SECRET',
plaintext: rawMetric.name,
encryptedValue: 'encrypted-metric-name',
path: ['metrics', '0', 'name'],
projected: projectedMetric.name,
},
{
name: 'METRIC_DESCRIPTION_SECRET',
plaintext: rawMetric.description,
encryptedValue: 'encrypted-metric-description',
path: ['metrics', '0', 'description'],
projected: projectedMetric.description,
},
{
name: 'METRIC_MIN_SECRET',
plaintext: rawMetric.range.min,
encryptedValue: 'encrypted-metric-min',
path: ['metrics', '0', 'range', 'min'],
projected: projectedMetric.range.min,
},
{
name: 'METRIC_MAX_SECRET',
plaintext: rawMetric.range.max,
encryptedValue: 'encrypted-metric-max',
path: ['metrics', '0', 'range', 'max'],
projected: projectedMetric.range.max,
},
] as const
const registry = new ResolvedSecretTraceRegistry([
...secrets.map(({ name, plaintext, encryptedValue }) => ({
name,
plaintext,
encryptedValue,
})),
{ name: 'UNUSED_SECRET', plaintext: 'x', encryptedValue: 'encrypted-unused' },
])
for (const secret of secrets) {
registry.recordResolvedAtInputPath(secret.name, secret.plaintext, secret.path)
registry.recordResolvedInputProjection(secret.path, secret.plaintext, secret.projected)
}
mockContext.resolvedSecretTraceRegistry = registry
mockFetch.mockResolvedValueOnce({
ok: true,
json: () =>
Promise.resolve({
content: JSON.stringify({ [projectedMetric.name.toLowerCase()]: 7 }),
model: 'mock-model',
tokens: {},
cost: 0,
}),
})
const result = await handler.execute(mockContext, mockBlock, {
content: 'Public x remains public.',
metrics: [rawMetric],
model: 'gpt-4o',
apiKey: 'test-api-key',
})
const requestBody = JSON.parse(mockFetch.mock.calls[0][1].body)
const serializedRequest = JSON.stringify(requestBody)
for (const secret of secrets) {
expect(serializedRequest).not.toContain(secret.plaintext)
expect(requestBody.systemPrompt).toContain(secret.projected)
}
expect(requestBody.systemPrompt).toContain('Public x remains public.')
expect(requestBody.responseFormat.schema.properties).toEqual({
[projectedMetric.name.toLowerCase()]: { type: 'number' },
})
expect(
requestBody[RESOLVED_SECRET_PROVENANCE_FIELD].entries
.map((entry: { name: string }) => entry.name)
.sort()
).toEqual(secrets.map((secret) => secret.name).sort())
expect(result).toMatchObject({ [rawMetric.name.toLowerCase()]: 7 })
})
it('keeps the evaluator request shape when no provenance registry exists', async () => {
await handler.execute(mockContext, mockBlock, {
content: 'Public evaluator content',
@@ -275,6 +382,7 @@ describe('EvaluatorBlockHandler', () => {
const requestBody = JSON.parse(request.body)
expect(Object.hasOwn(requestBody, RESOLVED_SECRET_PROVENANCE_FIELD)).toBe(false)
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBeNull()
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBeNull()
})
it('resolves sim-auto before executing evaluator and preserves its public identity', async () => {
apps/sim/executor/handlers/evaluator/evaluator-handler.ts
+57 -16
View File
@@ -2,6 +2,8 @@ import { createLogger } from '@sim/logger'
import {
addModelInputProvenanceToRequest,
createModelInputProvenanceRequestMetadata,
markModelInputProjected,
projectResolvedModelInput,
} from '@/lib/execution/model-input-provenance'
import {
type AutoRoutingResult,
@@ -16,10 +18,12 @@ import type { BlockHandler, ExecutionContext } from '@/executor/types'
import { buildAPIUrl, buildAuthHeaders, extractAPIErrorMessage } from '@/executor/utils/http'
import { isJSONString, parseJSON, stringifyJSON } from '@/executor/utils/json'
import { projectResolvedSecretDiagnosticError } from '@/executor/utils/resolved-secret-content-projection'
import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
import type {
ResolvedSecretInputPath,
ResolvedSecretTraceRegistry,
} from '@/executor/utils/resolved-secret-trace-registry'
import { resolveVertexCredential } from '@/executor/utils/vertex-credential'
import { resolveProxiedModelCost } from '@/providers/cost-policy'
import { collectProviderModelInputProvenanceValues } from '@/providers/model-input-provenance'
import { isAutoModel, SIM_AUTO_MODEL_ID } from '@/providers/models'
import type { ProviderRequest } from '@/providers/types'
import { getProviderFromModel } from '@/providers/utils'
@@ -51,8 +55,6 @@ export class EvaluatorBlockHandler implements BlockHandler {
bedrockRegion: inputs.bedrockRegion,
}
const processedContent = this.processContent(inputs.content)
let systemPromptObj: { systemPrompt: string; responseFormat: any } = {
systemPrompt: '',
responseFormat: null,
@@ -64,15 +66,43 @@ export class EvaluatorBlockHandler implements BlockHandler {
} else {
metrics = []
}
const modelInputPaths: ResolvedSecretInputPath[] = [
['content'],
...metrics.flatMap((_, index) => [
['metrics', String(index), 'name'],
['metrics', String(index), 'description'],
['metrics', String(index), 'range', 'min'],
['metrics', String(index), 'range', 'max'],
]),
]
const modelInputProjection = projectResolvedModelInput(
ctx.resolvedSecretTraceRegistry,
{ content: inputs.content, metrics: inputs.metrics },
modelInputPaths
)
if (!modelInputProjection.complete) {
throw new Error('Evaluator model input could not be safely projected')
}
const processedContent = this.processContent(modelInputProjection.value.content)
const projectedMetrics = Array.isArray(modelInputProjection.value.metrics)
? modelInputProjection.value.metrics
: []
const metricDescriptions = metrics
.filter((m: any) => m?.name && m.range)
.map((m: any) => `"${m.name}" (${m.range.min}-${m.range.max}): ${m.description || ''}`)
.map((metric: any, index: number) => ({ metric, projected: projectedMetrics[index] }))
.filter(({ metric, projected }) =>
Boolean(metric?.name && metric.range && projected?.name && projected.range)
)
.map(
({ projected }) =>
`"${projected.name}" (${projected.range.min}-${projected.range.max}): ${projected.description || ''}`
)
.join('\n')
const responseProperties: Record<string, any> = {}
metrics.forEach((m: any, metricIndex: number) => {
if (m?.name) {
responseProperties[m.name.toLowerCase()] = { type: 'number' }
const projectedMetric = projectedMetrics[metricIndex]
if (m?.name && projectedMetric?.name) {
responseProperties[projectedMetric.name.toLowerCase()] = { type: 'number' }
} else {
logger.warn('Skipping invalid metric entry during response format generation', {
metricIndex,
@@ -96,7 +126,10 @@ export class EvaluatorBlockHandler implements BlockHandler {
schema: {
type: 'object',
properties: responseProperties,
required: metrics.filter((m: any) => m?.name).map((m: any) => m.name.toLowerCase()),
required: metrics.flatMap((m: any, metricIndex: number) => {
const projectedName = projectedMetrics[metricIndex]?.name
return m?.name && projectedName ? [projectedName.toLowerCase()] : []
}),
additionalProperties: false,
},
strict: true,
@@ -181,14 +214,16 @@ export class EvaluatorBlockHandler implements BlockHandler {
}
const headers = new Headers(await buildAuthHeaders(ctx.userId))
const modelInputMetadata = createModelInputProvenanceRequestMetadata(
modelInputProjection.registry,
modelInputPaths
)
const requestBody = addModelInputProvenanceToRequest(
{ provider: providerId, ...providerRequest },
headers,
createModelInputProvenanceRequestMetadata(
ctx.resolvedSecretTraceRegistry,
collectProviderModelInputProvenanceValues(providerRequest, providerId)
)
modelInputMetadata
)
if (modelInputMetadata) markModelInputProjected(headers)
const response = await fetch(url.toString(), {
method: 'POST',
headers,
@@ -207,7 +242,7 @@ export class EvaluatorBlockHandler implements BlockHandler {
ctx.resolvedSecretTraceRegistry
)
const metricScores = this.extractMetricScores(parsedContent, inputs.metrics)
const metricScores = this.extractMetricScores(parsedContent, metrics, projectedMetrics)
const inputTokens = result.tokens?.input || result.tokens?.prompt || DEFAULTS.TOKENS.PROMPT
const outputTokens =
@@ -297,7 +332,8 @@ export class EvaluatorBlockHandler implements BlockHandler {
private extractMetricScores(
parsedContent: Record<string, any>,
metrics: any
metrics: any,
projectedMetrics: any
): Record<string, number> {
const metricScores: Record<string, number> = {}
let validMetrics: any[]
@@ -316,6 +352,7 @@ export class EvaluatorBlockHandler implements BlockHandler {
return metricScores
}
const validProjectedMetrics = Array.isArray(projectedMetrics) ? projectedMetrics : []
validMetrics.forEach((metric: any, metricIndex: number) => {
if (!metric?.name) {
logger.warn('Skipping invalid metric entry', {
@@ -325,7 +362,11 @@ export class EvaluatorBlockHandler implements BlockHandler {
return
}
const score = this.findMetricScore(parsedContent, metric.name)
const projectedName = validProjectedMetrics[metricIndex]?.name
const score = this.findMetricScore(
parsedContent,
typeof projectedName === 'string' && projectedName ? projectedName : metric.name
)
metricScores[metric.name.toLowerCase()] = score
})
apps/sim/executor/handlers/function/function-handler.test.ts
+43
View File
@@ -200,6 +200,49 @@ describe('FunctionBlockHandler', () => {
}
})
it('forwards an explicit selected secret scope without changing legacy unset blocks', async () => {
await handler.execute(mockContext, mockBlock, {
code: 'return {{API_KEY}}',
secretScope: 'selected',
mountedSecrets: [' API_KEY ', 42, 'SECOND_KEY', '', 'API_KEY'],
})
expect(mockExecuteTool).toHaveBeenCalledWith(
'function_execute',
expect.objectContaining({
secretScope: 'selected',
mountedSecrets: ['API_KEY', 'SECOND_KEY'],
}),
{ executionContext: mockContext }
)
vi.clearAllMocks()
mockExecuteTool.mockResolvedValue({ success: true, output: { result: 'Success' } })
await handler.execute(mockContext, mockBlock, { code: 'return {{API_KEY}}' })
const legacyParams = mockExecuteTool.mock.calls[0][1]
expect(legacyParams).not.toHaveProperty('secretScope')
expect(legacyParams).not.toHaveProperty('mountedSecrets')
})
it('fails closed for an invalid explicit secret scope', async () => {
await handler.execute(mockContext, mockBlock, {
code: 'return {{API_KEY}}',
secretScope: 'invalid',
mountedSecrets: ['API_KEY'],
})
expect(mockExecuteTool).toHaveBeenCalledWith(
'function_execute',
expect.objectContaining({
secretScope: 'selected',
mountedSecrets: [],
}),
{ executionContext: mockContext }
)
})
it('should handle execution errors from the tool', async () => {
const inputs = { code: 'throw new Error("Code failed");' }
const errorResult = { success: false, error: 'Function execution failed: Code failed' }
apps/sim/executor/handlers/function/function-handler.ts
+9
View File
@@ -1,3 +1,4 @@
import { normalizeSecretMountPolicy } from '@/lib/copilot/secret-mount-policy'
import { getRemainingExecutionMs } from '@/lib/core/execution-limits'
import {
normalizeRecord,
@@ -68,6 +69,13 @@ export class FunctionBlockHandler implements BlockHandler {
? remainingExecutionMs
: Math.min(requestedTimeout, remainingExecutionMs)
)
const secretMountPolicy =
inputs.secretScope === undefined
? undefined
: normalizeSecretMountPolicy({
secretScope: inputs.secretScope,
mountedSecrets: inputs.mountedSecrets,
})
const toolParams = {
code: codeContent,
@@ -75,6 +83,7 @@ export class FunctionBlockHandler implements BlockHandler {
language: inputs.language || DEFAULT_CODE_LANGUAGE,
timeout,
...(inputs.sandboxId ? { sandboxId: inputs.sandboxId } : {}),
...(secretMountPolicy ?? {}),
envVars: normalizeStringRecord(ctx.environmentVariables),
workflowVariables: normalizeWorkflowVariables(ctx.workflowVariables),
blockData: {},
apps/sim/executor/handlers/generic/generic-handler.test.ts
+392
View File
@@ -1,14 +1,19 @@
import '@sim/testing/mocks/executor'
import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest'
import { KnowledgeBlock } from '@/blocks/blocks/knowledge'
import { getBlock } from '@/blocks/index'
import { BlockType } from '@/executor/constants'
import { GenericBlockHandler } from '@/executor/handlers/generic/generic-handler'
import type { ExecutionContext } from '@/executor/types'
import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
import type { SerializedBlock } from '@/serializer/types'
import { executeTool } from '@/tools'
import { selectKnowledgeDocumentWriteSecretProvenance } from '@/tools/knowledge/secret-provenance'
import type { ToolConfig } from '@/tools/types'
import { getTool } from '@/tools/utils'
const mockGetBlock = vi.mocked(getBlock)
const mockGetTool = vi.mocked(getTool)
const mockExecuteTool = executeTool as Mock
@@ -60,6 +65,7 @@ describe('GenericBlockHandler', () => {
// Reset mocks using vi
vi.clearAllMocks()
mockGetBlock.mockReturnValue(undefined)
// Set up mockGetTool to return mockTool
mockGetTool.mockImplementation((toolId) => {
@@ -98,6 +104,392 @@ describe('GenericBlockHandler', () => {
expect(result).toEqual(expectedOutput)
})
it('preserves exact secret provenance when block params rename a selected input', async () => {
mockTool.request.modelInput = {
mode: 'private-provenance',
inputPaths: () => [['filePath']],
}
mockGetBlock.mockReturnValue({
tools: {
access: ['some_custom_tool'],
config: {
tool: () => 'some_custom_tool',
params: (params: Record<string, unknown>) => ({
filePath: String(params.document).trim(),
}),
},
},
inputs: {
document: { type: 'string', description: 'Document URL' },
},
} as never)
const registry = new ResolvedSecretTraceRegistry([
{
name: 'FILE_SECRET',
plaintext: 'secret-url',
encryptedValue: 'encrypted-file-secret',
},
])
registry.recordResolvedAtInputPath('FILE_SECRET', 'secret-url', ['document'])
registry.recordResolvedInputProjection(['document'], ' secret-url ', ' {{FILE_SECRET}} ')
mockContext.resolvedSecretTraceRegistry = registry
await handler.execute(mockContext, mockBlock, { document: ' secret-url ' })
expect(mockExecuteTool).toHaveBeenCalledWith(
'some_custom_tool',
expect.objectContaining({
document: ' secret-url ',
filePath: 'secret-url',
}),
{ executionContext: mockContext }
)
expect(registry.exportCommittedProvenanceForInputPaths([['filePath']])).toMatchObject({
complete: true,
entries: [{ name: 'FILE_SECRET', encryptedValue: 'encrypted-file-secret' }],
})
})
it('traces each secret path without letting secret-valued controls change another path', async () => {
mockTool.request.modelInput = {
mode: 'private-provenance',
inputPaths: () => [['input']],
}
mockGetBlock.mockReturnValue({
tools: {
access: ['some_custom_tool'],
config: {
tool: () => 'some_custom_tool',
params: (params: Record<string, unknown>) =>
params.operation === 'deep_research' ? { input: params.research_input } : {},
},
},
inputs: {
operation: { type: 'string', description: 'Operation' },
research_input: { type: 'string', description: 'Research input' },
},
} as never)
const registry = new ResolvedSecretTraceRegistry([
{
name: 'OPERATION',
plaintext: 'deep_research',
encryptedValue: 'encrypted-operation',
},
{ name: 'QUERY', plaintext: 'secret query', encryptedValue: 'encrypted-query' },
])
registry.recordResolvedAtInputPath('OPERATION', 'deep_research', ['operation'])
registry.recordResolvedInputProjection(['operation'], 'deep_research', '{{OPERATION}}')
registry.recordResolvedAtInputPath('QUERY', 'secret query', ['research_input'])
registry.recordResolvedInputProjection(['research_input'], 'secret query', '{{QUERY}}')
mockContext.resolvedSecretTraceRegistry = registry
await handler.execute(mockContext, mockBlock, {
operation: 'deep_research',
research_input: 'secret query',
})
expect(mockExecuteTool).toHaveBeenCalledWith(
'some_custom_tool',
expect.objectContaining({
operation: 'deep_research',
research_input: 'secret query',
input: 'secret query',
}),
{ executionContext: mockContext }
)
expect(registry.exportCommittedProvenanceForInputPaths([['input']])).toMatchObject({
complete: true,
entries: [{ name: 'QUERY', encryptedValue: 'encrypted-query' }],
})
})
it('preserves exact table leaf provenance across legacy unquoted JSON placeholders', async () => {
mockTool.request.secretProvenance = {
request: () => [{ key: 'data', inputPaths: [['data']] }],
}
mockGetBlock.mockReturnValue({
tools: {
access: ['some_custom_tool'],
config: {
tool: () => 'some_custom_tool',
params: (params: Record<string, unknown>) => ({
data: typeof params.data === 'string' ? JSON.parse(params.data) : params.data,
}),
},
},
inputs: {
data: { type: 'json', description: 'Row data' },
},
} as never)
const registry = new ResolvedSecretTraceRegistry([
{ name: '1BOOLEAN_SECRET', plaintext: 'true', encryptedValue: 'encrypted-boolean' },
{ name: 'UNUSED', plaintext: 'true', encryptedValue: 'encrypted-unused' },
])
registry.recordResolvedAtInputPath('1BOOLEAN_SECRET', 'true', ['data'])
registry.recordResolvedInputProjection(
['data'],
'{"secret":true,"public":true}',
'{"secret":{{1BOOLEAN_SECRET}},"public":true}'
)
mockContext.resolvedSecretTraceRegistry = registry
await handler.execute(mockContext, mockBlock, {
data: '{"secret":true,"public":true}',
})
expect(mockExecuteTool).toHaveBeenCalledWith(
'some_custom_tool',
expect.objectContaining({ data: { secret: true, public: true } }),
{ executionContext: mockContext }
)
expect(registry.exportCommittedProvenanceForInputPaths([['data', 'secret']])).toMatchObject({
complete: true,
entries: [{ name: '1BOOLEAN_SECRET', encryptedValue: 'encrypted-boolean' }],
})
expect(registry.exportCommittedProvenanceForInputPaths([['data', 'public']])).toMatchObject({
complete: true,
entries: [],
})
})
it('normalizes legacy JSON-string knowledge tags without mutating inputs or overbinding tag names', async () => {
mockTool.request.secretProvenance = {
request: selectKnowledgeDocumentWriteSecretProvenance,
}
mockGetBlock.mockReturnValue(KnowledgeBlock)
mockBlock.metadata = { id: 'knowledge', name: 'Knowledge' }
const documentTags = '[{"tagName":"team","value":"support"}]'
const projectedTags = '[{"tagName":"team","value":"{{TAG_VALUE}}"}]'
const inputs = {
operation: 'create_document',
knowledgeBaseId: 'kb-1',
name: 'doc.md',
content: 'content',
documentTags,
}
const registry = new ResolvedSecretTraceRegistry([
{ name: 'TAG_VALUE', plaintext: 'support', encryptedValue: 'encrypted-tag' },
])
registry.recordResolvedAtInputPath('TAG_VALUE', 'support', ['documentTags'])
registry.recordResolvedInputProjection(['documentTags'], documentTags, projectedTags)
mockContext.resolvedSecretTraceRegistry = registry
await handler.execute(mockContext, mockBlock, inputs)
expect(inputs.documentTags).toBe(documentTags)
expect(mockExecuteTool).toHaveBeenCalledWith(
'some_custom_tool',
expect.objectContaining({
documentTags: [{ tagName: 'team', value: 'support' }],
}),
{ executionContext: mockContext }
)
expect(
registry.exportCommittedProvenanceForInputPaths([['documentTags', '0', 'tagName']])
).toMatchObject({ complete: true, entries: [] })
expect(
registry.exportCommittedProvenanceForInputPaths([['documentTags', '0', 'value']])
).toMatchObject({
complete: true,
entries: [{ name: 'TAG_VALUE', encryptedValue: 'encrypted-tag' }],
})
})
it('preserves a whole structured secret without changing the raw parsed value', async () => {
mockTool.request.secretProvenance = {
request: () => [{ key: 'data', inputPaths: [['data']] }],
}
mockGetBlock.mockReturnValue({
tools: {
access: ['some_custom_tool'],
config: {
tool: () => 'some_custom_tool',
params: (params: Record<string, unknown>) => ({
data: typeof params.data === 'string' ? JSON.parse(params.data) : params.data,
}),
},
},
inputs: {
data: { type: 'json', description: 'Row data' },
},
} as never)
const rawStructuredSecret = '{"nested":"value","url":"https://example.com/data","count":1}'
const registry = new ResolvedSecretTraceRegistry([
{
name: 'JSON_SECRET',
plaintext: rawStructuredSecret,
encryptedValue: 'encrypted-json',
},
])
registry.recordResolvedAtInputPath('JSON_SECRET', rawStructuredSecret, ['data'])
registry.recordResolvedInputProjection(['data'], rawStructuredSecret, '{{JSON_SECRET}}')
mockContext.resolvedSecretTraceRegistry = registry
await handler.execute(mockContext, mockBlock, { data: rawStructuredSecret })
expect(mockExecuteTool).toHaveBeenCalledWith(
'some_custom_tool',
expect.objectContaining({
data: { nested: 'value', url: 'https://example.com/data', count: 1 },
}),
{ executionContext: mockContext }
)
expect(registry.exportCommittedProvenanceForInputPaths([['data', 'nested']])).toMatchObject({
complete: true,
entries: [{ name: 'JSON_SECRET', encryptedValue: 'encrypted-json' }],
})
expect(registry.exportCommittedProvenanceForInputPaths([['data', 'count']])).toMatchObject({
complete: true,
entries: [{ name: 'JSON_SECRET', encryptedValue: 'encrypted-json' }],
})
expect(registry.exportCommittedProvenanceForInputPaths([['data', 'url']])).toMatchObject({
complete: true,
entries: [{ name: 'JSON_SECRET', encryptedValue: 'encrypted-json' }],
})
})
it('preserves structured message roles while projecting only model-visible content', async () => {
const parseMessages = (value: unknown) => {
const parsed = typeof value === 'string' ? JSON.parse(value) : value
if (!Array.isArray(parsed)) throw new Error('Messages must be an array')
return parsed.map((message) => {
if (
!message ||
typeof message !== 'object' ||
!['user', 'assistant', 'system'].includes(String(message.role))
) {
throw new Error('Invalid message role')
}
return { role: String(message.role), content: String(message.content) }
})
}
mockTool.request.modelInput = {
mode: 'project',
select: (params) => ({
messages: parseMessages(params.messages).map((message) => message.content),
}),
applyProjected: (selectedParams, projectedSelection) => ({
messages: parseMessages(selectedParams.messages).map((message, index) => ({
...message,
content: (projectedSelection.messages as unknown[])[index],
})),
}),
}
mockGetBlock.mockReturnValue({
tools: {
access: ['some_custom_tool'],
config: {
tool: () => 'some_custom_tool',
params: (params: Record<string, unknown>) => ({
messages: parseMessages(params.messages),
}),
},
},
inputs: {
messages: { type: 'json', description: 'Messages' },
},
} as never)
const rawMessages = '[{"role":"user","content":"hello"}]'
const registry = new ResolvedSecretTraceRegistry([
{ name: 'MESSAGES', plaintext: rawMessages, encryptedValue: 'encrypted-messages' },
])
registry.recordResolvedAtInputPath('MESSAGES', rawMessages, ['messages'])
registry.recordResolvedInputProjection(['messages'], rawMessages, '{{MESSAGES}}')
mockContext.resolvedSecretTraceRegistry = registry
await handler.execute(mockContext, mockBlock, { messages: rawMessages })
expect(mockExecuteTool).toHaveBeenCalledWith(
'some_custom_tool',
expect.objectContaining({ messages: [{ role: 'user', content: 'hello' }] }),
{ executionContext: mockContext }
)
expect(
registry.exportCommittedProvenanceForInputPaths([['messages', '0', 'role']])
).toMatchObject({ complete: true, entries: [] })
expect(
registry.exportCommittedProvenanceForInputPaths([['messages', '0', 'content']])
).toMatchObject({
complete: true,
entries: [{ name: 'MESSAGES', encryptedValue: 'encrypted-messages' }],
})
})
it('preserves raw file execution while binding whole serialized descriptors to the file boundary', async () => {
mockTool.params.audioFile = { type: 'file' }
mockTool.params.audioUrl = { type: 'string' }
mockTool.request.modelInput = {
mode: 'private-provenance',
inputPaths: () => [['audioUrl']],
}
mockGetBlock.mockReturnValue({
tools: {
access: ['some_custom_tool'],
config: {
tool: () => 'some_custom_tool',
params: (params: Record<string, unknown>) => {
const file =
typeof params.audioFile === 'string' ? JSON.parse(params.audioFile) : params.audioFile
if (!file || typeof file !== 'object' || !String(file.url).startsWith('https://')) {
throw new Error('A valid HTTPS audio file is required')
}
return { audioUrl: String(file.url), audioFile: undefined }
},
},
},
inputs: {
audioFile: { type: 'json', description: 'Audio file' },
},
} as never)
const rawFile = '{"name":"audio.mp3","size":4,"url":"https://files.example/audio.mp3"}'
const registry = new ResolvedSecretTraceRegistry([
{ name: 'FILE', plaintext: rawFile, encryptedValue: 'encrypted-file' },
])
registry.recordResolvedAtInputPath('FILE', rawFile, ['audioFile'])
registry.recordResolvedInputProjection(['audioFile'], rawFile, '{{FILE}}')
mockContext.resolvedSecretTraceRegistry = registry
await handler.execute(mockContext, mockBlock, { audioFile: rawFile })
expect(mockExecuteTool).toHaveBeenCalledWith(
'some_custom_tool',
expect.objectContaining({
audioFile: undefined,
audioUrl: 'https://files.example/audio.mp3',
}),
{ executionContext: mockContext }
)
expect(registry.isComplete()).toBe(true)
expect(registry.exportCommittedProvenanceForInputPaths([['audioUrl']])).toMatchObject({
complete: true,
entries: [{ name: 'FILE', encryptedValue: 'encrypted-file' }],
})
})
it('does not replay block transforms for configured but unused secrets', async () => {
mockTool.request.modelInput = {
mode: 'project',
select: (params) => ({ param1: params.param1 }),
}
const transform = vi.fn((params: Record<string, unknown>) => params)
mockGetBlock.mockReturnValue({
tools: {
access: ['some_custom_tool'],
config: { tool: () => 'some_custom_tool', params: transform },
},
inputs: {
param1: { type: 'string', description: 'Value' },
},
} as never)
mockContext.resolvedSecretTraceRegistry = new ResolvedSecretTraceRegistry([
{ name: 'UNUSED', plaintext: 'value1', encryptedValue: 'encrypted-unused' },
])
await handler.execute(mockContext, mockBlock, { param1: 'value1' })
expect(transform).toHaveBeenCalledTimes(1)
})
it('should throw error if the associated tool is not found', async () => {
const inputs = { param1: 'value' }
apps/sim/executor/handlers/generic/generic-handler.ts
+209
View File
@@ -1,15 +1,154 @@
import { isDeepStrictEqual } from 'node:util'
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { isPlainRecord } from '@sim/utils/object'
import { getBlock } from '@/blocks/index'
import { isMcpTool } from '@/executor/constants'
import type { BlockHandler, ExecutionContext } from '@/executor/types'
import { readStatusCode } from '@/executor/utils/errors'
import { prepareResolvedSecretProjectedInputs } from '@/executor/utils/resolved-secret-input-projection'
import type { ResolvedSecretInputPath } from '@/executor/utils/resolved-secret-trace-registry'
import type { SerializedBlock } from '@/serializer/types'
import { executeTool } from '@/tools'
import type { ToolConfig } from '@/tools/types'
import { getTool } from '@/tools/utils'
const logger = createLogger('GenericBlockHandler')
interface BlockBoundaryPaths {
paths: ResolvedSecretInputPath[]
requiredProjectionRoots: Set<string>
}
function selectBlockBoundaryPaths(
tool: ToolConfig,
params: Record<string, unknown>
): BlockBoundaryPaths | undefined {
try {
const paths: ResolvedSecretInputPath[] = []
const requiredProjectionRoots = new Set<string>()
const modelInput = tool.request.modelInput
if (modelInput?.mode === 'project') {
const selected = modelInput.select(params)
if (!isPlainRecord(selected)) return undefined
for (const key of Object.keys(selected)) {
requiredProjectionRoots.add(key)
paths.push([key])
}
const privateInputPaths = modelInput.privateInputPaths?.(params) ?? []
paths.push(...privateInputPaths)
for (const path of privateInputPaths) {
if (path[0]) requiredProjectionRoots.add(path[0])
}
} else if (modelInput?.mode === 'private-provenance') {
const privateInputPaths = modelInput.inputPaths(params)
paths.push(...privateInputPaths)
for (const path of privateInputPaths) {
if (path[0]) requiredProjectionRoots.add(path[0])
}
}
const opaqueInputPaths = tool.request.opaqueModelInput?.inputPaths(params) ?? []
paths.push(...opaqueInputPaths)
for (const path of opaqueInputPaths) {
if (path[0]) requiredProjectionRoots.add(path[0])
}
for (const selection of tool.request.secretProvenance?.request?.(params) ?? []) {
paths.push(...selection.inputPaths)
for (const path of selection.inputPaths) {
if (path[0]) requiredProjectionRoots.add(path[0])
}
}
const uniquePaths = new Map<string, ResolvedSecretInputPath>()
for (const path of paths) {
if (path.length > 0) uniquePaths.set(JSON.stringify(path), path)
}
return { paths: [...uniquePaths.values()], requiredProjectionRoots }
} catch {
return undefined
}
}
function canonicalPlaceholder(value: unknown): string | undefined {
if (typeof value !== 'string') return undefined
const match = /^\{\{([A-Za-z0-9_]+)\}\}$/.exec(value.trim())
return match ? value.trim() : undefined
}
function isFileBoundaryPath(tool: ToolConfig, path: ResolvedSecretInputPath): boolean {
return Boolean(path[0] && tool.params[path[0]]?.type === 'file')
}
function projectScalarLeaves(
value: unknown,
placeholder: string
): { value: unknown; projectedLeaves: number } | undefined {
if (value === null || typeof value !== 'object') {
return { value: placeholder, projectedLeaves: 1 }
}
if (!Array.isArray(value) && !isPlainRecord(value)) return undefined
const root: unknown[] | Record<string, unknown> = Array.isArray(value) ? [] : {}
const pending: Array<{
source: unknown[] | Record<string, unknown>
target: unknown[] | Record<string, unknown>
}> = [{ source: value as unknown[] | Record<string, unknown>, target: root }]
const visited = new WeakSet<object>()
let projectedLeaves = 0
while (pending.length > 0) {
const { source, target } = pending.pop()!
if (visited.has(source)) return undefined
visited.add(source)
for (const [key, child] of Object.entries(source)) {
if (child !== null && typeof child === 'object') {
if (!Array.isArray(child) && !isPlainRecord(child)) return undefined
const projectedChild: unknown[] | Record<string, unknown> = Array.isArray(child) ? [] : {}
;(target as Record<string, unknown>)[key] = projectedChild
pending.push({
source: child as unknown[] | Record<string, unknown>,
target: projectedChild,
})
} else {
;(target as Record<string, unknown>)[key] = placeholder
projectedLeaves += 1
}
}
}
return { value: root, projectedLeaves }
}
function createStructuredModelProjection(
tool: ToolConfig,
finalInputs: Record<string, unknown>,
sourcePath: ResolvedSecretInputPath,
projectedSourceValue: unknown
): Record<string, unknown> | undefined {
const modelInput = tool.request.modelInput
const sourceKey = sourcePath.length === 1 ? sourcePath[0] : undefined
const placeholder = canonicalPlaceholder(projectedSourceValue)
if (modelInput?.mode !== 'project' || !modelInput.applyProjected || !sourceKey || !placeholder) {
return undefined
}
try {
const selected = modelInput.select(finalInputs)
if (!isPlainRecord(selected) || !Object.hasOwn(selected, sourceKey)) return undefined
const projectedValue = projectScalarLeaves(selected[sourceKey], placeholder)
if (!projectedValue || projectedValue.projectedLeaves === 0) return undefined
const projectedSelection = { ...selected, [sourceKey]: projectedValue.value }
const selectedParams = Object.fromEntries(
Object.keys(selected).map((key) => [key, finalInputs[key]])
)
const patch = modelInput.applyProjected(structuredClone(selectedParams), projectedSelection)
if (!isPlainRecord(patch)) return undefined
const projectedInputs = { ...finalInputs, ...patch }
if (!isDeepStrictEqual(modelInput.select(projectedInputs), projectedSelection)) return undefined
return projectedInputs
} catch {
return undefined
}
}
export class GenericBlockHandler implements BlockHandler {
canHandle(block: SerializedBlock): boolean {
return true
@@ -35,6 +174,8 @@ export class GenericBlockHandler implements BlockHandler {
const blockType = block.metadata?.id
if (blockType) {
const blockConfig = getBlock(blockType)
const registry = ctx.resolvedSecretTraceRegistry
if (blockConfig?.tools?.config?.params) {
const transformedParams = blockConfig.tools.config.params(inputs)
finalInputs = { ...inputs, ...transformedParams }
@@ -57,6 +198,74 @@ export class GenericBlockHandler implements BlockHandler {
}
}
}
const boundary = tool ? selectBlockBoundaryPaths(tool, finalInputs) : undefined
const projectedInputs =
boundary && boundary.paths.length > 0 && registry?.hasResolvedInputProjections()
? registry.projectResolvedInputSelections(inputs)
: undefined
if (projectedInputs?.complete === false) registry?.markIncomplete()
if (projectedInputs?.complete && boundary && tool && registry) {
for (const projection of projectedInputs.values) {
const preserveFileDescriptorGrammar =
isFileBoundaryPath(tool, projection.path) ||
boundary.paths.some((path) => isFileBoundaryPath(tool, path))
let projectedFinalInputs = prepareResolvedSecretProjectedInputs(
projection.value,
blockConfig?.inputs,
inputs,
{ preserveFileDescriptorGrammar }
)
try {
if (blockConfig?.tools?.config?.params) {
projectedFinalInputs = {
...projectedFinalInputs,
...blockConfig.tools.config.params(projectedFinalInputs),
}
}
} catch {
const structuredProjection = createStructuredModelProjection(
tool,
finalInputs,
projection.path,
projection.projectedValue
)
if (structuredProjection) {
registry.recordTransformedInputProjection(finalInputs, structuredProjection, {
targetPaths: boundary.paths,
})
continue
}
if (boundary.requiredProjectionRoots.has(projection.path[0])) {
registry.markIncomplete()
}
continue
}
if (blockConfig?.inputs) {
projectedFinalInputs = prepareResolvedSecretProjectedInputs(
projectedFinalInputs,
blockConfig.inputs,
finalInputs,
{ preserveFileDescriptorGrammar }
)
for (const [key, inputSchema] of Object.entries(blockConfig.inputs)) {
const value = projectedFinalInputs[key]
if (typeof value !== 'string' || value.trim().length === 0) continue
const inputType = typeof inputSchema === 'object' ? inputSchema.type : inputSchema
if (inputType !== 'json' && inputType !== 'array') continue
try {
projectedFinalInputs[key] = JSON.parse(value.trim())
} catch {}
}
}
registry.recordTransformedInputProjection(finalInputs, projectedFinalInputs, {
targetPaths: boundary.paths,
})
}
}
}
try {
apps/sim/executor/handlers/mothership/mothership-handler.test.ts
+549 -74
View File
@@ -5,8 +5,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import { BlockType } from '@/executor/constants'
import { MothershipBlockHandler } from '@/executor/handlers/mothership/mothership-handler'
import type { ExecutionContext, StreamingExecution } from '@/executor/types'
import { createResolvedSecretMatcher } from '@/executor/utils/resolved-secret-content-projection'
import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
import type { SerializedBlock } from '@/serializer/types'
const BILLING_ATTRIBUTION = {
@@ -121,21 +120,16 @@ async function readStreamText(stream: ReadableStream): Promise<string> {
}
function createTraceRegistryMock(): ResolvedSecretTraceRegistry & {
getModelEgressRevision: ReturnType<typeof vi.fn>
getModelEgressSnapshot: ReturnType<typeof vi.fn>
importProvenanceForValue: ReturnType<typeof vi.fn>
markIncomplete: ReturnType<typeof vi.fn>
} {
return {
getModelEgressRevision: vi.fn().mockReturnValue(0),
getModelEgressSnapshot: vi
.fn()
.mockReturnValue({ complete: true, matches: [], matcher: undefined }),
importProvenanceForValue: vi.fn().mockResolvedValue(true),
markIncomplete: vi.fn(),
} as unknown as ResolvedSecretTraceRegistry & {
getModelEgressRevision: ReturnType<typeof vi.fn>
getModelEgressSnapshot: ReturnType<typeof vi.fn>
const registry = new ResolvedSecretTraceRegistry()
return Object.assign(registry, {
importProvenanceForValue: vi
.spyOn(registry, 'importProvenanceForValue')
.mockResolvedValue(true),
markIncomplete: vi.spyOn(registry, 'markIncomplete'),
}) as ResolvedSecretTraceRegistry & {
importProvenanceForValue: ReturnType<typeof vi.fn>
markIncomplete: ReturnType<typeof vi.fn>
}
@@ -299,19 +293,19 @@ describe('MothershipBlockHandler', () => {
expect(JSON.stringify(result)).not.toContain('encrypted-secret')
})
it('projects parent execution secrets before sending a Mothership prompt', async () => {
const registry = createTraceRegistryMock()
const matches = [
{
plaintext: 'cross-workspace-secret',
replacement: '[REDACTED_SECRET]',
},
]
registry.getModelEgressSnapshot.mockReturnValue({
complete: true,
matches,
matcher: createResolvedSecretMatcher(matches),
})
it('projects only secrets resolved at the Mothership prompt input path', async () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'PROMPT_SECRET', plaintext: 'prompt-secret', encryptedValue: 'prompt-ciphertext' },
{ name: 'UNUSED', plaintext: 'x', encryptedValue: 'unused-ciphertext' },
])
registry.recordResolvedAtInputPath('PROMPT_SECRET', 'prompt-secret', ['prompt'])
registry.recordResolvedInputProjection(
['prompt'],
'Use prompt-secret while Box stays unchanged',
'Use {{PROMPT_SECRET}} while Box stays unchanged'
)
registry.recordResolved('UNUSED', 'x')
vi.spyOn(registry, 'importProvenanceForValue').mockResolvedValue(true)
context.resolvedSecretTraceRegistry = registry
mockGenerateId
.mockReturnValueOnce('chat-uuid')
@@ -320,25 +314,45 @@ describe('MothershipBlockHandler', () => {
fetchMock.mockResolvedValue(createJsonResponse({ content: 'done', toolCalls: [] }))
await handler.execute(context, block, {
prompt: 'Use cross-workspace-secret and __var_FOREIGN',
prompt: 'Use prompt-secret while Box stays unchanged',
})
const [, options] = fetchMock.mock.calls[0] as [string, RequestInit]
const body = String(options.body)
expect(body).toContain('[REDACTED_SECRET]')
expect(body).not.toContain('cross-workspace-secret')
expect(body).not.toContain('prompt-secret')
expect(JSON.parse(body)).toMatchObject({
messages: [{ content: 'Use [REDACTED_SECRET] and __var_FOREIGN' }],
messages: [{ content: 'Use {{PROMPT_SECRET}} while Box stays unchanged' }],
})
})
it('drops a legacy response without poisoning later provenance', async () => {
it('preserves a headerless legacy JSON response without poisoning later calls', async () => {
const registry = createTraceRegistryMock()
context.resolvedSecretTraceRegistry = registry
mockGenerateId
.mockReturnValueOnce('chat-uuid')
.mockReturnValueOnce('message-uuid')
.mockReturnValueOnce('request-uuid')
fetchMock.mockResolvedValue(
new Response(
JSON.stringify({
content: 'unchanged output',
toolCalls: [],
}),
{ status: 200, headers: { 'Content-Type': 'application/json' } }
)
)
await expect(handler.execute(context, block, { prompt: 'Hello' })).resolves.toMatchObject({
content: 'unchanged output',
})
expect(registry.importProvenanceForValue).not.toHaveBeenCalled()
expect(registry.markIncomplete).not.toHaveBeenCalled()
})
it('rejects a headerless response that contains a partial private envelope', async () => {
const registry = createTraceRegistryMock()
context.resolvedSecretTraceRegistry = registry
fetchMock.mockResolvedValue(
new Response(
JSON.stringify({
@@ -351,11 +365,11 @@ describe('MothershipBlockHandler', () => {
)
await expect(handler.execute(context, block, { prompt: 'Hello' })).rejects.toThrow(
'does not support private provenance metadata'
'provenance metadata is invalid'
)
expect(registry.importProvenanceForValue).not.toHaveBeenCalled()
expect(registry.markIncomplete).not.toHaveBeenCalled()
expect(registry.markIncomplete).toHaveBeenCalled()
})
it('poisons provenance when a declared response omits its private field', async () => {
@@ -379,13 +393,55 @@ describe('MothershipBlockHandler', () => {
expect(registry.markIncomplete).toHaveBeenCalledOnce()
})
it('fails closed before the request when the model-egress registry is unavailable', async () => {
it('preserves legacy request and response behavior when no registry is available', async () => {
context.resolvedSecretTraceRegistry = undefined
fetchMock.mockResolvedValue(
new Response(JSON.stringify({ content: 'legacy output', toolCalls: [] }), {
status: 200,
headers: { 'Content-Type': 'application/json' },
})
)
await expect(handler.execute(context, block, { prompt: 'Hello' })).resolves.toMatchObject({
content: 'legacy output',
})
expect(context.resolvedSecretTraceRegistry).toBeUndefined()
})
it('rejects an explicitly mismatched response metadata version', async () => {
const registry = createTraceRegistryMock()
context.resolvedSecretTraceRegistry = registry
fetchMock.mockResolvedValue(
new Response(JSON.stringify({ content: 'unsafe output', toolCalls: [] }), {
status: 200,
headers: {
'Content-Type': 'application/json',
'x-sim-private-tool-metadata': 'resolved-secret-provenance-v2',
},
})
)
await expect(handler.execute(context, block, { prompt: 'Hello' })).rejects.toThrow(
'Mothership input could not be safely projected'
'provenance metadata is invalid'
)
expect(fetchMock).not.toHaveBeenCalled()
expect(registry.markIncomplete).toHaveBeenCalled()
})
it('preserves a headerless legacy upstream error without poisoning later calls', async () => {
const registry = createTraceRegistryMock()
context.resolvedSecretTraceRegistry = registry
fetchMock.mockResolvedValue(
new Response(JSON.stringify({ error: 'legacy upstream error' }), {
status: 502,
headers: { 'Content-Type': 'application/json' },
})
)
await expect(handler.execute(context, block, { prompt: 'Hello' })).rejects.toThrow(
'Sim execution failed: legacy upstream error'
)
expect(registry.importProvenanceForValue).not.toHaveBeenCalled()
expect(registry.markIncomplete).not.toHaveBeenCalled()
})
it('imports provenance from a terminal NDJSON error without forcing structural fallback', async () => {
@@ -468,6 +524,100 @@ describe('MothershipBlockHandler', () => {
expect(JSON.stringify(result.execution.output)).not.toContain('encrypted-secret')
})
it('preserves a headerless legacy NDJSON final result without poisoning later calls', async () => {
const registry = createTraceRegistryMock()
context.resolvedSecretTraceRegistry = registry
const encoder = new TextEncoder()
fetchMock.mockResolvedValue(
new Response(
new ReadableStream({
start(controller) {
controller.enqueue(
encoder.encode(
`${JSON.stringify({
type: 'final',
data: { content: 'legacy final', toolCalls: [] },
})}\n`
)
)
controller.close()
},
}),
{ headers: { 'Content-Type': 'application/x-ndjson; charset=utf-8' } }
)
)
await expect(handler.execute(context, block, { prompt: 'Hello' })).resolves.toMatchObject({
content: 'legacy final',
})
expect(registry.importProvenanceForValue).not.toHaveBeenCalled()
expect(registry.markIncomplete).not.toHaveBeenCalled()
})
it('preserves headerless legacy selected-output streaming without poisoning later calls', async () => {
const registry = createTraceRegistryMock()
context.resolvedSecretTraceRegistry = registry
context.stream = true
context.selectedOutputs = [`${block.id}_content`]
const encoder = new TextEncoder()
fetchMock.mockResolvedValue(
new Response(
new ReadableStream({
start(controller) {
controller.enqueue(
encoder.encode(`${JSON.stringify({ type: 'chunk', content: 'legacy chunk' })}\n`)
)
controller.enqueue(
encoder.encode(
`${JSON.stringify({
type: 'final',
data: { content: 'legacy final', toolCalls: [] },
})}\n`
)
)
controller.close()
},
}),
{ headers: { 'Content-Type': 'application/x-ndjson; charset=utf-8' } }
)
)
const result = (await handler.execute(context, block, {
prompt: 'Hello',
})) as StreamingExecution
await expect(readStreamText(result.stream)).resolves.toBe('legacy chunk')
expect(result.execution.output).toMatchObject({ content: 'legacy final' })
expect(registry.importProvenanceForValue).not.toHaveBeenCalled()
expect(registry.markIncomplete).not.toHaveBeenCalled()
})
it('surfaces a headerless legacy NDJSON terminal error without poisoning later calls', async () => {
const registry = createTraceRegistryMock()
context.resolvedSecretTraceRegistry = registry
const encoder = new TextEncoder()
fetchMock.mockResolvedValue(
new Response(
new ReadableStream({
start(controller) {
controller.enqueue(
encoder.encode(
`${JSON.stringify({ type: 'error', error: 'legacy terminal error' })}\n`
)
)
controller.close()
},
}),
{ headers: { 'Content-Type': 'application/x-ndjson; charset=utf-8' } }
)
)
await expect(handler.execute(context, block, { prompt: 'Hello' })).rejects.toThrow(
'Sim execution failed: legacy terminal error'
)
expect(registry.importProvenanceForValue).not.toHaveBeenCalled()
expect(registry.markIncomplete).not.toHaveBeenCalled()
})
it('forwards workflow and execution metadata with generated UUID ids', async () => {
mockGenerateId.mockReturnValueOnce('chat-uuid')
mockGenerateId.mockReturnValueOnce('message-uuid')
@@ -661,16 +811,12 @@ describe('MothershipBlockHandler', () => {
expect(body.contexts).toEqual([{ kind: 'skill', skillId: 'skill-1', label: 'sales-playbook' }])
})
it('projects proven model metadata without rewriting arbitrary attachment names or payloads', async () => {
it('does not scan arbitrary Mothership metadata, attachment names, or payloads', async () => {
const secret = 'boundary-secret'
const replacement = '{{API_KEY}}'
const registry = createTraceRegistryMock()
const matches = [{ plaintext: secret, replacement }]
registry.getModelEgressSnapshot.mockReturnValue({
complete: true,
matches,
matcher: createResolvedSecretMatcher(matches),
})
const registry = new ResolvedSecretTraceRegistry([
{ name: 'UNUSED_SECRET', plaintext: secret, encryptedValue: 'encrypted-unused-secret' },
])
vi.spyOn(registry, 'importProvenanceForValue').mockResolvedValue(true)
context.resolvedSecretTraceRegistry = registry
mockGenerateId
.mockReturnValueOnce('chat-uuid')
@@ -734,33 +880,366 @@ describe('MothershipBlockHandler', () => {
usageControl: 'force',
schema: {
type: 'object',
title: `Query ${replacement}`,
description: `Search using ${replacement}`,
title: `Query ${secret}`,
description: `Search using ${secret}`,
properties: {
query: { type: 'string', description: `Find ${replacement}` },
query: { type: 'string', description: `Find ${secret}` },
},
},
params: {
serverId: 'mcp-server-1',
toolName: 'search',
serverName: `Docs ${replacement}`,
serverName: `Docs ${secret}`,
},
},
])
expect(body.contexts).toEqual([
{ kind: 'skill', skillId: 'skill-1', label: `Playbook ${replacement}` },
{ kind: 'skill', skillId: 'skill-1', label: `Playbook ${secret}` },
])
const attachmentMetadata = {
type: body.fileAttachments[0].type,
})
it('projects only resolver-recorded model-visible MCP and skill metadata', async () => {
const secrets = [
{
name: 'MCP_SERVER_LABEL',
plaintext: 'private server label',
encryptedValue: 'encrypted-server-label',
path: ['tools', '0', 'params', 'serverName'],
raw: 'Docs private server label',
projected: 'Docs {{MCP_SERVER_LABEL}}',
},
{
name: 'MCP_SCHEMA_DESCRIPTION',
plaintext: 'private schema text',
encryptedValue: 'encrypted-schema-description',
path: ['tools', '0', 'schema', 'description'],
raw: 'Search private schema text for Box',
projected: 'Search {{MCP_SCHEMA_DESCRIPTION}} for Box',
},
{
name: 'SKILL_LABEL',
plaintext: 'private skill label',
encryptedValue: 'encrypted-skill-label',
path: ['skills', '0', 'name'],
raw: 'Playbook private skill label',
projected: 'Playbook {{SKILL_LABEL}}',
},
{
name: 'DISABLED_SERVER_ID',
plaintext: 'disabled-server-secret',
encryptedValue: 'encrypted-disabled-server-id',
path: ['tools', '1', 'params', 'serverId'],
raw: 'disabled-server-secret',
projected: '{{DISABLED_SERVER_ID}}',
},
] as const
const registry = new ResolvedSecretTraceRegistry([
...secrets.map(({ name, plaintext, encryptedValue }) => ({
name,
plaintext,
encryptedValue,
})),
{ name: 'UNUSED_SECRET', plaintext: 'x', encryptedValue: 'encrypted-unused' },
])
for (const secret of secrets) {
registry.recordResolvedAtInputPath(secret.name, secret.plaintext, secret.path)
registry.recordResolvedInputProjection(secret.path, secret.raw, secret.projected)
}
expect(
JSON.stringify({
messages: body.messages,
attachmentMetadata,
mcpTools: body.mcpTools,
contexts: body.contexts,
vi.spyOn(registry, 'importProvenanceForValue').mockResolvedValue(true)
context.resolvedSecretTraceRegistry = registry
mockGenerateId
.mockReturnValueOnce('chat-uuid')
.mockReturnValueOnce('message-uuid')
.mockReturnValueOnce('request-uuid')
fetchMock.mockResolvedValue(createJsonResponse({ content: 'done', toolCalls: [] }))
const tools = [
{
type: 'mcp',
params: {
serverId: 'mcp-server-1',
toolName: 'search',
serverName: 'Docs private server label',
},
schema: {
type: 'object',
description: 'Search private schema text for Box',
properties: { query: { type: 'string' } },
},
},
{
type: 'mcp',
usageControl: 'none',
params: { serverId: 'disabled-server-secret', toolName: 'disabled' },
},
]
const skills = [{ skillId: 'skill-1', name: 'Playbook private skill label' }]
await handler.execute(context, block, {
prompt: 'Use Box without changing it',
tools,
skills,
})
const [, options] = fetchMock.mock.calls[0] as [string, RequestInit]
const body = JSON.parse(String(options.body))
expect(body.mcpTools).toEqual([
{
type: 'mcp',
schema: {
type: 'object',
description: 'Search {{MCP_SCHEMA_DESCRIPTION}} for Box',
properties: { query: { type: 'string' } },
},
params: {
serverId: 'mcp-server-1',
toolName: 'search',
serverName: 'Docs {{MCP_SERVER_LABEL}}',
},
},
])
expect(body.contexts).toEqual([
{ kind: 'skill', skillId: 'skill-1', label: 'Playbook {{SKILL_LABEL}}' },
])
expect(JSON.stringify(body)).not.toContain('private server label')
expect(JSON.stringify(body)).not.toContain('private schema text')
expect(JSON.stringify(body)).not.toContain('private skill label')
expect(tools[0].params.serverName).toBe('Docs private server label')
expect(tools[0].schema.description).toBe('Search private schema text for Box')
expect(skills[0].name).toBe('Playbook private skill label')
})
it('rejects only an enabled structural identifier with exact resolver provenance', async () => {
const registry = new ResolvedSecretTraceRegistry([
{
name: 'MCP_SERVER_ID',
plaintext: 'resolved-server-id',
encryptedValue: 'encrypted-server-id',
},
])
registry.recordResolvedAtInputPath('MCP_SERVER_ID', 'resolved-server-id', [
'tools',
'0',
'params',
'serverId',
])
registry.recordResolvedInputProjection(
['tools', '0', 'params', 'serverId'],
'resolved-server-id',
'{{MCP_SERVER_ID}}'
)
context.resolvedSecretTraceRegistry = registry
await expect(
handler.execute(context, block, {
prompt: 'Use the selected tool',
tools: [
{
type: 'mcp',
params: { serverId: 'resolved-server-id', toolName: 'search' },
},
],
})
).not.toContain(secret)
).rejects.toThrow('Mothership structural model inputs cannot contain secret references')
expect(fetchMock).not.toHaveBeenCalled()
})
it('rejects a resolver-derived MCP enum under a property named description', async () => {
const registry = new ResolvedSecretTraceRegistry([
{
name: 'ENUM_VALUE',
plaintext: 'private-option',
encryptedValue: 'encrypted-option',
},
])
const inputPath = ['tools', '0', 'schema', 'properties', 'description', 'enum', '0'] as const
registry.recordResolvedAtInputPath('ENUM_VALUE', 'private-option', inputPath)
registry.recordResolvedInputProjection(inputPath, 'private-option', '{{ENUM_VALUE}}')
context.resolvedSecretTraceRegistry = registry
await expect(
handler.execute(context, block, {
prompt: 'Use the selected tool',
tools: [
{
type: 'mcp',
params: { serverId: 'server-1', toolName: 'search' },
schema: {
type: 'object',
properties: {
description: { type: 'string', enum: ['private-option'] },
},
},
},
],
})
).rejects.toThrow('Mothership structural model inputs cannot contain secret references')
expect(fetchMock).not.toHaveBeenCalled()
})
it('projects a resolver-recorded attachment name without changing file materialization', async () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'FILE_TOKEN', plaintext: 'x', encryptedValue: 'encrypted-file-token' },
])
registry.recordResolvedAtInputPath('FILE_TOKEN', 'x', ['files', '0', 'name'])
registry.recordResolvedInputProjection(
['files', '0', 'name'],
'report-x.txt',
'report-{{FILE_TOKEN}}.txt'
)
vi.spyOn(registry, 'importProvenanceForValue').mockResolvedValue(true)
context.resolvedSecretTraceRegistry = registry
mockGenerateId
.mockReturnValueOnce('chat-uuid')
.mockReturnValueOnce('message-uuid')
.mockReturnValueOnce('request-uuid')
const attachmentData = Buffer.from('ordinary bytes', 'utf8').toString('base64')
mockReadUserFileContent.mockResolvedValueOnce(attachmentData)
fetchMock.mockResolvedValue(createJsonResponse({ content: 'done', toolCalls: [] }))
await handler.execute(context, block, {
prompt: 'Read the attachment',
files: [
{
name: 'report-x.txt',
key: 'workspace/workspace-1/report.txt',
size: 32,
type: 'text/plain',
},
],
})
expect(mockReadUserFileContent).toHaveBeenCalledWith(
expect.objectContaining({
name: 'report-x.txt',
key: 'workspace/workspace-1/report.txt',
}),
expect.any(Object)
)
const [, options] = fetchMock.mock.calls[0] as [string, RequestInit]
const body = JSON.parse(String(options.body))
expect(body.fileAttachments).toEqual([
{
type: 'document',
source: {
type: 'base64',
media_type: 'text/plain',
data: attachmentData,
},
filename: 'report-{{FILE_TOKEN}}.txt',
},
])
})
it('rejects resolver-derived inline attachment bytes before materialization', async () => {
const encodedSecret = 'aW1hZ2U='
const registry = new ResolvedSecretTraceRegistry([
{
name: 'FILE_BYTES',
plaintext: encodedSecret,
encryptedValue: 'encrypted-file-bytes',
},
])
const inputPath = ['files', '0', 'base64'] as const
registry.recordResolvedAtInputPath('FILE_BYTES', encodedSecret, inputPath)
registry.recordResolvedInputProjection(inputPath, encodedSecret, '{{FILE_BYTES}}')
context.resolvedSecretTraceRegistry = registry
await expect(
handler.execute(context, block, {
prompt: 'Read the attachment',
files: [
{
name: 'example.png',
key: 'workspace/workspace-1/example.png',
size: 5,
type: 'image/png',
base64: encodedSecret,
},
],
})
).rejects.toThrow('Mothership inline file content cannot contain secret references')
expect(mockReadUserFileContent).not.toHaveBeenCalled()
expect(fetchMock).not.toHaveBeenCalled()
})
it('rejects resolver-derived inline bytes inside a serialized file input', async () => {
const rawFiles = JSON.stringify([
{
name: 'example.png',
key: 'workspace/workspace-1/example.png',
size: 5,
type: 'image/png',
base64: 'aW1hZ2U=',
},
])
const projectedFiles = JSON.stringify([
{
name: 'example.png',
key: 'workspace/workspace-1/example.png',
size: 5,
type: 'image/png',
base64: '{{FILE_BYTES}}',
},
])
const registry = new ResolvedSecretTraceRegistry([
{
name: 'FILE_BYTES',
plaintext: 'aW1hZ2U=',
encryptedValue: 'encrypted-file-bytes',
},
])
registry.recordResolvedAtInputPath('FILE_BYTES', 'aW1hZ2U=', ['files'])
registry.recordResolvedInputProjection(['files'], rawFiles, projectedFiles)
context.resolvedSecretTraceRegistry = registry
await expect(
handler.execute(context, block, {
prompt: 'Read the attachment',
files: rawFiles,
})
).rejects.toThrow('Mothership inline file content cannot contain secret references')
expect(mockReadUserFileContent).not.toHaveBeenCalled()
expect(fetchMock).not.toHaveBeenCalled()
})
it('keeps dormant inline attachment bytes unchanged', async () => {
const encodedBytes = 'aW1hZ2U='
context.resolvedSecretTraceRegistry = new ResolvedSecretTraceRegistry([
{
name: 'UNUSED_FILE_BYTES',
plaintext: encodedBytes,
encryptedValue: 'encrypted-unused-file-bytes',
},
])
mockGenerateId
.mockReturnValueOnce('chat-uuid')
.mockReturnValueOnce('message-uuid')
.mockReturnValueOnce('request-uuid')
mockReadUserFileContent.mockResolvedValueOnce(encodedBytes)
fetchMock.mockResolvedValue(
new Response(JSON.stringify({ content: 'done', toolCalls: [] }), {
status: 200,
headers: { 'Content-Type': 'application/json' },
})
)
await handler.execute(context, block, {
prompt: 'Read the attachment',
files: [
{
name: 'example.png',
key: 'workspace/workspace-1/example.png',
size: 5,
type: 'image/png',
base64: encodedBytes,
},
],
})
expect(mockReadUserFileContent).toHaveBeenCalledTimes(1)
const [, options] = fetchMock.mock.calls[0] as [string, RequestInit]
const body = JSON.parse(String(options.body))
expect(body.fileAttachments[0].source.data).toBe(encodedBytes)
})
it('rejects a canonical tracked file whose exact byte provenance is not model-safe', async () => {
@@ -787,15 +1266,9 @@ describe('MothershipBlockHandler', () => {
expect(fetchMock).not.toHaveBeenCalled()
})
it('drops Mothership selections whose protocol identifiers or schema semantics contain secrets', async () => {
it('does not infer secret provenance from matching protocol or schema literals', async () => {
const secret = 'boundary-secret'
const registry = createTraceRegistryMock()
const matches = [{ plaintext: secret, replacement: '{{API_KEY}}' }]
registry.getModelEgressSnapshot.mockReturnValue({
complete: true,
matches,
matcher: createResolvedSecretMatcher(matches),
})
context.resolvedSecretTraceRegistry = registry
mockGenerateId
.mockReturnValueOnce('chat-uuid')
@@ -838,14 +1311,16 @@ describe('MothershipBlockHandler', () => {
const [, options] = fetchMock.mock.calls[0] as [string, RequestInit]
const body = JSON.parse(String(options.body))
expect(body.mcpTools).toEqual([
{
type: 'mcp',
params: { serverId: 'mcp-server-1', toolName: 'search' },
},
expect(body.mcpTools).toHaveLength(5)
expect(body.mcpTools[0]).toEqual({
type: 'mcp',
params: { serverId: secret, toolName: 'search' },
})
expect(body.mcpTools[2].schema).toEqual({ type: 'string', enum: [secret] })
expect(body.contexts).toEqual([
{ kind: 'skill', skillId: secret, label: 'Unsafe' },
{ kind: 'skill', skillId: 'skill-1', label: 'Safe skill' },
])
expect(body.contexts).toEqual([{ kind: 'skill', skillId: 'skill-1', label: 'Safe skill' }])
expect(JSON.stringify(body)).not.toContain(secret)
})
it('consumes mothership execute heartbeat streams until the final result', async () => {
apps/sim/executor/handlers/mothership/mothership-handler.ts
+206 -119
View File
@@ -6,13 +6,14 @@ import {
BILLING_ATTRIBUTION_HEADER,
serializeBillingAttributionHeader,
} from '@/lib/billing/core/billing-attribution'
import {
collectModelVisibleSchemaContent,
restoreModelVisibleSchemaValues,
} from '@/lib/copilot/model-visible-schema'
import { normalizeSecretMountPolicy } from '@/lib/copilot/secret-mount-policy'
import { env } from '@/lib/core/config/env'
import { isExecutionCancelled, isRedisCancellationEnabled } from '@/lib/execution/cancellation'
import {
projectModelSchemaAnnotations,
projectResolvedModelInput,
selectModelSchemaInputPaths,
} from '@/lib/execution/model-input-provenance'
import { readUserFileContent } from '@/lib/execution/payloads/materialization.server'
import {
inspectPrivateToolMetadataEnvelope,
@@ -31,6 +32,7 @@ import {
processSingleFileToUserFile,
type RawFileInput,
} from '@/lib/uploads/utils/file-utils'
import { selectModelBoundFileInputPaths } from '@/lib/uploads/utils/model-input'
import type { BlockOutput } from '@/blocks/types'
import { normalizeFileInput } from '@/blocks/utils'
import { BlockType } from '@/executor/constants'
@@ -41,11 +43,10 @@ import type {
StreamingExecution,
} from '@/executor/types'
import { buildAPIUrl, buildAuthHeaders, extractAPIErrorMessage } from '@/executor/utils/http'
import {
isResolvedSecretModelContentUnchanged,
projectResolvedSecretModelContent,
} from '@/executor/utils/resolved-secret-content-projection'
import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
import type {
ResolvedSecretInputPath,
ResolvedSecretTraceRegistry,
} from '@/executor/utils/resolved-secret-trace-registry'
import type { SerializedBlock } from '@/serializer/types'
const logger = createLogger('MothershipBlockHandler')
@@ -75,6 +76,17 @@ interface MothershipSkillContext {
label: string
}
interface IndexedMothershipMcpToolSelection {
inputIndex: number
selection: MothershipMcpToolSelection
}
interface IndexedMothershipSkillContext {
inputIndex: number
context: MothershipSkillContext
hasExplicitLabel: boolean
}
type MothershipExecuteResult = {
content?: string
model?: string
@@ -92,67 +104,21 @@ type MothershipExecuteStreamEvent =
Record<typeof RESOLVED_SECRET_PROVENANCE_FIELD, unknown>
>)
function projectMothershipPrompt(
prompt: string,
registry: ResolvedSecretTraceRegistry | undefined
): string {
const projection = projectResolvedSecretModelContent(prompt, registry)
if (!projection.safe || typeof projection.value !== 'string') {
throw new Error('Mothership input could not be safely projected')
}
return projection.value
}
function projectMothershipMcpTools(
tools: unknown,
registry: ResolvedSecretTraceRegistry | undefined
): MothershipMcpToolSelection[] {
function selectIndexedMothershipMcpTools(tools: unknown): IndexedMothershipMcpToolSelection[] {
if (!Array.isArray(tools)) return []
return tools.flatMap((candidate) => {
return tools.flatMap((candidate, inputIndex) => {
if (!isPlainRecord(candidate) || candidate.type !== 'mcp') return []
if (candidate.usageControl === 'none' || !isPlainRecord(candidate.params)) return []
const { serverId, toolName } = candidate.params
if (
typeof serverId !== 'string' ||
!serverId ||
typeof toolName !== 'string' ||
!toolName ||
!isResolvedSecretModelContentUnchanged([serverId, toolName], registry)
) {
if (typeof serverId !== 'string' || !serverId || typeof toolName !== 'string' || !toolName) {
return []
}
const serverName =
typeof candidate.params.serverName === 'string' ? candidate.params.serverName : undefined
const schema = isPlainRecord(candidate.schema) ? candidate.schema : undefined
const schemaContent = schema
? collectModelVisibleSchemaContent(schema)
: { projectedValues: [], guardedValues: [] }
if (!isResolvedSecretModelContentUnchanged(schemaContent.guardedValues, registry)) return []
const projection = projectResolvedSecretModelContent(
[serverName, schemaContent.projectedValues],
registry
)
if (!projection.safe || !Array.isArray(projection.value) || projection.value.length !== 2) {
throw new Error('Mothership MCP tool metadata could not be safely projected')
}
const [projectedServerName, projectedSchemaValues] = projection.value
if (serverName === undefined && projectedServerName !== undefined) {
throw new Error('Mothership MCP tool metadata could not be safely projected')
}
if (serverName !== undefined && typeof projectedServerName !== 'string') {
throw new Error('Mothership MCP tool metadata could not be safely projected')
}
const projectedSchema = schema
? restoreModelVisibleSchemaValues(schema, projectedSchemaValues)
: undefined
if (projectedSchema !== undefined && !isPlainRecord(projectedSchema)) {
throw new Error('Mothership MCP tool metadata could not be safely projected')
}
const usageControl =
candidate.usageControl === 'auto' || candidate.usageControl === 'force'
@@ -161,54 +127,116 @@ function projectMothershipMcpTools(
const selection: MothershipMcpToolSelection = {
type: 'mcp',
...(usageControl ? { usageControl } : {}),
...(projectedSchema ? { schema: projectedSchema } : {}),
...(schema ? { schema } : {}),
params: {
serverId,
toolName,
...(projectedServerName !== undefined ? { serverName: projectedServerName } : {}),
...(serverName !== undefined ? { serverName } : {}),
},
}
return [selection]
return [{ inputIndex, selection }]
})
}
function projectMothershipSkillContexts(
skills: unknown,
registry: ResolvedSecretTraceRegistry | undefined
): MothershipSkillContext[] {
function selectMothershipMcpTools(tools: unknown): MothershipMcpToolSelection[] {
return selectIndexedMothershipMcpTools(tools).map(({ selection }) => selection)
}
function selectIndexedMothershipSkillContexts(skills: unknown): IndexedMothershipSkillContext[] {
if (!Array.isArray(skills)) return []
const selected = skills.flatMap((candidate) => {
return skills.flatMap((candidate, inputIndex) => {
if (!isPlainRecord(candidate) || typeof candidate.skillId !== 'string' || !candidate.skillId) {
return []
}
if (!isResolvedSecretModelContentUnchanged(candidate.skillId, registry)) return []
const explicitLabel = typeof candidate.name === 'string' ? candidate.name : undefined
const hasExplicitLabel = explicitLabel !== undefined
const label = explicitLabel ?? candidate.skillId
return [
{
skillId: candidate.skillId,
label: typeof candidate.name === 'string' ? candidate.name : candidate.skillId,
inputIndex,
hasExplicitLabel,
context: {
kind: 'skill' as const,
skillId: candidate.skillId,
label,
},
},
]
})
const projection = projectResolvedSecretModelContent(
selected.map((skill) => skill.label),
registry
)
if (!projection.safe) {
throw new Error('Mothership skill metadata could not be safely projected')
}
const projectedLabels = projection.value
if (!Array.isArray(projectedLabels) || projectedLabels.length !== selected.length) {
throw new Error('Mothership skill metadata could not be safely projected')
}
function selectMothershipSkillContexts(skills: unknown): MothershipSkillContext[] {
return selectIndexedMothershipSkillContexts(skills).map(({ context }) => context)
}
function selectMothershipMetadataModelInputPaths(
tools: unknown,
skills: unknown
): {
modelInputPaths: ResolvedSecretInputPath[]
structuralInputPaths: ResolvedSecretInputPath[]
} {
const modelInputPaths: ResolvedSecretInputPath[] = []
const structuralInputPaths: ResolvedSecretInputPath[] = []
for (const { inputIndex, selection } of selectIndexedMothershipMcpTools(tools)) {
const root = ['tools', String(inputIndex)] as const
structuralInputPaths.push([...root, 'params', 'serverId'], [...root, 'params', 'toolName'])
if (selection.schema) {
const schemaPaths = selectModelSchemaInputPaths(selection.schema, [...root, 'schema'])
modelInputPaths.push(...schemaPaths.annotationInputPaths)
structuralInputPaths.push(...schemaPaths.semanticInputPaths)
}
if (selection.params.serverName !== undefined) {
modelInputPaths.push([...root, 'params', 'serverName'])
}
}
return selected.map((skill, index) => {
const label = projectedLabels[index]
if (typeof label !== 'string') {
throw new Error('Mothership skill metadata could not be safely projected')
for (const { inputIndex, hasExplicitLabel } of selectIndexedMothershipSkillContexts(skills)) {
const root = ['skills', String(inputIndex)] as const
structuralInputPaths.push([...root, 'skillId'])
if (hasExplicitLabel) modelInputPaths.push([...root, 'name'])
}
return { modelInputPaths, structuralInputPaths }
}
function assertMothershipToolSchemaProjectionsAreSafe(
registry: ResolvedSecretTraceRegistry,
tools: unknown
): void {
if (!Array.isArray(tools)) return
const projection = registry.projectResolvedInputSelection({ tools })
if (!projection.complete || !Array.isArray(projection.value.tools)) {
throw new Error('Mothership input could not be safely projected')
}
for (const { inputIndex, selection } of selectIndexedMothershipMcpTools(tools)) {
if (!selection.schema) continue
const projectedCandidate = projection.value.tools[inputIndex]
if (!isPlainRecord(projectedCandidate)) {
throw new Error('Mothership input could not be safely projected')
}
return { kind: 'skill', skillId: skill.skillId, label }
})
const projectedSchema = projectedCandidate.schema ?? selection.schema
const schemaProjection = projectModelSchemaAnnotations(selection.schema, projectedSchema)
if (!schemaProjection.safe) {
throw new Error('Mothership input could not be safely projected')
}
}
}
function assertMothershipStructuralInputsDoNotResolveSecrets(
registry: ResolvedSecretTraceRegistry,
inputPaths: readonly ResolvedSecretInputPath[]
): void {
const provenance = registry.exportCommittedProvenanceForInputPaths(inputPaths)
if (!provenance.complete) {
throw new Error('Mothership input could not be safely projected')
}
if (provenance.entries.length > 0) {
throw new Error('Mothership structural model inputs cannot contain secret references')
}
}
async function consumeMothershipProvenance(
@@ -216,8 +244,6 @@ async function consumeMothershipProvenance(
response: Response,
registry?: ResolvedSecretTraceRegistry
): Promise<boolean> {
if (!registry) throw new Error('Mothership model-egress provenance registry is unavailable')
const inspection = inspectPrivateToolMetadataEnvelope(
response.headers,
payload,
@@ -226,31 +252,35 @@ async function consumeMothershipProvenance(
const provenance = payload[RESOLVED_SECRET_PROVENANCE_FIELD]
payload[RESOLVED_SECRET_PROVENANCE_FIELD] = undefined
if (inspection.status === 'unsupported') {
throw new Error('Mothership response does not support private provenance metadata')
return false
}
if (inspection.status === 'invalid') {
registry.markIncomplete()
registry?.markIncomplete()
throw new Error('Mothership response provenance metadata is invalid')
}
if (!registry) return false
const imported = await registry.importProvenanceForValue(provenance, payload, { trusted: true })
if (!imported) throw new Error('Mothership response provenance metadata is invalid')
return true
}
function assertMothershipResponseCapability(
function inspectMothershipResponseCapability(
response: Response,
registry: ResolvedSecretTraceRegistry | undefined
): void {
if (!registry) throw new Error('Mothership model-egress provenance registry is unavailable')
): boolean {
const capability = inspectPrivateToolMetadataResponseCapability(
response.headers,
RESOLVED_SECRET_PROVENANCE_METADATA_V1
)
if (capability.status === 'supported') return
if (capability.status === 'mismatched') registry.markIncomplete()
throw new Error('Mothership response does not support private provenance metadata')
if (capability.status === 'supported') return true
if (capability.status === 'unsupported') {
return false
}
registry?.markIncomplete()
throw new Error('Mothership response provenance metadata is invalid')
}
function parseMothershipExecuteStreamLine(line: string): MothershipExecuteStreamEvent | undefined {
@@ -306,15 +336,18 @@ async function readMothershipExecuteResponse(
response: Response,
registry?: ResolvedSecretTraceRegistry
): Promise<MothershipExecuteResult> {
assertMothershipResponseCapability(response, registry)
const expectsProvenance = inspectMothershipResponseCapability(response, registry)
const contentType = response.headers.get('content-type') || ''
if (!contentType.includes('application/x-ndjson')) {
let result: MothershipExecuteResult
try {
result = (await response.json()) as MothershipExecuteResult
} catch {
registry?.markIncomplete()
throw new Error('Mothership response provenance metadata is invalid')
} catch (error) {
if (expectsProvenance) {
registry?.markIncomplete()
throw new Error('Mothership response provenance metadata is invalid')
}
throw error
}
await consumeMothershipProvenance(result, response, registry)
return result
@@ -374,7 +407,9 @@ async function readMothershipExecuteResponse(
return finalResult
} finally {
if (!finalResult && !receivedTerminalProvenance) registry?.markIncomplete()
if (expectsProvenance && !finalResult && !receivedTerminalProvenance) {
registry?.markIncomplete()
}
reader.releaseLock()
}
}
@@ -389,7 +424,7 @@ function createMothershipStreamingExecution(
registry?: ResolvedSecretTraceRegistry
} = {}
): StreamingExecution {
assertMothershipResponseCapability(response, options.registry)
const expectsProvenance = inspectMothershipResponseCapability(response, options.registry)
if (!response.body) {
throw new Error('Sim execution stream ended without a response body')
}
@@ -476,7 +511,9 @@ function createMothershipStreamingExecution(
controller.error(error)
}
} finally {
if (!sawFinal && !receivedTerminalProvenance) options.registry?.markIncomplete()
if (expectsProvenance && !sawFinal && !receivedTerminalProvenance) {
options.registry?.markIncomplete()
}
cleanup()
reader?.releaseLock()
}
@@ -507,6 +544,7 @@ function createMothershipStreamingExecution(
async function buildMothershipFileAttachments(
filesInput: unknown,
projectedFilesInput: unknown,
ctx: ExecutionContext,
requestId: string
): Promise<MothershipFileAttachment[] | undefined> {
@@ -518,6 +556,10 @@ async function buildMothershipFileAttachments(
if (!ctx.userId) {
throw new Error('Mothership file attachments require an authenticated user.')
}
const projectedFiles = normalizeFileInput(projectedFilesInput)
if (!projectedFiles || projectedFiles.length !== files.length) {
throw new Error('Mothership input could not be safely projected')
}
const userFiles = files.map((file) =>
processSingleFileToUserFile(file as RawFileInput, requestId, logger)
@@ -529,7 +571,17 @@ async function buildMothershipFileAttachments(
if (!modelSafe) throw new Error(MODEL_UNSAFE_WORKSPACE_FILE_ERROR_MESSAGE)
const attachments: MothershipFileAttachment[] = []
for (const userFile of userFiles) {
for (let fileIndex = 0; fileIndex < userFiles.length; fileIndex++) {
const userFile = userFiles[fileIndex]
const rawFile = files[fileIndex]
const projectedFile = projectedFiles[fileIndex]
if (
isPlainRecord(rawFile) &&
isPlainRecord(projectedFile) &&
!Object.is(rawFile.base64, projectedFile.base64)
) {
throw new Error('Mothership inline file content cannot contain secret references')
}
const base64 = await readUserFileContent(userFile, {
encoding: 'base64',
userId: ctx.userId,
@@ -551,7 +603,11 @@ async function buildMothershipFileAttachments(
throw new Error(`File type is not supported for Mothership attachments: ${userFile.name}`)
}
attachments.push({ ...content, filename: userFile.name })
const projectedName = isPlainRecord(projectedFile) ? projectedFile.name : undefined
attachments.push({
...content,
filename: typeof projectedName === 'string' ? projectedName : userFile.name,
})
}
return attachments
@@ -584,10 +640,34 @@ export class MothershipBlockHandler implements BlockHandler {
if (!prompt || typeof prompt !== 'string') {
throw new Error('Prompt input is required')
}
const metadataInputPaths = selectMothershipMetadataModelInputPaths(inputs.tools, inputs.skills)
if (ctx.resolvedSecretTraceRegistry) {
assertMothershipStructuralInputsDoNotResolveSecrets(
ctx.resolvedSecretTraceRegistry,
metadataInputPaths.structuralInputPaths
)
assertMothershipToolSchemaProjectionsAreSafe(ctx.resolvedSecretTraceRegistry, inputs.tools)
}
const modelInputPaths: ResolvedSecretInputPath[] = [
['prompt'],
...selectModelBoundFileInputPaths(inputs.files, ['files'], {
includeInlineBase64: true,
includeName: true,
}),
...metadataInputPaths.modelInputPaths,
]
const modelInputProjection = projectResolvedModelInput(
ctx.resolvedSecretTraceRegistry,
{ prompt, files: inputs.files, tools: inputs.tools, skills: inputs.skills },
modelInputPaths
)
if (!modelInputProjection.complete || typeof modelInputProjection.value.prompt !== 'string') {
throw new Error('Mothership input could not be safely projected')
}
const messages = [
{
role: 'user' as const,
content: projectMothershipPrompt(prompt, ctx.resolvedSecretTraceRegistry),
content: modelInputProjection.value.prompt,
},
]
const providedConversationId =
@@ -599,12 +679,14 @@ export class MothershipBlockHandler implements BlockHandler {
secretScope: inputs.secretScope,
mountedSecrets: inputs.mountedSecrets,
})
const fileAttachments = await buildMothershipFileAttachments(inputs.files, ctx, requestId)
const mcpTools = projectMothershipMcpTools(inputs.tools, ctx.resolvedSecretTraceRegistry)
const skillContexts = projectMothershipSkillContexts(
inputs.skills,
ctx.resolvedSecretTraceRegistry
const fileAttachments = await buildMothershipFileAttachments(
inputs.files,
modelInputProjection.value.files,
ctx,
requestId
)
const mcpTools = selectMothershipMcpTools(modelInputProjection.value.tools)
const skillContexts = selectMothershipSkillContexts(modelInputProjection.value.skills)
const url = buildAPIUrl('/api/mothership/execute')
const headers = await buildAuthHeaders(ctx.userId)
@@ -706,15 +788,20 @@ export class MothershipBlockHandler implements BlockHandler {
})
if (!response.ok) {
assertMothershipResponseCapability(response, ctx.resolvedSecretTraceRegistry)
let payload: MothershipExecuteResult
try {
payload = (await response.clone().json()) as MothershipExecuteResult
} catch {
ctx.resolvedSecretTraceRegistry?.markIncomplete()
throw new Error('Mothership response provenance metadata is invalid')
const expectsProvenance = inspectMothershipResponseCapability(
response,
ctx.resolvedSecretTraceRegistry
)
if (expectsProvenance) {
let payload: MothershipExecuteResult
try {
payload = (await response.clone().json()) as MothershipExecuteResult
} catch {
ctx.resolvedSecretTraceRegistry?.markIncomplete()
throw new Error('Mothership response provenance metadata is invalid')
}
await consumeMothershipProvenance(payload, response, ctx.resolvedSecretTraceRegistry)
}
await consumeMothershipProvenance(payload, response, ctx.resolvedSecretTraceRegistry)
const errorMsg = await extractAPIErrorMessage(response)
throw new Error(`Sim execution failed: ${errorMsg}`)
}
apps/sim/executor/handlers/pi/pi-handler.test.ts
+51 -16
View File
@@ -190,29 +190,39 @@ describe('PiBlockHandler', () => {
it('projects activated task secrets at the final Pi input boundary', async () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'API_KEY', plaintext: 'secret-value', encryptedValue: 'ciphertext' },
{ name: 'UNUSED', plaintext: 'x', encryptedValue: 'unused-ciphertext' },
])
registry.recordResolved('API_KEY', 'secret-value')
registry.recordResolvedAtInputPath('API_KEY', 'secret-value', ['task'])
registry.recordResolvedInputProjection(
['task'],
'Use secret-value without changing Box.',
'Use {{API_KEY}} without changing Box.'
)
registry.recordResolved('UNUSED', 'x')
await handler.execute(
ctx({ resolvedSecretTraceRegistry: registry }),
block,
localInputs({ task: 'Use secret-value without changing the rest.' })
localInputs({ task: 'Use secret-value without changing Box.' })
)
expect(mockRunLocal.mock.calls[0][0].task).toBe('Use {{API_KEY}} without changing the rest.')
expect(mockRunLocal.mock.calls[0][0].task).toBe('Use {{API_KEY}} without changing Box.')
})
it.each([
['missing', undefined],
[
'incomplete',
(() => {
const registry = new ResolvedSecretTraceRegistry()
registry.markIncomplete()
return registry
})(),
],
])('fails closed when task provenance is %s', async (_label, registry) => {
it('preserves legacy task behavior when no provenance registry exists', async () => {
await handler.execute(
ctx({ resolvedSecretTraceRegistry: undefined }),
block,
localInputs({ task: 'ordinary task' })
)
expect(mockRunLocal.mock.calls[0][0].task).toBe('ordinary task')
})
it('fails closed when task provenance is incomplete', async () => {
const registry = new ResolvedSecretTraceRegistry()
registry.markIncomplete()
await expect(
handler.execute(
ctx({ resolvedSecretTraceRegistry: registry }),
@@ -682,7 +692,8 @@ describe('PiBlockHandler', () => {
expect(mockBuildSearchTool).toHaveBeenCalledWith(
expect.anything(),
{ provider: 'exa', apiKey: 'search-key' },
'local'
'local',
'search-key'
)
expect(mockRunLocal.mock.calls[0][0].search).toEqual({
provider: 'exa',
@@ -691,6 +702,29 @@ describe('PiBlockHandler', () => {
})
})
it('replays search-key normalization on the resolver-recorded projection', async () => {
mockParseSearchProvider.mockReturnValue('exa')
mockResolveSearchKey.mockImplementation(({ apiKey }: { apiKey?: string }) => apiKey?.trim())
const registry = new ResolvedSecretTraceRegistry([
{ name: 'SEARCH_KEY', plaintext: ' key\n', encryptedValue: 'ciphertext' },
])
registry.recordResolvedAtInputPath('SEARCH_KEY', ' key\n', ['searchApiKey'])
registry.recordResolvedInputProjection(['searchApiKey'], ' key\n', '{{SEARCH_KEY}}')
await handler.execute(
ctx({ resolvedSecretTraceRegistry: registry }),
block,
localInputs({ searchProvider: 'exa', searchApiKey: ' key\n' })
)
expect(mockBuildSearchTool).toHaveBeenCalledWith(
expect.anything(),
{ provider: 'exa', apiKey: 'key' },
'local',
'{{SEARCH_KEY}}'
)
})
it('builds the host tool for Review Code too', async () => {
mockParseSearchProvider.mockReturnValue('serper')
@@ -708,7 +742,8 @@ describe('PiBlockHandler', () => {
expect(mockBuildSearchTool).toHaveBeenCalledWith(
expect.anything(),
expect.objectContaining({ provider: 'serper' }),
'cloud_review'
'cloud_review',
'search-key'
)
expect(mockRunCloudReview.mock.calls[0][0].search.tool).toEqual({ name: 'web_search' })
})
apps/sim/executor/handlers/pi/pi-handler.ts
+30 -11
View File
@@ -7,6 +7,7 @@
*/
import { createLogger } from '@sim/logger'
import { projectResolvedModelInput } from '@/lib/execution/model-input-provenance'
import type { BlockOutput } from '@/blocks/types'
import { parseOptionalNumberInput } from '@/blocks/utils'
import {
@@ -50,7 +51,6 @@ import type {
NormalizedBlockOutput,
StreamingExecution,
} from '@/executor/types'
import { projectResolvedSecretModelContent } from '@/executor/utils/resolved-secret-content-projection'
import { isPiSupportedProvider, resolvePiModelId } from '@/providers/pi-providers'
import { getProviderFromModel } from '@/providers/utils'
import type { SerializedBlock } from '@/serializer/types'
@@ -163,14 +163,15 @@ export class PiBlockHandler implements BlockHandler {
const mode = parsePiMode(inputs.mode)
const resolvedTask = asOptString(inputs.task)
if (!resolvedTask) throw new Error('Task is required')
const taskProjection = projectResolvedSecretModelContent(
resolvedTask,
ctx.resolvedSecretTraceRegistry
const taskProjection = projectResolvedModelInput(
ctx.resolvedSecretTraceRegistry,
{ task: resolvedTask },
[['task']]
)
if (!taskProjection.safe || typeof taskProjection.value !== 'string') {
if (!taskProjection.complete || typeof taskProjection.value.task !== 'string') {
throw new Error('Pi input could not be safely projected')
}
const task = taskProjection.value
const task = taskProjection.value.task
const model = asOptString(inputs.model) ?? DEFAULT_MODEL
const providerId = getProviderFromModel(model)
@@ -400,15 +401,33 @@ export class PiBlockHandler implements BlockHandler {
throw error
}
const rawSearchApiKey = inputs.searchApiKey
const apiKey = resolvePiSearchKey({
provider,
apiKey: asOptString(inputs.searchApiKey),
apiKey: asOptString(rawSearchApiKey),
})
const credentials = { provider, apiKey }
return mode === 'cloud' || mode === 'cloud_branch'
? credentials
: { ...credentials, tool: buildPiSearchToolSpec(ctx, credentials, mode) }
if (mode === 'cloud' || mode === 'cloud_branch') return credentials
const searchInputProjection = projectResolvedModelInput(
ctx.resolvedSecretTraceRegistry,
{ searchApiKey: rawSearchApiKey },
[['searchApiKey']]
)
if (!searchInputProjection.complete) {
throw new Error('Pi search input could not be safely projected')
}
const projectedApiKey = Object.is(searchInputProjection.value.searchApiKey, rawSearchApiKey)
? apiKey
: resolvePiSearchKey({
provider,
apiKey: asOptString(searchInputProjection.value.searchApiKey),
})
return {
...credentials,
tool: buildPiSearchToolSpec(ctx, credentials, mode, projectedApiKey),
}
}
private isContentSelectedForStreaming(ctx: ExecutionContext, block: SerializedBlock): boolean {
apps/sim/executor/handlers/pi/search/tool.test.ts
+134 -31
View File
@@ -1,11 +1,15 @@
/**
* @vitest-environment node
*/
import { encryptionMockFns } from '@sim/testing'
import { beforeEach, describe, expect, it, vi } from 'vitest'
const { mockExecuteTool } = vi.hoisted(() => ({ mockExecuteTool: vi.fn() }))
vi.mock('@/tools', () => ({ executeTool: mockExecuteTool }))
vi.mock('@/lib/core/security/encryption', () => ({
decryptSecret: encryptionMockFns.mockDecryptSecret,
}))
import {
PI_SEARCH_BUDGET_MESSAGE,
@@ -44,6 +48,7 @@ async function run(
beforeEach(() => {
vi.clearAllMocks()
encryptionMockFns.mockDecryptSecret.mockReset()
})
describe('buildPiSearchToolSpec', () => {
@@ -117,20 +122,30 @@ describe('buildPiSearchToolSpec', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'SEARCH_QUERY', plaintext: secret, encryptedValue: 'ciphertext' },
])
registry.recordResolved('SEARCH_QUERY', secret)
const mergeSpy = vi.spyOn(registry, 'mergeToolCallRegistry')
const context = executionContext(registry)
mockExecuteTool.mockResolvedValue({
success: true,
output: {
results: [
{
title: secret,
url: 'https://example.com/docs',
text: `Bearer ${secret}`,
},
],
},
encryptionMockFns.mockDecryptSecret.mockResolvedValue({ decrypted: secret })
mockExecuteTool.mockImplementation(async (_toolId, _params, options) => {
await options.resolvedSecretTraceRegistry.importProvenance(
{
version: 1,
complete: true,
entries: [{ name: 'SEARCH_QUERY', encryptedValue: 'ciphertext' }],
},
{ trusted: true }
)
return {
success: true,
output: {
results: [
{
title: secret,
url: 'https://example.com/docs',
text: `Bearer ${secret}`,
},
],
},
}
})
const result = await buildTool('exa', context).execute({ query: secret })
@@ -176,13 +191,67 @@ describe('buildPiSearchToolSpec', () => {
})
})
it('projects only the exact resolver-recorded search key and leaves the raw result unchanged', async () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'SEARCH_KEY', plaintext: 'key-123', encryptedValue: 'search-ciphertext' },
{ name: 'UNRELATED', plaintext: 'Test', encryptedValue: 'unrelated-ciphertext' },
])
registry.recordResolvedAtInputPath('SEARCH_KEY', 'key-123', ['searchApiKey'])
registry.recordResolvedInputProjection(['searchApiKey'], 'key-123', '{{SEARCH_KEY}}')
registry.recordResolvedAtInputPath('UNRELATED', 'Test', ['task'])
registry.recordResolvedInputProjection(['task'], 'Test', '{{UNRELATED}}')
const output = {
results: [
{
title: 'key-123',
url: 'https://example.com/docs',
text: 'Test',
},
],
}
mockExecuteTool.mockResolvedValue({ success: true, output })
const result = await buildPiSearchToolSpec(
executionContext(registry),
{ provider: 'exa', apiKey: 'key-123' },
'local',
'{{SEARCH_KEY}}'
).execute({ query: 'pi' })
expect(JSON.parse(result.text).results[0]).toEqual({
title: '{{SEARCH_KEY}}',
url: 'https://example.com/docs',
snippet: 'Test',
})
expect(
mockExecuteTool.mock.calls[0][2].resolvedSecretTraceRegistry
.exportCommittedProvenanceForInputPaths([['apiKey']])
.entries.map((entry: { name?: string }) => entry.name)
).toEqual(['SEARCH_KEY'])
expect(output).toEqual({
results: [
{
title: 'key-123',
url: 'https://example.com/docs',
text: 'Test',
},
],
})
})
it('projects anonymous provenance learned by the isolated search call', async () => {
const registry = new ResolvedSecretTraceRegistry()
encryptionMockFns.mockDecryptSecret.mockResolvedValue({ decrypted: 'foreign-secret' })
mockExecuteTool.mockImplementation(async (_toolId, _params, options) => {
vi.spyOn(options.resolvedSecretTraceRegistry, 'getModelEgressSnapshot').mockReturnValue({
complete: true,
matches: [{ plaintext: 'foreign-secret', replacement: '[REDACTED_SECRET]' }],
})
await options.resolvedSecretTraceRegistry.importProvenance(
{
version: 1,
complete: true,
entries: [{ name: 'FOREIGN', encryptedValue: 'foreign-ciphertext' }],
scope: { userId: 'foreign-user', workspaceId: 'foreign-workspace' },
},
{ trusted: true }
)
return {
success: true,
output: {
@@ -202,24 +271,31 @@ describe('buildPiSearchToolSpec', () => {
expect(JSON.parse(result.text).results[0].snippet).toBe('[REDACTED_SECRET]')
})
it.each([
['missing', undefined],
[
'incomplete',
(() => {
const registry = new ResolvedSecretTraceRegistry()
registry.markIncomplete()
return registry
})(),
],
])('fails closed before search when provenance is %s', async (_label, registry) => {
const context = registry
? executionContext(registry)
: ({ executionId: 'exec-1', workspaceId: 'ws-1' } as ExecutionContext)
it('preserves legacy search behavior when no provenance registry exists', async () => {
const output = {
results: [{ title: 'Docs', url: 'https://example.com/docs', text: 'Page text' }],
}
mockExecuteTool.mockResolvedValue({ success: true, output })
const context = { executionId: 'exec-1', workspaceId: 'ws-1' } as ExecutionContext
const result = await buildTool('exa', context).execute({ query: 'pi' })
expect(result.isError).toBe(false)
expect(JSON.parse(result.text).results[0].snippet).toBe('Page text')
expect(mockExecuteTool.mock.calls[0][2].resolvedSecretTraceRegistry).toBeUndefined()
expect(output.results[0].text).toBe('Page text')
})
it('fails closed before search when provenance is incomplete', async () => {
const registry = new ResolvedSecretTraceRegistry()
registry.markIncomplete()
const result = await buildTool('exa', executionContext(registry)).execute({ query: 'pi' })
expect(result.isError).toBe(true)
expect(result.text).toContain('could not be returned safely')
expect(result.text).toBe(
'Web search settled, but its result could not be returned safely. Do not retry automatically.'
)
expect(mockExecuteTool).not.toHaveBeenCalled()
})
@@ -243,6 +319,33 @@ describe('buildPiSearchToolSpec', () => {
expect(mergeSpy).not.toHaveBeenCalled()
})
it('keeps the fixed unavailable message unchanged when active provenance contains one character', async () => {
encryptionMockFns.mockDecryptSecret.mockResolvedValue({ decrypted: 'W' })
const registry = new ResolvedSecretTraceRegistry([
{ name: 'LETTER', plaintext: 'W', encryptedValue: 'encrypted-letter' },
])
const cyclic: Record<string, unknown> = {}
cyclic.self = cyclic
mockExecuteTool.mockImplementation(async (_toolId, _params, options) => {
await options.resolvedSecretTraceRegistry.importProvenance(
{
version: 1,
complete: true,
entries: [{ name: 'LETTER', encryptedValue: 'encrypted-letter' }],
},
{ trusted: true }
)
return { success: true, output: cyclic }
})
const result = await buildTool('exa', executionContext(registry)).execute({ query: 'pi' })
expect(result).toEqual({
text: 'Web search settled, but its result could not be returned safely. Do not retry automatically.',
isError: true,
})
})
it('reports an empty search as a successful no-results envelope', async () => {
mockExecuteTool.mockResolvedValue({ success: true, output: { results: [] } })
apps/sim/executor/handlers/pi/search/tool.ts
+30 -21
View File
@@ -25,25 +25,19 @@ import {
serializePiSearchEnvelope,
} from '@/executor/handlers/pi/search/normalize'
import type { ExecutionContext } from '@/executor/types'
import {
projectResolvedSecretModelContent,
projectResolvedSecretModelControlMessage,
} from '@/executor/utils/resolved-secret-content-projection'
import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
import { projectResolvedSecretModelContent } from '@/executor/utils/resolved-secret-content-projection'
import { executeTool } from '@/tools'
const logger = createLogger('PiSearchTool')
const SEARCH_RESULT_UNAVAILABLE_MESSAGE =
'Web search settled, but its result could not be returned safely. Do not retry automatically.'
function unavailableSearchResult(registry: ResolvedSecretTraceRegistry | undefined): {
function unavailableSearchResult(): {
text: string
isError: true
} {
return {
text:
projectResolvedSecretModelControlMessage(SEARCH_RESULT_UNAVAILABLE_MESSAGE, registry) ??
SEARCH_RESULT_UNAVAILABLE_MESSAGE,
text: SEARCH_RESULT_UNAVAILABLE_MESSAGE,
isError: true,
}
}
@@ -87,7 +81,8 @@ function describePiSearchFailure(label: string, status: unknown, error: unknown)
export function buildPiSearchToolSpec(
ctx: ExecutionContext,
search: Pick<PiSearchConfig, 'provider' | 'apiKey'>,
mode: 'local' | 'cloud_review'
mode: 'local' | 'cloud_review',
projectedApiKey = search.apiKey
): PiToolSpec {
const { label, toolId } = PI_SEARCH_PROVIDERS[search.provider]
const logContext = {
@@ -122,9 +117,18 @@ export function buildPiSearchToolSpec(
timeout: PI_SEARCH_TIMEOUT_MS,
}
const registry = ctx.resolvedSecretTraceRegistry
const toolCallRegistry = registry?.forkForToolInputValues(Object.values(providerParams))
if (!registry || !toolCallRegistry?.isComplete()) {
return unavailableSearchResult(toolCallRegistry)
const toolCallRegistry = registry?.forkForInputPaths([['searchApiKey']], {
propagated: true,
})
if (toolCallRegistry && !toolCallRegistry.isComplete()) {
return unavailableSearchResult()
}
if (toolCallRegistry) {
toolCallRegistry.recordTransformedInputProjection(providerParams, {
...providerParams,
apiKey: projectedApiKey,
})
if (!toolCallRegistry.isComplete()) return unavailableSearchResult()
}
const result = await executeTool(toolId, providerParams, {
@@ -133,18 +137,18 @@ export function buildPiSearchToolSpec(
})
if (!result.success) {
if (!toolCallRegistry.isComplete()) {
return unavailableSearchResult(toolCallRegistry)
if (toolCallRegistry && !toolCallRegistry.isComplete()) {
return unavailableSearchResult()
}
if (result.error === PARALLEL_EMPTY_RESULTS_ERROR) {
logger.info('Pi search returned no results', { ...logContext, resultCount: 0 })
registry.mergeToolCallRegistry(toolCallRegistry)
if (registry && toolCallRegistry) registry.mergeToolCallRegistry(toolCallRegistry)
return { text: serializePiSearchEnvelope([]), isError: false }
}
const status = (result.output as { status?: unknown } | undefined)?.status
logger.warn('Pi search failed', { ...logContext, status })
registry.mergeToolCallRegistry(toolCallRegistry)
if (registry && toolCallRegistry) registry.mergeToolCallRegistry(toolCallRegistry)
return {
// Classified rather than quoted: `result.error` can carry provider-response-derived text
// for all four providers, which the untrusted-results guideline does not cover. Only the
@@ -155,9 +159,14 @@ export function buildPiSearchToolSpec(
}
}
const outputProjection = projectResolvedSecretModelContent(result.output, toolCallRegistry)
if (!outputProjection.safe || !toolCallRegistry.isComplete()) {
return unavailableSearchResult(toolCallRegistry)
const outputProjection = toolCallRegistry
? projectResolvedSecretModelContent(
result.output,
toolCallRegistry.forkForPropagatedEntries()
)
: ({ safe: true, value: result.output } as const)
if (!outputProjection.safe || (toolCallRegistry && !toolCallRegistry.isComplete())) {
return unavailableSearchResult()
}
const results = normalizePiSearchRecords(
search.provider,
@@ -165,7 +174,7 @@ export function buildPiSearchToolSpec(
numResults
)
logger.info('Pi search completed', { ...logContext, resultCount: results.length })
registry.mergeToolCallRegistry(toolCallRegistry)
if (registry && toolCallRegistry) registry.mergeToolCallRegistry(toolCallRegistry)
return { text: serializePiSearchEnvelope(results), isError: false }
},
}
apps/sim/executor/handlers/pi/sim-tools.test.ts
+145 -25
View File
@@ -1,6 +1,7 @@
/**
* @vitest-environment node
*/
import { encryptionMockFns } from '@sim/testing'
import { beforeEach, describe, expect, it, vi } from 'vitest'
const { mockTransformBlockTool, mockExecuteTool } = vi.hoisted(() => ({
@@ -12,6 +13,9 @@ vi.mock('@/providers/utils', () => ({ transformBlockTool: mockTransformBlockTool
vi.mock('@/tools', () => ({ executeTool: mockExecuteTool }))
vi.mock('@/tools/utils', () => ({ getTool: vi.fn() }))
vi.mock('@/tools/utils.server', () => ({ getToolAsync: vi.fn() }))
vi.mock('@/lib/core/security/encryption', () => ({
decryptSecret: encryptionMockFns.mockDecryptSecret,
}))
import { buildSimToolSpecs } from '@/executor/handlers/pi/sim-tools'
import type { ExecutionContext } from '@/executor/types'
@@ -44,6 +48,7 @@ function mockToolAdapter(params: Record<string, unknown> = {}): void {
describe('buildSimToolSpecs', () => {
beforeEach(() => {
vi.clearAllMocks()
encryptionMockFns.mockDecryptSecret.mockReset()
})
it('names the Pi tool with the snake_case tool id, not the human label', async () => {
@@ -122,9 +127,20 @@ describe('buildSimToolSpecs', () => {
it('projects named provenance in successful Sim tool output', async () => {
mockToolAdapter({ apiKey: 'secret-value' })
mockExecuteTool.mockResolvedValue({
success: true,
output: { authorization: 'Bearer secret-value' },
encryptionMockFns.mockDecryptSecret.mockResolvedValue({ decrypted: 'secret-value' })
mockExecuteTool.mockImplementation(async (_toolId, _params, options) => {
await options.resolvedSecretTraceRegistry.importProvenance(
{
version: 1,
complete: true,
entries: [{ name: 'API_KEY', encryptedValue: 'ciphertext' }],
},
{ trusted: true }
)
return {
success: true,
output: { authorization: 'Bearer secret-value' },
}
})
const registry = new ResolvedSecretTraceRegistry([
{ name: 'API_KEY', plaintext: 'secret-value', encryptedValue: 'ciphertext' },
@@ -141,11 +157,17 @@ describe('buildSimToolSpecs', () => {
it('uses the anonymous fallback for cross-scope provenance', async () => {
mockToolAdapter()
encryptionMockFns.mockDecryptSecret.mockResolvedValue({ decrypted: 'foreign-secret' })
mockExecuteTool.mockImplementation(async (_toolId, _params, options) => {
vi.spyOn(options.resolvedSecretTraceRegistry, 'getModelEgressSnapshot').mockReturnValue({
complete: true,
matches: [{ plaintext: 'foreign-secret', replacement: '[REDACTED_SECRET]' }],
})
await options.resolvedSecretTraceRegistry.importProvenance(
{
version: 1,
complete: true,
entries: [{ name: 'FOREIGN', encryptedValue: 'foreign-ciphertext' }],
scope: { userId: 'foreign-user', workspaceId: 'foreign-workspace' },
},
{ trusted: true }
)
return {
success: true,
output: { token: 'foreign-secret' },
@@ -189,54 +211,152 @@ describe('buildSimToolSpecs', () => {
await expect(spec.execute({})).resolves.toEqual({ text: 'Test', isError: false })
})
it('projects only the selected tool params by original array index and leaves raw output unchanged', async () => {
const selectedTool = {
type: 'exa',
operation: 'exa_search',
usageControl: 'auto',
params: { apiKey: 'secret-value' },
}
const tools = [{ type: 'exa', operation: 'exa_search', usageControl: 'none' }, selectedTool]
mockToolAdapter(selectedTool.params)
const output = { selected: 'secret-value', unrelated: 'Test' }
mockExecuteTool.mockResolvedValue({ success: true, output })
const registry = new ResolvedSecretTraceRegistry([
{ name: 'API_KEY', plaintext: 'secret-value', encryptedValue: 'ciphertext' },
{ name: 'UNRELATED', plaintext: 'Test', encryptedValue: 'unrelated-ciphertext' },
])
registry.recordResolvedAtInputPath('API_KEY', 'secret-value', [
'tools',
'1',
'params',
'apiKey',
])
registry.recordResolvedInputProjection(
['tools', '1', 'params', 'apiKey'],
'secret-value',
'{{API_KEY}}'
)
registry.recordResolvedAtInputPath('UNRELATED', 'Test', ['task'])
registry.recordResolvedInputProjection(['task'], 'Test', '{{UNRELATED}}')
const [spec] = await buildSimToolSpecs(executionContext(registry), tools)
await expect(spec.execute({ query: 'pi' })).resolves.toEqual({
text: JSON.stringify({ selected: '{{API_KEY}}', unrelated: 'Test' }),
isError: false,
})
expect(
mockExecuteTool.mock.calls[0][2].resolvedSecretTraceRegistry
.exportCommittedProvenanceForInputPaths([['apiKey']])
.entries.map((entry: { name?: string }) => entry.name)
).toEqual(['API_KEY'])
expect(output).toEqual({ selected: 'secret-value', unrelated: 'Test' })
})
it('projects error text returned or thrown by a Sim tool', async () => {
mockToolAdapter({ apiKey: 'secret-value' })
encryptionMockFns.mockDecryptSecret.mockResolvedValue({ decrypted: 'secret-value' })
const registry = new ResolvedSecretTraceRegistry([
{ name: 'API_KEY', plaintext: 'secret-value', encryptedValue: 'ciphertext' },
])
registry.recordResolved('API_KEY', 'secret-value')
const [spec] = await buildSimToolSpecs(executionContext(registry), toolInput)
mockExecuteTool.mockResolvedValueOnce({
success: false,
output: {},
error: 'provider rejected secret-value',
mockExecuteTool.mockImplementationOnce(async (_toolId, _params, options) => {
await options.resolvedSecretTraceRegistry.importProvenance(
{
version: 1,
complete: true,
entries: [{ name: 'API_KEY', encryptedValue: 'ciphertext' }],
},
{ trusted: true }
)
return {
success: false,
output: {},
error: 'provider rejected secret-value',
}
})
await expect(spec.execute({})).resolves.toEqual({
text: 'provider rejected {{API_KEY}}',
isError: true,
})
mockExecuteTool.mockRejectedValueOnce(new Error('transport exposed secret-value'))
mockExecuteTool.mockImplementationOnce(async (_toolId, _params, options) => {
await options.resolvedSecretTraceRegistry.importProvenance(
{
version: 1,
complete: true,
entries: [{ name: 'API_KEY', encryptedValue: 'ciphertext' }],
},
{ trusted: true }
)
throw new Error('transport exposed secret-value')
})
await expect(spec.execute({})).resolves.toEqual({
text: 'transport exposed {{API_KEY}}',
isError: true,
})
})
it.each([
['missing', undefined],
[
'incomplete',
(() => {
const registry = new ResolvedSecretTraceRegistry()
registry.markIncomplete()
return registry
})(),
],
])('fails closed when Sim tool result provenance is %s', async (_label, registry) => {
it('preserves legacy Sim tool behavior when no provenance registry exists', async () => {
mockToolAdapter()
const output = { result: 'ordinary output' }
mockExecuteTool.mockResolvedValue({ success: true, output })
const [spec] = await buildSimToolSpecs(executionContext(undefined), toolInput)
await expect(spec.execute({})).resolves.toEqual({
text: JSON.stringify(output),
isError: false,
})
expect(mockExecuteTool.mock.calls[0][2].resolvedSecretTraceRegistry).toBeUndefined()
expect(output).toEqual({ result: 'ordinary output' })
})
it('fails closed when Sim tool result provenance is incomplete', async () => {
mockToolAdapter()
mockExecuteTool.mockResolvedValue({
success: true,
output: { result: 'untrusted output' },
})
const registry = new ResolvedSecretTraceRegistry()
registry.markIncomplete()
const [spec] = await buildSimToolSpecs(executionContext(registry), toolInput)
const result = await spec.execute({})
expect(result.isError).toBe(true)
expect(result.text).toContain('could not be returned safely')
expect(result.text).toBe(
'Tool execution settled, but its result could not be returned safely. Do not retry a mutation automatically.'
)
expect(result.text).not.toContain('untrusted output')
expect(mockExecuteTool).not.toHaveBeenCalled()
})
it('keeps the fixed unavailable message unchanged when active provenance contains one character', async () => {
mockToolAdapter()
encryptionMockFns.mockDecryptSecret.mockResolvedValue({ decrypted: 'T' })
const registry = new ResolvedSecretTraceRegistry([
{ name: 'LETTER', plaintext: 'T', encryptedValue: 'encrypted-letter' },
])
const cyclic: Record<string, unknown> = {}
cyclic.self = cyclic
mockExecuteTool.mockImplementation(async (_toolId, _params, options) => {
await options.resolvedSecretTraceRegistry.importProvenance(
{
version: 1,
complete: true,
entries: [{ name: 'LETTER', encryptedValue: 'encrypted-letter' }],
},
{ trusted: true }
)
return { success: true, output: cyclic }
})
const [spec] = await buildSimToolSpecs(executionContext(registry), toolInput)
await expect(spec.execute({})).resolves.toEqual({
text: 'Tool execution settled, but its result could not be returned safely. Do not retry a mutation automatically.',
isError: true,
})
})
})
apps/sim/executor/handlers/pi/sim-tools.ts
+68 -24
View File
@@ -14,10 +14,7 @@ import { getAllBlocks } from '@/blocks/registry'
import type { ToolInput } from '@/executor/handlers/agent/types'
import type { PiToolResult, PiToolSpec } from '@/executor/handlers/pi/backend'
import type { ExecutionContext } from '@/executor/types'
import {
projectResolvedSecretModelContent,
projectResolvedSecretModelControlMessage,
} from '@/executor/utils/resolved-secret-content-projection'
import { projectResolvedSecretModelContent } from '@/executor/utils/resolved-secret-content-projection'
import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
import { transformBlockTool } from '@/providers/utils'
import { executeTool } from '@/tools'
@@ -30,12 +27,11 @@ import { getToolAsync } from '@/tools/utils.server'
const logger = createLogger('PiSimTools')
const TOOL_RESULT_UNAVAILABLE_MESSAGE =
'Tool execution settled, but its result could not be returned safely. Do not retry a mutation automatically.'
const TOOL_EXECUTION_FAILED_MESSAGE = 'Tool execution failed'
function unavailableToolResult(registry: ResolvedSecretTraceRegistry | undefined): PiToolResult {
function unavailableToolResult(): PiToolResult {
return {
text:
projectResolvedSecretModelControlMessage(TOOL_RESULT_UNAVAILABLE_MESSAGE, registry) ??
TOOL_RESULT_UNAVAILABLE_MESSAGE,
text: TOOL_RESULT_UNAVAILABLE_MESSAGE,
isError: true,
}
}
@@ -49,10 +45,28 @@ function projectToolResult(
registry: ResolvedSecretTraceRegistry | undefined
): PiToolResultProjection {
try {
if (!registry) {
if (!result.success) {
return {
safe: true,
result: {
text: result.error || TOOL_EXECUTION_FAILED_MESSAGE,
isError: true,
},
}
}
const text =
typeof result.output === 'string' ? result.output : JSON.stringify(result.output ?? {})
return typeof text === 'string'
? { safe: true, result: { text, isError: false } }
: { safe: false, result: unavailableToolResult() }
}
if (result.success) {
const projection = projectResolvedSecretModelContent(result.output, registry)
if (!projection.safe) {
return { safe: false, result: unavailableToolResult(registry) }
return { safe: false, result: unavailableToolResult() }
}
const text =
@@ -61,18 +75,24 @@ function projectToolResult(
: JSON.stringify(projection.value ?? {})
return typeof text === 'string'
? { safe: true, result: { text, isError: false } }
: { safe: false, result: unavailableToolResult(registry) }
: { safe: false, result: unavailableToolResult() }
}
const projection = projectResolvedSecretModelContent(
result.error || 'Tool execution failed',
registry
)
if (!result.error) {
return registry.isComplete()
? {
safe: true,
result: { text: TOOL_EXECUTION_FAILED_MESSAGE, isError: true },
}
: { safe: false, result: unavailableToolResult() }
}
const projection = projectResolvedSecretModelContent(result.error, registry)
return projection.safe && typeof projection.value === 'string'
? { safe: true, result: { text: projection.value, isError: true } }
: { safe: false, result: unavailableToolResult(registry) }
: { safe: false, result: unavailableToolResult() }
} catch {
return { safe: false, result: unavailableToolResult(registry) }
return { safe: false, result: unavailableToolResult() }
}
}
@@ -88,7 +108,7 @@ export async function buildSimToolSpecs(
const specs: PiToolSpec[] = []
for (const tool of inputTools as ToolInput[]) {
for (const [toolIndex, tool] of (inputTools as ToolInput[]).entries()) {
if ((tool.usageControl || 'auto') === 'none') continue
if (!tool.type || tool.type === 'mcp' || tool.type === 'custom-tool') continue
@@ -123,9 +143,30 @@ export async function buildSimToolSpecs(
execute: async (args) => {
const params = mergeToolParameters(preseededParams, args as Record<string, unknown>)
const registry = ctx.resolvedSecretTraceRegistry
const toolCallRegistry = registry?.forkForToolInputValues(Object.values(params))
if (!registry || !toolCallRegistry?.isComplete()) {
return unavailableToolResult(toolCallRegistry)
const sourcePath = ['tools', String(toolIndex), 'params'] as const
const toolCallRegistry = registry?.forkForInputPaths([sourcePath], {
propagated: true,
})
if (toolCallRegistry && !toolCallRegistry.isComplete()) {
return unavailableToolResult()
}
if (toolCallRegistry) {
const inputProjection = toolCallRegistry.projectResolvedInputSelection({
tools: inputTools,
})
const projectedTool = inputProjection.complete
? (inputProjection.value.tools as ToolInput[] | undefined)?.[toolIndex]
: undefined
if (!inputProjection.complete || !projectedTool) {
return unavailableToolResult()
}
const projectedParams = mergeToolParameters(
projectedTool.params || {},
args as Record<string, unknown>
)
toolCallRegistry.recordTransformedInputProjection(params, projectedParams)
if (!toolCallRegistry.isComplete()) return unavailableToolResult()
}
try {
@@ -156,8 +197,11 @@ export async function buildSimToolSpecs(
resolvedSecretTraceRegistry: toolCallRegistry,
}
)
const projection = projectToolResult(result, toolCallRegistry)
if (projection.safe && toolCallRegistry.isComplete()) {
const projection = projectToolResult(
result,
toolCallRegistry?.forkForPropagatedEntries()
)
if (projection.safe && registry && toolCallRegistry?.isComplete()) {
registry.mergeToolCallRegistry(toolCallRegistry)
}
return projection.result
@@ -168,9 +212,9 @@ export async function buildSimToolSpecs(
output: {},
error: getErrorMessage(error, 'Tool execution failed'),
},
toolCallRegistry
toolCallRegistry?.forkForPropagatedEntries()
)
if (projection.safe && toolCallRegistry.isComplete()) {
if (projection.safe && registry && toolCallRegistry?.isComplete()) {
registry.mergeToolCallRegistry(toolCallRegistry)
}
return projection.result
apps/sim/executor/handlers/router/router-handler.test.ts
+120 -4
View File
@@ -1,7 +1,12 @@
import '@sim/testing/mocks/executor'
import { createLogger } from '@sim/logger'
import { authOAuthUtilsMock, authOAuthUtilsMockFns } from '@sim/testing'
import {
authOAuthUtilsMock,
authOAuthUtilsMockFns,
encryptionMock,
encryptionMockFns,
} from '@sim/testing'
import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest'
const { mockResolveAutoModel } = vi.hoisted(() => ({
@@ -9,6 +14,7 @@ const { mockResolveAutoModel } = vi.hoisted(() => ({
}))
vi.mock('@/app/api/auth/oauth/utils', () => authOAuthUtilsMock)
vi.mock('@/lib/core/security/encryption', () => encryptionMock)
vi.mock('@/lib/credentials/access', () => ({
getCredentialActorContext: vi.fn().mockResolvedValue({
@@ -32,7 +38,11 @@ vi.mock('@/lib/model-router/resolve', () => ({
SIM_AUTO_SYSTEM_PREAMBLE: 'Sim auto system preamble',
}))
import { PRIVATE_MODEL_INPUT_PROVENANCE_HEADER } from '@/lib/execution/model-input-provenance'
import {
PRIVATE_MODEL_INPUT_PROVENANCE_HEADER,
PRIVATE_MODEL_INPUT_STATE_HEADER,
PROJECTED_MODEL_INPUT_PATHS_V1,
} from '@/lib/execution/model-input-provenance'
import {
RESOLVED_SECRET_PROVENANCE_FIELD,
RESOLVED_SECRET_PROVENANCE_METADATA_V1,
@@ -124,6 +134,7 @@ describe('RouterBlockHandler', () => {
}
vi.clearAllMocks()
encryptionMockFns.mockDecryptSecret.mockResolvedValue({ decrypted: 'test-decrypted' })
// unstubGlobals removes any module-scope fetch stub before each test, so re-stub here
vi.stubGlobal('fetch', mockFetch)
@@ -258,7 +269,8 @@ describe('RouterBlockHandler', () => {
encryptedValue: 'encrypted-router-credential',
},
])
registry.recordResolved('PROMPT_SECRET', promptSecret)
registry.recordResolvedAtInputPath('PROMPT_SECRET', promptSecret, ['prompt'])
registry.recordResolvedInputProjection(['prompt'], promptSecret, '{{PROMPT_SECRET}}')
registry.recordResolved('API_KEY', credentialSecret)
mockContext.resolvedSecretTraceRegistry = registry
@@ -273,6 +285,9 @@ describe('RouterBlockHandler', () => {
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBe(
RESOLVED_SECRET_PROVENANCE_METADATA_V1
)
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBe(
PROJECTED_MODEL_INPUT_PATHS_V1
)
expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({
version: 1,
complete: true,
@@ -284,6 +299,100 @@ describe('RouterBlockHandler', () => {
],
})
expect(requestBody.apiKey).toBe(credentialSecret)
expect(mockGenerateRouterPrompt).toHaveBeenCalledWith('{{PROMPT_SECRET}}', expect.any(Array))
})
it('omits a prior target state when only aggregate secret provenance is available', async () => {
const stateSecret = 'x'
const encryptedStateSecret = 'encrypted-router-state'
const rawState = { result: stateSecret, ordinary: 'Box remains raw state' }
const registry = new ResolvedSecretTraceRegistry([
{
name: 'STATE_SECRET',
plaintext: stateSecret,
encryptedValue: encryptedStateSecret,
},
])
mockContext.resolvedSecretTraceRegistry = registry
mockContext.blockStates = new Map([
[
mockTargetBlock1.id,
{
output: rawState,
executed: true,
executionTime: 1,
resolvedSecretTraceProvenance: {
version: 1,
complete: true,
entries: [{ name: 'STATE_SECRET', encryptedValue: encryptedStateSecret }],
},
},
],
])
encryptionMockFns.mockDecryptSecret.mockImplementation(async (encryptedValue: string) => ({
decrypted: encryptedValue === encryptedStateSecret ? stateSecret : 'test-decrypted',
}))
await handler.execute(mockContext, mockBlock, {
prompt: 'Choose the best option.',
model: 'gpt-4o',
})
expect(mockGenerateRouterPrompt).toHaveBeenCalledWith(
'Choose the best option.',
expect.arrayContaining([
expect.objectContaining({
id: mockTargetBlock1.id,
subBlocks: expect.objectContaining({ p: 'a' }),
currentState: undefined,
}),
])
)
expect(rawState).toEqual({ result: stateSecret, ordinary: 'Box remains raw state' })
expect(mockTargetBlock1.config.params).toEqual({ p: 'a' })
const requestBody = JSON.parse(mockFetch.mock.calls[0][1].body)
expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({
version: 1,
complete: true,
entries: [],
})
})
it('keeps an ordinary prior target state with exact-empty provenance unchanged', async () => {
const rawState = { result: 'x', ordinary: 'Box remains raw state' }
mockContext.resolvedSecretTraceRegistry = new ResolvedSecretTraceRegistry([])
mockContext.blockStates = new Map([
[
mockTargetBlock1.id,
{
output: rawState,
executed: true,
executionTime: 1,
resolvedSecretTraceProvenance: {
version: 1,
complete: true,
entries: [],
},
},
],
])
await handler.execute(mockContext, mockBlock, {
prompt: 'Choose the best option.',
model: 'gpt-4o',
})
expect(mockGenerateRouterPrompt).toHaveBeenCalledWith(
'Choose the best option.',
expect.arrayContaining([
expect.objectContaining({
id: mockTargetBlock1.id,
currentState: rawState,
}),
])
)
expect(rawState).toEqual({ result: 'x', ordinary: 'Box remains raw state' })
})
it('keeps the legacy router request shape when no provenance registry exists', async () => {
@@ -297,6 +406,7 @@ describe('RouterBlockHandler', () => {
const requestBody = JSON.parse(request.body)
expect(Object.hasOwn(requestBody, RESOLVED_SECRET_PROVENANCE_FIELD)).toBe(false)
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBeNull()
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBeNull()
})
it('bills the cost the provider proxy decided rather than recomputing it', async () => {
@@ -665,7 +775,8 @@ describe('RouterBlockHandler V2', () => {
encryptedValue: 'encrypted-router-v2-credential',
},
])
registry.recordResolved('CONTEXT_SECRET', contextSecret)
registry.recordResolvedAtInputPath('CONTEXT_SECRET', contextSecret, ['context'])
registry.recordResolvedInputProjection(['context'], contextSecret, '{{CONTEXT_SECRET}}')
registry.recordResolved('API_KEY', credentialSecret)
mockContext.resolvedSecretTraceRegistry = registry
mockFetch.mockResolvedValueOnce({
@@ -690,6 +801,9 @@ describe('RouterBlockHandler V2', () => {
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBe(
RESOLVED_SECRET_PROVENANCE_METADATA_V1
)
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBe(
PROJECTED_MODEL_INPUT_PATHS_V1
)
expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({
version: 1,
complete: true,
@@ -701,6 +815,7 @@ describe('RouterBlockHandler V2', () => {
],
})
expect(requestBody.apiKey).toBe(credentialSecret)
expect(mockGenerateRouterV2Prompt).toHaveBeenCalledWith('{{CONTEXT_SECRET}}', expect.any(Array))
})
it('keeps the router V2 request shape when no provenance registry exists', async () => {
@@ -725,6 +840,7 @@ describe('RouterBlockHandler V2', () => {
const requestBody = JSON.parse(request.body)
expect(Object.hasOwn(requestBody, RESOLVED_SECRET_PROVENANCE_FIELD)).toBe(false)
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBeNull()
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBeNull()
})
it('resolves sim-auto before executing router V2 and preserves its public identity', async () => {
apps/sim/executor/handlers/router/router-handler.ts
+87 -39
View File
@@ -3,6 +3,8 @@ import { getInternalApiBaseUrl } from '@/lib/core/utils/urls'
import {
addModelInputProvenanceToRequest,
createModelInputProvenanceRequestMetadata,
markModelInputProjected,
projectResolvedModelInput,
} from '@/lib/execution/model-input-provenance'
import {
type AutoRoutingResult,
@@ -22,9 +24,9 @@ import {
} from '@/executor/constants'
import type { BlockHandler, ExecutionContext } from '@/executor/types'
import { buildAuthHeaders } from '@/executor/utils/http'
import type { ResolvedSecretInputPath } from '@/executor/utils/resolved-secret-trace-registry'
import { resolveVertexCredential } from '@/executor/utils/vertex-credential'
import { resolveProxiedModelCost } from '@/providers/cost-policy'
import { collectProviderModelInputProvenanceValues } from '@/providers/model-input-provenance'
import { isAutoModel, SIM_AUTO_MODEL_ID } from '@/providers/models'
import type { ProviderRequest } from '@/providers/types'
import { getProviderFromModel } from '@/providers/utils'
@@ -71,10 +73,19 @@ export class RouterBlockHandler implements BlockHandler {
block: SerializedBlock,
inputs: Record<string, any>
): Promise<BlockOutput> {
const promptModelInputPaths: ResolvedSecretInputPath[] = [['prompt']]
const modelInputProjection = projectResolvedModelInput(
ctx.resolvedSecretTraceRegistry,
{ prompt: inputs.prompt },
promptModelInputPaths
)
if (!modelInputProjection.complete) {
throw new Error('Router model input could not be safely projected')
}
const targetBlocks = this.getTargetBlocks(ctx, block)
const routerConfig = {
prompt: inputs.prompt,
prompt: modelInputProjection.value.prompt,
model: inputs.model || ROUTER.DEFAULT_MODEL,
apiKey: inputs.apiKey,
vertexProject: inputs.vertexProject,
@@ -131,14 +142,16 @@ export class RouterBlockHandler implements BlockHandler {
}
const headers = new Headers(await buildAuthHeaders(ctx.userId))
const modelInputMetadata = createModelInputProvenanceRequestMetadata(
modelInputProjection.registry,
promptModelInputPaths
)
const requestBody = addModelInputProvenanceToRequest(
{ provider: providerId, ...providerRequest },
headers,
createModelInputProvenanceRequestMetadata(
ctx.resolvedSecretTraceRegistry,
collectProviderModelInputProvenanceValues(providerRequest, providerId)
)
modelInputMetadata
)
if (modelInputMetadata) markModelInputProjected(headers)
const response = await fetch(url.toString(), {
method: 'POST',
headers,
@@ -226,8 +239,31 @@ export class RouterBlockHandler implements BlockHandler {
throw new Error('No routes defined for router')
}
const modelInputPaths: ResolvedSecretInputPath[] = [
['context'],
...(Array.isArray(inputs.routes)
? inputs.routes.map((_, index) => ['routes', String(index), 'value'] as const)
: [['routes'] as const]),
]
const modelInputProjection = projectResolvedModelInput(
ctx.resolvedSecretTraceRegistry,
{ context: inputs.context, routes: inputs.routes },
modelInputPaths
)
if (!modelInputProjection.complete) {
throw new Error('Router model input could not be safely projected')
}
const projectedRoutes = this.parseRoutes(modelInputProjection.value.routes)
if (projectedRoutes.length !== routes.length) {
throw new Error('Router model input could not be safely projected')
}
const modelRoutes = routes.map((route, index) => ({
...route,
value: projectedRoutes[index]?.value ?? route.value,
}))
const routerConfig = {
context: inputs.context,
context: modelInputProjection.value.context,
model: inputs.model || ROUTER.DEFAULT_MODEL,
apiKey: inputs.apiKey,
vertexProject: inputs.vertexProject,
@@ -243,7 +279,7 @@ export class RouterBlockHandler implements BlockHandler {
if (ctx.userId) url.searchParams.set('userId', ctx.userId)
const messages = [{ role: 'user', content: routerConfig.context }]
const systemPrompt = generateRouterV2Prompt(routerConfig.context, routes)
const systemPrompt = generateRouterV2Prompt(routerConfig.context, modelRoutes)
const resolved = await this.resolveModel(
ctx,
block.id,
@@ -303,14 +339,16 @@ export class RouterBlockHandler implements BlockHandler {
}
const headers = new Headers(await buildAuthHeaders(ctx.userId))
const modelInputMetadata = createModelInputProvenanceRequestMetadata(
modelInputProjection.registry,
modelInputPaths
)
const requestBody = addModelInputProvenanceToRequest(
{ provider: providerId, ...providerRequest },
headers,
createModelInputProvenanceRequestMetadata(
ctx.resolvedSecretTraceRegistry,
collectProviderModelInputProvenanceValues(providerRequest, providerId)
)
modelInputMetadata
)
if (modelInputMetadata) markModelInputProjected(headers)
const response = await fetch(url.toString(), {
method: 'POST',
headers,
@@ -495,35 +533,45 @@ export class RouterBlockHandler implements BlockHandler {
}
private getTargetBlocks(ctx: ExecutionContext, block: SerializedBlock) {
return ctx.workflow?.connections
.filter((conn) => conn.source === block.id)
.map((conn) => {
const targetBlock = ctx.workflow?.blocks.find((b) => b.id === conn.target)
if (!targetBlock) {
throw new Error(`Target block ${conn.target} not found`)
}
const targetBlocks = []
const connections = ctx.workflow?.connections.filter((conn) => conn.source === block.id) ?? []
let systemPrompt = ''
if (isAgentBlockType(targetBlock.metadata?.id)) {
const paramsPrompt = targetBlock.config?.params?.systemPrompt
const inputsPrompt = targetBlock.inputs?.systemPrompt
systemPrompt =
(typeof paramsPrompt === 'string' ? paramsPrompt : '') ||
(typeof inputsPrompt === 'string' ? inputsPrompt : '') ||
''
}
for (const conn of connections) {
const targetBlock = ctx.workflow?.blocks.find((candidate) => candidate.id === conn.target)
if (!targetBlock) {
throw new Error(`Target block ${conn.target} not found`)
}
return {
id: targetBlock.id,
type: targetBlock.metadata?.id,
title: targetBlock.metadata?.name,
description: targetBlock.metadata?.description,
subBlocks: {
...targetBlock.config.params,
systemPrompt: systemPrompt,
},
currentState: ctx.blockStates.get(targetBlock.id)?.output,
}
let systemPrompt = ''
if (isAgentBlockType(targetBlock.metadata?.id)) {
const paramsPrompt = targetBlock.config?.params?.systemPrompt
const inputsPrompt = targetBlock.inputs?.systemPrompt
systemPrompt =
(typeof paramsPrompt === 'string' ? paramsPrompt : '') ||
(typeof inputsPrompt === 'string' ? inputsPrompt : '') ||
''
}
const targetState = ctx.blockStates.get(targetBlock.id)
const stateProvenance = targetState?.resolvedSecretTraceProvenance
const currentState =
stateProvenance && (!stateProvenance.complete || stateProvenance.entries.length > 0)
? undefined
: targetState?.output
targetBlocks.push({
id: targetBlock.id,
type: targetBlock.metadata?.id,
title: targetBlock.metadata?.name,
description: targetBlock.metadata?.description,
subBlocks: {
...targetBlock.config.params,
systemPrompt,
},
currentState,
})
}
return targetBlocks
}
}
apps/sim/executor/handlers/workflow/custom-block-tool-runner.test.ts
+1
View File
@@ -101,6 +101,7 @@ describe('runCustomBlockTool', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'API_KEY', plaintext: secret, encryptedValue: 'ciphertext' },
])
registry.recordResolved('API_KEY', secret)
mockExecute.mockRejectedValue(new Error(message))
const projected = await runCustomBlockTool(
apps/sim/executor/handlers/workflow/workflow-handler.test.ts
+4 -1
View File
@@ -1247,6 +1247,7 @@ describe('WorkflowBlockHandler', () => {
mockExecutorExecute.mockImplementationOnce(async () => {
childRegistry = executorOptions.at(-1)?.contextExtensions
.resolvedSecretTraceRegistry as ResolvedSecretTraceRegistry
expect(childRegistry.recordResolved('SECRET', 'publisher-secret')).toBe(true)
return {
success: true,
output: {},
@@ -1273,7 +1274,9 @@ describe('WorkflowBlockHandler', () => {
replacement: ANONYMOUS_SECRET_TRACE_REPLACEMENT,
},
])
expect(childRegistry?.getActiveMatches()).toEqual([])
expect(childRegistry?.getActiveMatches()).toEqual([
{ plaintext: 'publisher-secret', replacement: '{{SECRET}}' },
])
expect(mockSetResolvedSecretTraceRegistry).toHaveBeenCalledTimes(1)
})
apps/sim/executor/handlers/workflow/workflow-handler.ts
+1 -1
View File
@@ -705,7 +705,7 @@ export class WorkflowBlockHandler implements BlockHandler {
const exposedOutput = this.projectCustomBlockOutput(executionResult, exposedOutputs)
if (ctx.resolvedSecretTraceRegistry && childResolvedSecretTraceRegistry) {
const crossingProvenance =
childResolvedSecretTraceRegistry.exportCatalogProvenanceForValue(exposedOutput, {
childResolvedSecretTraceRegistry.exportCommittedProvenanceForValue(exposedOutput, {
anonymous: true,
})
await ctx.resolvedSecretTraceRegistry.importProvenance(crossingProvenance, {
apps/sim/executor/orchestrators/loop.ts
+1 -1
View File
@@ -143,7 +143,7 @@ export class LoopOrchestrator {
}
let items: any[]
const parentRegistry = ctx.resolvedSecretTraceRegistry
const resolutionRegistry = parentRegistry?.forkForToolInputValues([])
const resolutionRegistry = parentRegistry?.forkForInputPaths([])
const resolutionCtx = resolutionRegistry
? { ...ctx, resolvedSecretTraceRegistry: resolutionRegistry }
: ctx
apps/sim/executor/orchestrators/parallel.ts
+1 -1
View File
@@ -71,7 +71,7 @@ export class ParallelOrchestrator {
let branchCount: number
let isEmpty = false
const parentRegistry = ctx.resolvedSecretTraceRegistry
const resolutionRegistry = parentRegistry?.forkForToolInputValues([])
const resolutionRegistry = parentRegistry?.forkForInputPaths([])
const resolutionCtx = resolutionRegistry
? { ...ctx, resolvedSecretTraceRegistry: resolutionRegistry }
: ctx
apps/sim/executor/types.ts
+1 -1
View File
@@ -332,7 +332,7 @@ export interface BlockState {
output: NormalizedBlockOutput
executed: boolean
executionTime: number
/** Encrypted provenance filtered to this exact output. Absent means legacy/untracked state. */
/** Encrypted candidates active in this block call. Consumers filter them to the selected value. */
resolvedSecretTraceProvenance?: ResolvedSecretTraceProvenanceV1
}
apps/sim/executor/utils/resolved-secret-content-projection.test.ts
+31 -8
View File
@@ -133,6 +133,30 @@ describe('projectResolvedSecretModelContent', () => {
})
})
it('does not apply provenance traversal limits when no secret was active', () => {
const registry = new ResolvedSecretTraceRegistry()
const value = new Array<null>(100_001).fill(null)
const projection = projectResolvedSecretModelContent(value, registry)
expect(projection.safe).toBe(true)
if (projection.safe) expect(projection.value).toBe(value)
expect(isResolvedSecretModelContentUnchanged(value, registry)).toBe(true)
const jsonProjection = projectResolvedSecretModelJsonContent(value, registry)
expect(jsonProjection.safe).toBe(true)
if (jsonProjection.safe) {
expect(jsonProjection.value).toHaveLength(value.length)
expect((jsonProjection.value as null[]).at(-1)).toBeNull()
expect(jsonProjection.value).not.toBe(value)
}
const jsonString = '{\n "preserve": true\n}'
expect(projectResolvedSecretModelJsonStrings([jsonString, undefined], registry)).toEqual({
safe: true,
value: [jsonString, undefined],
})
})
it('keeps longest-match semantics when a known opaque placeholder is nested in a secret', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'Test', plaintext: 'Test', encryptedValue: 'test-ciphertext' },
@@ -397,33 +421,32 @@ describe('projectResolvedSecretDiagnosticError', () => {
})
})
it('uses a known compiler alias as diagnostic-only provenance', () => {
const secret = 'diagnostic-secret-value'
it('sanitizes an inactive compiler alias without activating or scanning its secret', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'API_KEY', plaintext: secret, encryptedValue: 'ciphertext' },
{ name: 'X', plaintext: 'x', encryptedValue: 'ciphertext' },
])
const error = new Error(`request failed: ${secret} __var_API_KEY`)
const error = new Error('Box __var_X')
expect(projectResolvedSecretDiagnosticError(error, registry)).toEqual(
expect.objectContaining({ error: 'request failed: {{API_KEY}} {{API_KEY}}' })
expect.objectContaining({ error: 'Box [REDACTED_SECRET]' })
)
expect(registry.getActiveMatches()).toEqual([])
})
it('falls back to text-free diagnostics for unknown or runtime-only aliases', () => {
it('lexically sanitizes unknown and runtime-only aliases without catalog inference', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'API_KEY', plaintext: 'secret-value', encryptedValue: 'ciphertext' },
])
expect(
projectResolvedSecretDiagnosticError(new Error('secret-value __var_UNKNOWN'), registry)
).toEqual({ errorType: 'error', hasStack: true })
).toEqual(expect.objectContaining({ error: 'secret-value [REDACTED_SECRET]' }))
expect(
projectResolvedSecretDiagnosticError(
new Error('secret-value __sim_code_1_binding_0'),
registry
)
).toEqual({ errorType: 'error', hasStack: true })
).toEqual(expect.objectContaining({ error: 'secret-value [RUNTIME_BINDING]' }))
})
it('falls back to text-free structure when provenance is missing or incomplete', () => {
apps/sim/executor/utils/resolved-secret-content-projection.ts
+19 -99
View File
@@ -89,11 +89,6 @@ interface ProjectionState {
maxBytes: number
}
interface InternalDiagnosticIdentifierScan {
aliases: Set<string>
foundInternalIdentifier: boolean
}
export interface ResolvedSecretContentProjectionOptions {
/** Values already materialized and verified by a boundary-specific projector. */
isOpaqueSafeObject?: (value: object) => boolean
@@ -194,65 +189,6 @@ function* arrayDataEntries(value: readonly unknown[]): Generator<[number, unknow
}
}
function collectInternalDiagnosticIdentifiers(
value: unknown
): InternalDiagnosticIdentifierScan | undefined {
const result: InternalDiagnosticIdentifierScan = {
aliases: new Set<string>(),
foundInternalIdentifier: false,
}
const ancestors = new WeakSet<object>()
let nodes = 0
const scanString = (candidate: string): void => {
for (const identifier of candidate.match(INTERNAL_DIAGNOSTIC_IDENTIFIER_PATTERN) ?? []) {
result.foundInternalIdentifier = true
if (identifier.startsWith('__var_')) result.aliases.add(identifier)
}
}
const visit = (candidate: unknown, depth: number): boolean => {
nodes += 1
if (nodes > MAX_CONTENT_NODES || depth > MAX_CONTENT_DEPTH) return false
if (typeof candidate === 'string') {
scanString(candidate)
return true
}
if (
candidate === null ||
candidate === undefined ||
typeof candidate === 'number' ||
typeof candidate === 'boolean'
) {
return true
}
if (typeof candidate !== 'object') return false
if (!Array.isArray(candidate) && !isPlainRecord(candidate)) return false
if (ancestors.has(candidate)) return false
ancestors.add(candidate)
try {
if (Array.isArray(candidate)) {
for (const [, item] of arrayDataEntries(candidate)) {
if (!visit(item, depth + 1)) return false
}
return true
}
for (const [key, item] of enumerableDataEntries(candidate)) {
scanString(key)
if (!visit(item, depth + 1)) return false
}
return true
} catch {
return false
} finally {
ancestors.delete(candidate)
}
}
return visit(value, 0) ? result : undefined
}
function sanitizeContent(
value: unknown,
matcher: ResolvedSecretMatcher | undefined,
@@ -451,29 +387,6 @@ export function getResolvedSecretModelMatcher(
}
}
/** Produces a nonempty model control message only when the registry can prove it secret-free. */
export function projectResolvedSecretModelControlMessage(
message: string,
registry: ResolvedSecretTraceRegistry | undefined
): string | undefined {
const projection = projectResolvedSecretModelContent(message, registry)
if (projection.safe && typeof projection.value === 'string' && projection.value.length > 0) {
return projection.value
}
const snapshot = getResolvedSecretModelMatcher(registry)
if (!snapshot.complete) return undefined
for (let codePoint = 0x21; codePoint <= 0x10ffff; codePoint += 1) {
if (codePoint >= 0xd800 && codePoint <= 0xdfff) {
codePoint = 0xdfff
continue
}
const candidate = String.fromCodePoint(codePoint)
if (!snapshot.matcher || !containsResolvedSecret(candidate, snapshot.matcher)) return candidate
}
return undefined
}
/**
* Projects content that is about to become model-visible using committed active provenance.
* Trusted runtime boundaries activate exact secret-bearing values before this point; unrelated
@@ -488,6 +401,9 @@ export function projectResolvedSecretModelContent(
): ResolvedSecretContentProjection {
const snapshot = getResolvedSecretModelMatcher(registry)
if (!snapshot.complete) return { safe: false }
if (!snapshot.matcher && options.sanitizeInternalIdentifiers !== true) {
return { safe: true, value }
}
return projectContent(value, snapshot.matcher, maxBytes, {
projectPrimitiveLiterals: true,
@@ -506,7 +422,8 @@ export function projectResolvedSecretModelJsonContent(
maxBytes = MAX_INLINE_MATERIALIZATION_BYTES,
options: ResolvedSecretContentProjectionOptions = {}
): ResolvedSecretContentProjection {
if (!getResolvedSecretModelMatcher(registry).complete) return { safe: false }
const snapshot = getResolvedSecretModelMatcher(registry)
if (!snapshot.complete) return { safe: false }
try {
const encoded = JSON.stringify(value)
@@ -514,6 +431,9 @@ export function projectResolvedSecretModelJsonContent(
return { safe: false }
}
const normalized: unknown = JSON.parse(encoded)
if (!snapshot.matcher && options.sanitizeInternalIdentifiers !== true) {
return { safe: true, value: normalized }
}
const projection = projectResolvedSecretModelContent(normalized, registry, maxBytes, options)
if (!projection.safe) return projection
@@ -537,17 +457,7 @@ export function projectResolvedSecretDiagnosticContent(
registry: ResolvedSecretTraceRegistry | undefined,
maxBytes = MAX_INLINE_MATERIALIZATION_BYTES
): ResolvedSecretContentProjection {
const identifiers = collectInternalDiagnosticIdentifiers(value)
if (!identifiers) return { safe: false }
let diagnosticRegistry = registry
if (identifiers.foundInternalIdentifier) {
if (!registry || identifiers.aliases.size === 0) return { safe: false }
diagnosticRegistry = registry.forkForDiagnosticAliases(identifiers.aliases)
if (!diagnosticRegistry) return { safe: false }
}
return projectResolvedSecretModelContent(value, diagnosticRegistry, maxBytes, {
return projectResolvedSecretModelContent(value, registry, maxBytes, {
sanitizeInternalIdentifiers: true,
})
}
@@ -597,6 +507,7 @@ export function isResolvedSecretModelContentUnchanged(
): boolean {
const snapshot = getResolvedSecretModelMatcher(registry)
if (!snapshot.complete) return false
if (!snapshot.matcher) return true
const projection = projectContent(value, snapshot.matcher, MAX_INLINE_MATERIALIZATION_BYTES, {
projectPrimitiveLiterals: true,
@@ -616,6 +527,15 @@ export function projectResolvedSecretModelJsonStrings(
): ResolvedSecretContentProjection {
const snapshot = getResolvedSecretModelMatcher(registry)
if (!snapshot.complete) return { safe: false }
if (!snapshot.matcher) {
let outputBytes = 0
for (const value of values) {
if (value === undefined) continue
outputBytes += Buffer.byteLength(value, 'utf8')
if (outputBytes > maxBytes) return { safe: false }
}
return { safe: true, value: [...values] }
}
const projected: Array<string | undefined> = []
let outputBytes = 0
apps/sim/executor/utils/resolved-secret-input-projection.ts
+182
View File
@@ -0,0 +1,182 @@
function makeCanonicalPlaceholdersJsonParseable(value: string): string {
let result = ''
let inString = false
let escaped = false
for (let index = 0; index < value.length; index++) {
const character = value[index]
if (inString) {
result += character
if (escaped) {
escaped = false
} else if (character === '\\') {
escaped = true
} else if (character === '"') {
inString = false
}
continue
}
if (character === '"') {
inString = true
result += character
continue
}
if (character === '{' && value[index + 1] === '{') {
const end = value.indexOf('}}', index + 2)
if (end !== -1) {
const placeholder = value.slice(index, end + 2)
const name = value.slice(index + 2, end).trim()
if (/^[A-Za-z0-9_]+$/.test(name)) {
result += JSON.stringify(placeholder)
index = end + 1
continue
}
}
}
result += character
}
return result
}
function canonicalPlaceholder(value: string): string | undefined {
const match = /^\{\{([^{}]+)\}\}$/.exec(value.trim())
if (!match || !/^[A-Za-z0-9_]+$/.test(match[1].trim())) return undefined
return value.trim()
}
function projectStructuredLeaves(value: unknown, placeholder: string): unknown {
if (value === null || typeof value !== 'object') return placeholder
const projectedRoot: unknown[] | Record<string, unknown> = Array.isArray(value) ? [] : {}
const pending: Array<{
source: unknown[] | Record<string, unknown>
target: unknown[] | Record<string, unknown>
}> = [{ source: value as unknown[] | Record<string, unknown>, target: projectedRoot }]
while (pending.length > 0) {
const { source, target } = pending.pop()!
for (const [key, child] of Object.entries(source)) {
if (child !== null && typeof child === 'object') {
const projectedChild: unknown[] | Record<string, unknown> = Array.isArray(child) ? [] : {}
;(target as Record<string, unknown>)[key] = projectedChild
pending.push({
source: child as unknown[] | Record<string, unknown>,
target: projectedChild,
})
} else {
;(target as Record<string, unknown>)[key] = placeholder
}
}
}
return projectedRoot
}
function projectFileReferenceString(
key: string,
value: string,
placeholder: string
): string | undefined {
if (key !== 'key' && key !== 'path' && key !== 'url' && key !== 'base64') return undefined
if ((key === 'url' || key === 'path') && value.startsWith('https://')) {
return `https://${placeholder}`
}
if ((key === 'url' || key === 'path') && value.startsWith('http://')) {
return `http://${placeholder}`
}
if (key === 'path' && value.startsWith('/')) return `/${placeholder}`
return placeholder
}
function isFileDescriptor(value: unknown): boolean {
if (Array.isArray(value)) return value.length > 0 && value.every(isFileDescriptor)
return (
value !== null &&
typeof value === 'object' &&
['key', 'path', 'url'].some(
(key) => typeof (value as Record<string, unknown>)[key] === 'string'
)
)
}
function projectFileReferenceLeaves(
value: unknown,
placeholder: string
): { value: unknown; projectedReferences: number } | undefined {
if (!isFileDescriptor(value)) return undefined
const projectedRoot: unknown[] | Record<string, unknown> = Array.isArray(value) ? [] : {}
const pending: Array<{
source: unknown[] | Record<string, unknown>
target: unknown[] | Record<string, unknown>
}> = [{ source: value as unknown[] | Record<string, unknown>, target: projectedRoot }]
let projectedReferences = 0
while (pending.length > 0) {
const { source, target } = pending.pop()!
for (const [key, child] of Object.entries(source)) {
if (child !== null && typeof child === 'object') {
const projectedChild: unknown[] | Record<string, unknown> = Array.isArray(child) ? [] : {}
;(target as Record<string, unknown>)[key] = projectedChild
pending.push({
source: child as unknown[] | Record<string, unknown>,
target: projectedChild,
})
continue
}
const projectedReference =
typeof child === 'string' ? projectFileReferenceString(key, child, placeholder) : undefined
;(target as Record<string, unknown>)[key] = projectedReference ?? child
if (projectedReference !== undefined) projectedReferences += 1
}
}
return { value: projectedRoot, projectedReferences }
}
/**
* Makes only schema-declared structured inputs parseable on a private placeholder projection.
* Raw execution inputs are never passed to or changed by this helper.
*/
export function prepareResolvedSecretProjectedInputs(
inputs: Record<string, unknown>,
inputSchemas: Record<string, unknown> | undefined,
rawInputs?: Record<string, unknown>,
options: { preserveFileDescriptorGrammar?: boolean } = {}
): Record<string, unknown> {
if (!inputSchemas) return inputs
const prepared = { ...inputs }
for (const [key, inputSchema] of Object.entries(inputSchemas)) {
const inputType =
inputSchema && typeof inputSchema === 'object'
? (inputSchema as { type?: unknown }).type
: inputSchema
if (inputType !== 'json' && inputType !== 'array') continue
const value = prepared[key]
if (typeof value === 'string') {
const placeholder = canonicalPlaceholder(value)
const rawValue = rawInputs?.[key]
if (placeholder && typeof rawValue === 'string') {
try {
const parsedRawValue = JSON.parse(rawValue.trim())
if (parsedRawValue !== null && typeof parsedRawValue === 'object') {
const fileProjection = options.preserveFileDescriptorGrammar
? projectFileReferenceLeaves(parsedRawValue, placeholder)
: undefined
prepared[key] = JSON.stringify(
fileProjection && fileProjection.projectedReferences > 0
? fileProjection.value
: projectStructuredLeaves(parsedRawValue, placeholder)
)
continue
}
} catch {}
}
prepared[key] = makeCanonicalPlaceholdersJsonParseable(value)
}
}
return prepared
}
apps/sim/executor/utils/resolved-secret-matcher.test.ts
+23
View File
@@ -8,11 +8,34 @@ import {
OPAQUE_RESOLVED_SECRET_REPLACEMENT,
sanitizeResolvedSecretPrimitive,
sanitizeResolvedSecretString,
scanResolvedSecretString,
} from '@/executor/utils/resolved-secret-matcher'
const PRESERVE_NAMED_PROVENANCE = { preserveNamedProvenanceLabels: true } as const
describe('resolved secret matcher', () => {
it('reports each matched literal once across large repeated content', () => {
const matcher = createResolvedSecretMatcher([
{ plaintext: 'x', replacement: '{{SHORT}}' },
{ plaintext: 'xx', replacement: '{{OVERLAP}}' },
{ plaintext: 'abc', replacement: '{{PREFIX}}' },
{ plaintext: 'bc', replacement: '{{SUFFIX}}' },
])
const matches: string[] = []
expect(matcher).toBeDefined()
if (!matcher) return
expect(
scanResolvedSecretString(
`${'x'.repeat(1_000_001)}abcabc`,
matcher,
(match) => matches.push(match),
4
)
).toBe(4)
expect(matches).toEqual(['x', 'xx', 'abc', 'bc'])
})
it('uses exact matching for typed primitive renderings', () => {
const matcher = createResolvedSecretMatcher([{ plaintext: '23', replacement: '{{TOKEN}}' }])
apps/sim/executor/utils/resolved-secret-matcher.ts
+27 -3
View File
@@ -231,7 +231,7 @@ export function containsResolvedSecretLiteral(
return false
}
/** Visits exact secret literals with the same bounded automaton used by content projection. */
/** Visits each distinct exact secret literal once with the content-projection automaton. */
export function scanResolvedSecretString(
value: string,
matcher: ResolvedSecretMatcher,
@@ -240,16 +240,40 @@ export function scanResolvedSecretString(
): number {
let node = matcher.root
let matchEvents = 0
const matchedPlaintexts = new Set<string>()
const nextUnmatchedOutput = new WeakMap<SecretTrieNode, SecretTrieNode | null>()
const findNextUnmatchedOutput = (
candidate: SecretTrieNode | undefined
): SecretTrieNode | undefined => {
let current = candidate
const exhaustedPath: SecretTrieNode[] = []
while (current?.replacement && matchedPlaintexts.has(current.replacement.plaintext)) {
const cached = nextUnmatchedOutput.get(current)
if (cached !== undefined) {
current = cached ?? undefined
continue
}
exhaustedPath.push(current)
current = current.outputLink
}
for (const exhausted of exhaustedPath) {
nextUnmatchedOutput.set(exhausted, current ?? null)
}
return current
}
for (let index = 0; index < value.length; index += 1) {
node = advanceMatcher(matcher, node, value[index])
let outputNode: SecretTrieNode | undefined = node.replacement ? node : node.outputLink
let outputNode = findNextUnmatchedOutput(node.replacement ? node : node.outputLink)
while (outputNode?.replacement) {
matchEvents += 1
if (matchEvents > maxMatchEvents) {
throw new ResolvedSecretMatcherError('Secret matcher event limit exceeded')
}
matchedPlaintexts.add(outputNode.replacement.plaintext)
onMatch(outputNode.replacement.plaintext)
outputNode = outputNode.outputLink
outputNode = findNextUnmatchedOutput(outputNode)
}
}
return matchEvents
apps/sim/executor/utils/resolved-secret-trace-registry.test.ts
+399 -47
View File
@@ -178,6 +178,234 @@ describe('ResolvedSecretTraceRegistry', () => {
}
})
it('projects only resolver-recorded leaves and never rewrites sibling bytes or object keys', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'TOKEN', plaintext: 'x', encryptedValue: 'encrypted-token' },
])
registry.recordResolvedAtInputPath('TOKEN', 'x', ['prompt'])
registry.recordResolvedInputProjection(['prompt'], 'Box x', 'Box {{TOKEN}}')
const unrelatedNonCloneableInput = () => 'unchanged'
expect(
registry.projectResolvedInputSelection({
prompt: 'Box x',
auxiliary: 'xylophone',
Box: 'unchanged',
unrelatedNonCloneableInput,
})
).toEqual({
complete: true,
value: {
prompt: 'Box {{TOKEN}}',
auxiliary: 'xylophone',
Box: 'unchanged',
unrelatedNonCloneableInput,
},
})
})
it('keeps equal secret values causally bound to their own resolver paths', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'FIRST', plaintext: 'true', encryptedValue: 'encrypted-first' },
{ name: 'SECOND', plaintext: 'true', encryptedValue: 'encrypted-second' },
])
registry.recordResolvedAtInputPath('FIRST', 'true', ['first'])
registry.recordResolvedInputProjection(['first'], 'true', '{{FIRST}}')
registry.recordResolvedAtInputPath('SECOND', 'true', ['second'])
registry.recordResolvedInputProjection(['second'], 'true', '{{SECOND}}')
expect(registry.projectResolvedInputSelection({ first: 'true', second: 'true' })).toEqual({
complete: true,
value: { first: '{{FIRST}}', second: '{{SECOND}}' },
})
expect(registry.exportCommittedProvenanceForInputPaths([['first']])).toMatchObject({
complete: true,
entries: [{ name: 'FIRST', encryptedValue: 'encrypted-first' }],
})
})
it('isolates incomplete provenance to its known input path', async () => {
const registry = new ResolvedSecretTraceRegistry()
expect(
await registry.importProvenanceForValueAtInputPath(
{ version: 1 },
'unknown-value',
['tools', '0', 'params', 'apiKey'],
{ trusted: true }
)
).toEqual({ success: false, matched: false })
expect(registry.isComplete()).toBe(false)
expect(registry.projectResolvedInputSelection({ userPrompt: 'Public prompt' })).toEqual({
complete: true,
value: { userPrompt: 'Public prompt' },
})
expect(registry.exportCommittedProvenanceForInputPaths([['userPrompt']])).toEqual({
version: 1,
complete: true,
entries: [],
})
expect(registry.forkForInputPaths([['userPrompt']]).isComplete()).toBe(true)
})
it('propagates authenticated incomplete provenance without classifying it as malformed', async () => {
const scope = { userId: 'user-1', workspaceId: 'workspace-1' }
const registry = new ResolvedSecretTraceRegistry([], scope)
expect(
await registry.importProvenanceForValueAtInputPath(
{ version: 1, complete: false, entries: [], scope },
'untrusted value',
['toolResult'],
{ trusted: true }
)
).toEqual({ success: true, matched: false })
expect(registry.forkForInputPaths([['publicPrompt']]).isComplete()).toBe(true)
expect(registry.forkForInputPaths([['toolResult']]).isComplete()).toBe(false)
})
it('localizes a failed exact resolution when the resolver supplies its input path', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'TOKEN', plaintext: 'expected', encryptedValue: 'encrypted-token' },
])
expect(
registry.recordResolvedAtInputPath('TOKEN', 'unexpected', ['tools', '0', 'params', 'apiKey'])
).toBe(false)
expect(registry.projectResolvedInputSelection({ userPrompt: 'Public prompt' })).toEqual({
complete: true,
value: { userPrompt: 'Public prompt' },
})
expect(registry.forkForInputPaths([['userPrompt']]).isComplete()).toBe(true)
expect(registry.forkForInputPaths([['tools', '0', 'params']]).isComplete()).toBe(false)
expect(registry.getModelEgressSnapshot()).toEqual({ complete: false })
})
it('fails closed for a selected unknown path and arbitrary output projection', async () => {
const scope = { userId: 'user-1', workspaceId: 'workspace-1' }
const registry = new ResolvedSecretTraceRegistry([], scope)
await registry.importProvenanceForValueAtInputPath(
{ version: 1 },
'unknown-value',
['tools', '0', 'params', 'apiKey'],
{ trusted: true }
)
expect(
registry.projectResolvedInputSelection({ tools: [{ params: { apiKey: 'value' } }] })
).toEqual({ complete: false })
expect(registry.exportCommittedProvenanceForInputPaths([['tools', '0', 'params']])).toEqual({
version: 1,
complete: false,
entries: [],
scope,
})
expect(registry.forkForInputPaths([['tools', '0', 'params']]).isComplete()).toBe(false)
expect(registry.getModelEgressSnapshot()).toEqual({ complete: false })
expect(registry.exportProvenanceForValue('arbitrary output')).toEqual({
version: 1,
complete: false,
entries: [],
scope,
})
})
it('propagates only selected input-path entries when explicitly requested', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'SELECTED', plaintext: 'selected', encryptedValue: 'encrypted-selected' },
{ name: 'UNSELECTED', plaintext: 'unselected', encryptedValue: 'encrypted-unselected' },
])
registry.recordResolvedAtInputPath('SELECTED', 'selected', ['selected'])
registry.recordResolvedAtInputPath('UNSELECTED', 'unselected', ['unselected'])
const ordinaryFork = registry.forkForInputPaths([['selected']])
expect(ordinaryFork.forkForPropagatedEntries().exportProvenance().entries).toEqual([])
const propagatedFork = registry.forkForInputPaths([['selected']], { propagated: true })
expect(propagatedFork.forkForPropagatedEntries().exportProvenance().entries).toEqual([
{ name: 'SELECTED', encryptedValue: 'encrypted-selected' },
])
expect(propagatedFork.exportProvenance().entries).not.toContainEqual({
name: 'UNSELECTED',
encryptedValue: 'encrypted-unselected',
})
})
it('preserves exact paths through renamed and parsed parameter transforms', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'FIRST', plaintext: 'true', encryptedValue: 'encrypted-first' },
{ name: 'SECOND', plaintext: 'true', encryptedValue: 'encrypted-second' },
{ name: 'UNUSED', plaintext: 'true', encryptedValue: 'encrypted-unused' },
])
registry.recordResolvedAtInputPath('FIRST', 'true', ['rowTemplate'])
registry.recordResolvedAtInputPath('SECOND', 'true', ['rowTemplate'])
registry.recordResolvedInputProjection(
['rowTemplate'],
'{"first":true,"second":true,"public":true}',
'{"first":{{FIRST}},"second":{{SECOND}},"public":true}'
)
registry.recordTransformedInputProjection(
{ data: { first: true, second: true, public: true } },
{ data: { first: '{{FIRST}}', second: '{{SECOND}}', public: true } }
)
expect(
registry.projectResolvedInputSelection({
data: { first: true, second: true, public: true },
})
).toEqual({
complete: true,
value: {
data: { first: '{{FIRST}}', second: '{{SECOND}}', public: true },
},
})
expect(registry.exportCommittedProvenanceForInputPaths([['data', 'first']])).toMatchObject({
complete: true,
entries: [{ name: 'FIRST', encryptedValue: 'encrypted-first' }],
})
expect(registry.exportCommittedProvenanceForInputPaths([['data', 'second']])).toMatchObject({
complete: true,
entries: [{ name: 'SECOND', encryptedValue: 'encrypted-second' }],
})
expect(registry.exportCommittedProvenanceForInputPaths([['data', 'public']])).toMatchObject({
complete: true,
entries: [],
})
})
it('fails closed when independent secret paths collapse into one transformed string', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'FIRST', plaintext: 'first', encryptedValue: 'encrypted-first' },
{ name: 'SECOND', plaintext: 'second', encryptedValue: 'encrypted-second' },
])
registry.recordResolvedAtInputPath('FIRST', 'first', ['first'])
registry.recordResolvedInputProjection(['first'], 'first', '{{FIRST}}')
registry.recordResolvedAtInputPath('SECOND', 'second', ['second'])
registry.recordResolvedInputProjection(['second'], 'second', '{{SECOND}}')
registry.recordTransformedInputProjection(
{ combined: 'first:second' },
{ combined: '{{FIRST}}:second' }
)
registry.recordTransformedInputProjection(
{ combined: 'first:second' },
{ combined: 'first:{{SECOND}}' }
)
expect(registry.isComplete()).toBe(false)
expect(registry.projectResolvedInputSelection({ unrelated: 'public' })).toEqual({
complete: true,
value: { unrelated: 'public' },
})
expect(registry.projectResolvedInputSelection({ combined: 'first:second' })).toEqual({
complete: false,
})
expect(registry.getModelEgressSnapshot()).toEqual({ complete: false })
})
it('does not invalidate the model matcher for duplicate activations', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'API_KEY', plaintext: 'secret-value', encryptedValue: 'encrypted-value' },
@@ -248,6 +476,42 @@ describe('ResolvedSecretTraceRegistry', () => {
}
})
it('keeps a named anonymous secret distinct from anonymous provenance', async () => {
mockDecryptSecret.mockResolvedValueOnce({ decrypted: 'same-secret' })
const registry = new ResolvedSecretTraceRegistry(
[{ name: 'anonymous', plaintext: 'same-secret', encryptedValue: 'shared-ciphertext' }],
{ userId: 'user-1', workspaceId: 'workspace-1' }
)
registry.recordResolved('anonymous', 'same-secret')
await registry.importProvenance(
{
version: 1,
complete: true,
entries: [{ encryptedValue: 'shared-ciphertext' }],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
},
{ trusted: true, anonymous: true }
)
expect(registry.exportCommittedProvenanceForValue('same-secret')).toEqual({
version: 1,
complete: true,
entries: [
{ encryptedValue: 'shared-ciphertext' },
{ name: 'anonymous', encryptedValue: 'shared-ciphertext' },
],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
})
const snapshot = registry.getModelEgressSnapshot()
expect(snapshot.complete).toBe(true)
if (snapshot.complete) {
expect(snapshot.matches).toContainEqual({
plaintext: 'same-secret',
replacement: ANONYMOUS_SECRET_TRACE_REPLACEMENT,
})
}
})
it('projects committed provenance while temporary activations are pending', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'API_KEY', plaintext: 'secret-value', encryptedValue: 'encrypted-value' },
@@ -292,45 +556,6 @@ describe('ResolvedSecretTraceRegistry', () => {
})
})
it('seeds a tool child only with active provenance present in that tool input', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'INPUT', plaintext: 'input-secret', encryptedValue: 'input-ciphertext' },
{ name: 'UNRELATED', plaintext: 'Test', encryptedValue: 'unrelated-ciphertext' },
])
registry.recordResolved('INPUT', 'input-secret')
registry.recordResolved('UNRELATED', 'Test')
const child = registry.forkForToolInput({ authorization: 'Bearer input-secret' })
expect(child.getActiveMatches()).toEqual([
{ plaintext: 'input-secret', replacement: '{{INPUT}}' },
])
expect(child.recordResolved('UNRELATED', 'Test')).toBe(true)
})
it('forks independent roots without treating static param names or array indexes as data', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'PROMPT', plaintext: 'prompt', encryptedValue: 'prompt-ciphertext' },
{ name: 'ZERO', plaintext: '0', encryptedValue: 'zero-ciphertext' },
{ name: 'VALUE', plaintext: 'input-secret', encryptedValue: 'value-ciphertext' },
])
registry.recordResolved('PROMPT', 'prompt')
registry.recordResolved('ZERO', '0')
registry.recordResolved('VALUE', 'input-secret')
const child = registry.forkForToolInputValues(['safe', { nested: 'input-secret' }])
expect(child.getActiveMatches()).toEqual([
{ plaintext: 'input-secret', replacement: '{{VALUE}}' },
])
expect(registry.forkForToolInputValues([{ prompt: 'safe' }]).getActiveMatches()).toEqual([
{ plaintext: 'prompt', replacement: '{{PROMPT}}' },
])
expect(registry.forkForToolInputValues([0]).getActiveMatches()).toEqual([
{ plaintext: '0', replacement: '{{ZERO}}' },
])
})
it('uses the workspace catalog entry when personal and workspace names conflict', async () => {
const registry = await createResolvedSecretTraceRegistry({
personalEncrypted: { SHARED: 'personal-encrypted' },
@@ -701,6 +926,102 @@ describe('ResolvedSecretTraceRegistry', () => {
})
})
it('exports active provenance for a registered legacy runtime alias', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'API-KEY', plaintext: 'secret-value', encryptedValue: 'present-ciphertext' },
])
registry.recordResolved('API-KEY', 'secret-value')
expect(registry.exportCommittedProvenanceForValue('prefix __var_API_KEY suffix')).toEqual({
version: 1,
complete: true,
entries: [{ name: 'API-KEY', encryptedValue: 'present-ciphertext' }],
})
})
it('ignores repeated unrelated runtime aliases without exhausting the scan budget', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'API_KEY', plaintext: 'secret-value', encryptedValue: 'present-ciphertext' },
])
registry.recordResolved('API_KEY', 'secret-value')
expect(
registry.exportCommittedProvenanceForValue(`${'__var_Z '.repeat(1_000_001)}__var_API_KEY`)
).toEqual({
version: 1,
complete: true,
entries: [{ name: 'API_KEY', encryptedValue: 'present-ciphertext' }],
})
})
it('matches legacy runtime aliases as complete tokens instead of prefixes', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'A', plaintext: 'secret-a', encryptedValue: 'ciphertext-a' },
{ name: 'API_KEY', plaintext: 'secret-api', encryptedValue: 'ciphertext-api' },
])
registry.recordResolved('A', 'secret-a')
registry.recordResolved('API_KEY', 'secret-api')
expect(registry.exportCommittedProvenanceForValue('__var_API_KEY')).toEqual({
version: 1,
complete: true,
entries: [{ name: 'API_KEY', encryptedValue: 'ciphertext-api' }],
})
})
it('conservatively retains every secret mapped to a colliding runtime alias', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'API-KEY', plaintext: 'first-secret', encryptedValue: 'first-ciphertext' },
{ name: 'API_KEY', plaintext: 'second-secret', encryptedValue: 'second-ciphertext' },
])
registry.recordResolved('API-KEY', 'first-secret')
registry.recordResolved('API_KEY', 'second-secret')
expect(registry.exportCommittedProvenanceForValue('__var_API_KEY')).toEqual({
version: 1,
complete: true,
entries: [
{ name: 'API-KEY', encryptedValue: 'first-ciphertext' },
{ name: 'API_KEY', encryptedValue: 'second-ciphertext' },
],
})
})
it('retains an alias-specific entry when multiple names share one plaintext', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'FIRST', plaintext: 'shared-secret', encryptedValue: 'first-ciphertext' },
{ name: 'SECOND', plaintext: 'shared-secret', encryptedValue: 'second-ciphertext' },
])
registry.recordResolved('FIRST', 'shared-secret')
registry.recordResolved('SECOND', 'shared-secret')
expect(registry.exportCommittedProvenanceForValue('__var_SECOND')).toEqual({
version: 1,
complete: true,
entries: [{ name: 'SECOND', encryptedValue: 'second-ciphertext' }],
})
})
it('conservatively retains every active secret that shares a raw plaintext literal', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'FIRST', plaintext: 'true', encryptedValue: 'first-ciphertext' },
{ name: 'SECOND', plaintext: 'true', encryptedValue: 'second-ciphertext' },
])
registry.recordResolved('FIRST', 'true')
registry.recordResolved('SECOND', 'true')
const expected = {
version: 1 as const,
complete: true,
entries: [
{ name: 'FIRST', encryptedValue: 'first-ciphertext' },
{ name: 'SECOND', encryptedValue: 'second-ciphertext' },
],
}
expect(registry.exportCommittedProvenanceForValue('true')).toEqual(expected)
expect(registry.exportCommittedProvenanceForValue(true)).toEqual(expected)
})
it('exports active numeric, boolean, and null literals crossing a value boundary', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'NUMBER', plaintext: '1234', encryptedValue: 'number-ciphertext' },
@@ -800,11 +1121,13 @@ describe('ResolvedSecretTraceRegistry', () => {
registry.recordResolved('A_TOKEN', 'same')
registry.recordResolved('EMPTY', '')
expect(registry.getActiveMatches()).toEqual([{ plaintext: 'same', replacement: '{{A_TOKEN}}' }])
expect(registry.getActiveMatches()).toEqual([
{ plaintext: 'same', replacement: ANONYMOUS_SECRET_TRACE_REPLACEMENT },
])
registry.recordResolved('A', 'A')
expect(registry.getActiveMatches()).toEqual([
{ plaintext: 'same', replacement: '{{A_TOKEN}}' },
{ plaintext: 'same', replacement: ANONYMOUS_SECRET_TRACE_REPLACEMENT },
{ plaintext: 'A', replacement: '{{A}}' },
])
})
@@ -825,6 +1148,22 @@ describe('ResolvedSecretTraceRegistry', () => {
expect(registry.exportProvenance().entries).toEqual([])
})
it('does not poison unrelated inputs when dormant catalog entries exceed the hard cap', () => {
const entries = Array.from({ length: 10_001 }, (_, index) => ({
name: `SECRET_${index}`,
plaintext: `value-${index}`,
encryptedValue: `ciphertext-${index}`,
}))
const registry = new ResolvedSecretTraceRegistry(entries)
expect(registry.isComplete()).toBe(true)
expect(registry.recordResolvedAtInputPath('SECRET_10000', 'value-10000', ['userPrompt'])).toBe(
false
)
expect(registry.forkForInputPaths([['systemPrompt']]).isComplete()).toBe(true)
expect(registry.forkForInputPaths([['userPrompt']]).isComplete()).toBe(false)
})
it('bounds provenance by serialized JSON bytes including control-character escapes', () => {
const encryptedValue = '\u0000'.repeat(1_400_000)
const provenance: ResolvedSecretTraceProvenanceV1 = {
@@ -892,7 +1231,7 @@ describe('ResolvedSecretTraceRegistry', () => {
)
})
it('stops consuming a dormant catalog when its entry cap is exceeded', () => {
it('does not let a large dormant catalog poison unrelated execution provenance', () => {
let yieldedEntries = 0
function* catalogEntries() {
for (let index = 0; index < 20_000; index++) {
@@ -908,11 +1247,11 @@ describe('ResolvedSecretTraceRegistry', () => {
const registry = new ResolvedSecretTraceRegistry(catalogEntries())
expect(yieldedEntries).toBe(10_001)
expect(registry.isComplete()).toBe(false)
expect(registry.isComplete()).toBe(true)
expect(registry.exportProvenance().entries).toEqual([])
})
it('marks an oversized dormant catalog value incomplete without retaining it', () => {
it('keeps an oversized dormant value inert until that exact secret is resolved', () => {
const oversizedPlaintext = 'x'.repeat(8 * 1024 * 1024)
const registry = new ResolvedSecretTraceRegistry([
{
@@ -920,10 +1259,23 @@ describe('ResolvedSecretTraceRegistry', () => {
plaintext: oversizedPlaintext,
encryptedValue: 'ciphertext',
},
{
name: 'NORMAL',
plaintext: 'normal-secret',
encryptedValue: 'normal-ciphertext',
},
])
expect(registry.isComplete()).toBe(false)
expect(registry.recordResolved('OVERSIZED', oversizedPlaintext)).toBe(false)
expect(registry.getActiveMatches()).toEqual([])
expect(registry.isComplete()).toBe(true)
expect(registry.getModelEgressSnapshot()).toEqual({ complete: true, matches: [] })
expect(registry.recordResolvedAtInputPath('NORMAL', 'normal-secret', ['systemPrompt'])).toBe(
true
)
expect(
registry.recordResolvedAtInputPath('OVERSIZED', oversizedPlaintext, ['userPrompt'])
).toBe(false)
expect(registry.forkForInputPaths([['systemPrompt']]).isComplete()).toBe(true)
expect(registry.forkForInputPaths([['userPrompt']]).isComplete()).toBe(false)
expect(registry.getModelEgressSnapshot()).toEqual({ complete: false })
})
})
apps/sim/executor/utils/resolved-secret-trace-registry.ts
+760 -160
View File
File diff suppressed because it is too large Load Diff
apps/sim/executor/variables/resolver.test.ts
+109
View File
@@ -3,6 +3,7 @@
*/
import { loggerMock } from '@sim/testing'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import { projectResolvedModelInput } from '@/lib/execution/model-input-provenance'
import {
LARGE_ARRAY_MANIFEST_VERSION,
type LargeArrayManifest,
@@ -26,6 +27,10 @@ vi.mock('@/lib/execution/payloads/store', () => ({
materializeLargeValueRef: vi.fn(),
}))
vi.mock('@/lib/core/security/encryption', () => ({
decryptSecret: vi.fn(async (encryptedValue: string) => ({ decrypted: encryptedValue })),
}))
function createBlock(id: string, name: string, type: string, params = {}): SerializedBlock {
return {
id,
@@ -115,6 +120,7 @@ function createResolver(
return {
block: functionBlock,
ctx,
state,
resolver: new VariableResolver(workflow, {}, state, options),
}
}
@@ -187,6 +193,109 @@ describe('VariableResolver function block inputs', () => {
])
})
it('binds propagated references to exact model-selected inputs without changing runtime values', async () => {
const secret = 'x'
const provenance = {
version: 1 as const,
complete: true,
entries: [{ name: 'TOKEN', encryptedValue: secret }],
}
const producer = createBlock('producer', 'Producer', BlockType.API)
const loop = createBlock('loop-1', 'Loop1', BlockType.LOOP)
const parallel = createBlock('parallel-1', 'Parallel1', BlockType.PARALLEL)
const consumer = createBlock('consumer', 'Consumer', BlockType.API)
const workflowVariables = {
'var-1': { id: 'var-1', name: 'token', type: 'string', value: secret },
}
const workflow: SerializedWorkflow = {
version: '1',
blocks: [producer, loop, parallel, consumer],
connections: [],
loops: {
'loop-1': { id: 'loop-1', nodes: [], iterations: 1, loopType: 'for' },
},
parallels: {
'parallel-1': {
id: 'parallel-1',
nodes: [],
parallelType: 'count',
count: 1,
},
},
}
const state = new ExecutionState()
state.setBlockOutput('producer', { result: secret }, 0, provenance)
state.setBlockOutput('loop-1', { results: [secret] }, 0, provenance)
state.setBlockOutput('parallel-1', { results: [secret] }, 0, provenance)
const registry = new ResolvedSecretTraceRegistry()
const ctx = {
blockStates: state.getBlockStates(),
blockLogs: [],
environmentVariables: {},
workflowVariables,
workflowVariableResolvedSecretTraceProvenance: { 'var-1': provenance },
resolvedSecretTraceRegistry: registry,
decisions: { router: new Map(), condition: new Map() },
loopExecutions: new Map(),
parallelExecutions: new Map(),
executedBlocks: new Set(),
activeExecutionPath: new Set(),
completedLoops: new Set(),
metadata: {},
} as ExecutionContext
const resolver = new VariableResolver(workflow, workflowVariables, state, {
navigatePathAsync,
})
const inputs = {
blockPrompt: 'Box: <Producer.result>',
workflowPrompt: 'Workflow: <variable.token>',
loopPrompt: 'Loop: <Loop1.results[0]>',
parallelPrompt: 'Parallel: <Parallel1.results[0]>',
}
const resolved = await resolver.resolveInputs(ctx, consumer.id, inputs, consumer)
expect(resolved).toEqual({
blockPrompt: `Box: ${secret}`,
workflowPrompt: `Workflow: ${secret}`,
loopPrompt: `Loop: ${secret}`,
parallelPrompt: `Parallel: ${secret}`,
})
const projection = projectResolvedModelInput(
registry,
resolved,
Object.keys(inputs).map((key) => [key])
)
expect(projection.complete).toBe(true)
if (!projection.complete) throw new Error('Expected complete model projection')
expect(projection.value).toEqual(inputs)
})
it('preserves the destination path when resolving one whole reference directly', async () => {
const secret = 'resolved-secret'
const provenance = {
version: 1 as const,
complete: true,
entries: [{ name: 'TOKEN', encryptedValue: secret }],
}
const { ctx, resolver, state } = createResolver()
state.setBlockOutput('producer', { result: secret }, 0, provenance)
ctx.blockStates = state.getBlockStates()
const registry = new ResolvedSecretTraceRegistry()
ctx.resolvedSecretTraceRegistry = registry
await expect(
resolver.resolveSingleReference(ctx, 'function', '<Producer.result>', undefined, {
inputPath: ['prompt'],
})
).resolves.toBe(secret)
expect(registry.exportCommittedProvenanceForInputPaths([['prompt']])).toEqual(provenance)
expect(registry.projectResolvedInputSelection({ prompt: secret })).toEqual({
complete: true,
value: { prompt: '<Producer.result>' },
})
})
it('returns empty inputs when params are missing', async () => {
const { block, ctx, resolver } = createResolver()
apps/sim/executor/variables/resolver.ts
+107 -28
View File
@@ -215,11 +215,15 @@ export class VariableResolver {
resolved[key] = resolvedItems
display[key] = displayItems
} else {
resolved[key] = await this.resolveValue(ctx, currentNodeId, value, undefined, block)
resolved[key] = await this.resolveValue(ctx, currentNodeId, value, undefined, block, {
inputPath: [key],
})
display[key] = resolved[key]
}
} else {
resolved[key] = await this.resolveValue(ctx, currentNodeId, value, undefined, block)
resolved[key] = await this.resolveValue(ctx, currentNodeId, value, undefined, block, {
inputPath: [key],
})
display[key] = resolved[key]
}
}
@@ -256,14 +260,20 @@ export class VariableResolver {
if (Array.isArray(conditions)) {
resolved.conditions = await Promise.all(
conditions.map(async (condition) => {
conditions.map(async (condition, conditionIndex) => {
if (!condition || typeof condition !== 'object') return condition
const value = Reflect.get(condition, 'value')
return {
...condition,
value:
typeof value === 'string'
? await this.resolveTemplateWithoutConditionFormatting(ctx, currentNodeId, value)
? await this.resolveTemplateWithoutConditionFormatting(
ctx,
currentNodeId,
value,
undefined,
['conditions', String(conditionIndex), 'value']
)
: value,
}
})
@@ -274,7 +284,8 @@ export class VariableResolver {
currentNodeId,
conditions,
undefined,
block
block,
{ inputPath: ['conditions'] }
)
}
}
@@ -285,6 +296,7 @@ export class VariableResolver {
}
resolved[key] = await this.resolveValue(ctx, currentNodeId, value, undefined, block, {
allowLargeValueRefs: this.canResolveInputToLargeValueRef(block, key),
inputPath: [key],
})
}
return resolved
@@ -307,7 +319,7 @@ export class VariableResolver {
currentNodeId: string,
reference: string,
loopScope?: LoopScope,
options: { allowLargeValueRefs?: boolean } = {}
options: { allowLargeValueRefs?: boolean; inputPath?: readonly string[] } = {}
): Promise<any> {
if (typeof reference === 'string') {
const trimmed = reference.trim()
@@ -318,17 +330,21 @@ export class VariableResolver {
currentNodeId,
loopScope,
allowLargeValueRefs: options.allowLargeValueRefs,
inputPath: options.inputPath,
}
const result = await this.resolveReference(trimmed, resolutionContext)
if (result === RESOLVED_EMPTY) {
return null
}
return result
const resolved = result === RESOLVED_EMPTY ? null : result
ctx.resolvedSecretTraceRegistry?.recordResolvedInputProjection(
options.inputPath,
resolved,
trimmed
)
return resolved
}
}
return this.resolveValue(ctx, currentNodeId, reference, loopScope)
return this.resolveValue(ctx, currentNodeId, reference, loopScope, undefined, options)
}
private async resolveValue(
@@ -337,7 +353,7 @@ export class VariableResolver {
value: any,
loopScope?: LoopScope,
block?: SerializedBlock,
options: { allowLargeValueRefs?: boolean } = {}
options: { allowLargeValueRefs?: boolean; inputPath?: readonly string[] } = {}
): Promise<any> {
if (value === null || value === undefined) {
return value
@@ -345,16 +361,24 @@ export class VariableResolver {
if (Array.isArray(value)) {
return Promise.all(
value.map((v) => this.resolveValue(ctx, currentNodeId, v, loopScope, block, options))
value.map((v, index) =>
this.resolveValue(ctx, currentNodeId, v, loopScope, block, {
...options,
inputPath: options.inputPath ? [...options.inputPath, String(index)] : undefined,
})
)
)
}
if (typeof value === 'object') {
const entries = await Promise.all(
Object.entries(value).map(async ([key, val]) => [
key,
await this.resolveValue(ctx, currentNodeId, val, loopScope, block, options),
])
Object.entries(value).map(async ([key, val]) => {
const resolvedValue = await this.resolveValue(ctx, currentNodeId, val, loopScope, block, {
...options,
inputPath: options.inputPath ? [...options.inputPath, key] : undefined,
})
return [key, resolvedValue]
})
)
return Object.fromEntries(entries)
}
@@ -1237,7 +1261,7 @@ export class VariableResolver {
template: string,
loopScope?: LoopScope,
block?: SerializedBlock,
options: { allowLargeValueRefs?: boolean } = {}
options: { allowLargeValueRefs?: boolean; inputPath?: readonly string[] } = {}
): Promise<string> {
const resolutionContext: ResolutionContext = {
executionContext: ctx,
@@ -1245,6 +1269,7 @@ export class VariableResolver {
currentNodeId,
loopScope,
allowLargeValueRefs: options.allowLargeValueRefs,
inputPath: options.inputPath,
}
let replacementError: Error | null = null
@@ -1257,28 +1282,48 @@ export class VariableResolver {
| undefined)
: undefined
let result = await replaceValidReferencesAsync(template, async (match) => {
let projectedReferenceResult = ''
let projectedReferenceCursor = 0
let result = await replaceValidReferencesAsync(template, async (match, index) => {
if (replacementError) return match
projectedReferenceResult += template.slice(projectedReferenceCursor, index)
projectedReferenceCursor = index + match.length
let containsResolvedSecret = false
const referenceContext: ResolutionContext = {
...resolutionContext,
onResolvedSecretReference: () => {
containsResolvedSecret = true
},
}
try {
const resolved = await this.resolveReference(match, resolutionContext)
const resolved = await this.resolveReference(match, referenceContext)
if (resolved === undefined) {
projectedReferenceResult += match
return match
}
if (resolved === RESOLVED_EMPTY) {
if (blockType === BlockType.FUNCTION) {
return this.blockResolver.formatValueForBlock(null, blockType, language)
const formatted = this.blockResolver.formatValueForBlock(null, blockType, language)
projectedReferenceResult += formatted
return formatted
}
projectedReferenceResult += ''
return ''
}
return this.blockResolver.formatValueForBlock(resolved, blockType, language)
const formatted = this.blockResolver.formatValueForBlock(resolved, blockType, language)
projectedReferenceResult += containsResolvedSecret ? match : formatted
return formatted
} catch (error) {
replacementError = toError(error)
projectedReferenceResult += match
return match
}
})
projectedReferenceResult += template.slice(projectedReferenceCursor)
if (replacementError !== null) {
throw replacementError
@@ -1288,6 +1333,11 @@ export class VariableResolver {
const resolved = await this.resolveReference(match, resolutionContext)
return typeof resolved === 'string' ? resolved : match
})
ctx.resolvedSecretTraceRegistry?.recordResolvedInputProjection(
options.inputPath,
result,
projectedReferenceResult
)
return result
}
@@ -1295,27 +1345,43 @@ export class VariableResolver {
ctx: ExecutionContext,
currentNodeId: string,
template: string,
loopScope?: LoopScope
loopScope?: LoopScope,
inputPath?: readonly string[]
): Promise<string> {
const resolutionContext: ResolutionContext = {
executionContext: ctx,
executionState: this.state,
currentNodeId,
loopScope,
inputPath,
}
let replacementError: Error | null = null
let result = await replaceValidReferencesAsync(template, async (match) => {
let projectedReferenceResult = ''
let projectedReferenceCursor = 0
let result = await replaceValidReferencesAsync(template, async (match, index) => {
if (replacementError) return match
projectedReferenceResult += template.slice(projectedReferenceCursor, index)
projectedReferenceCursor = index + match.length
let containsResolvedSecret = false
const referenceContext: ResolutionContext = {
...resolutionContext,
onResolvedSecretReference: () => {
containsResolvedSecret = true
},
}
try {
const resolved = await this.resolveReference(match, resolutionContext)
const resolved = await this.resolveReference(match, referenceContext)
if (resolved === undefined) {
projectedReferenceResult += match
return match
}
if (resolved === RESOLVED_EMPTY) {
projectedReferenceResult += 'null'
return 'null'
}
@@ -1327,17 +1393,25 @@ export class VariableResolver {
.replace(/\r/g, '\\r')
.replace(/\u2028/g, '\\u2028')
.replace(/\u2029/g, '\\u2029')
return `'${escaped}'`
const formatted = `'${escaped}'`
projectedReferenceResult += containsResolvedSecret ? match : formatted
return formatted
}
if (typeof resolved === 'object' && resolved !== null) {
return JSON.stringify(resolved)
const formatted = JSON.stringify(resolved)
projectedReferenceResult += containsResolvedSecret ? match : formatted
return formatted
}
return String(resolved)
const formatted = String(resolved)
projectedReferenceResult += containsResolvedSecret ? match : formatted
return formatted
} catch (error) {
replacementError = toError(error)
projectedReferenceResult += match
return match
}
})
projectedReferenceResult += template.slice(projectedReferenceCursor)
if (replacementError !== null) {
throw replacementError
@@ -1347,6 +1421,11 @@ export class VariableResolver {
const resolved = await this.resolveReference(match, resolutionContext)
return typeof resolved === 'string' ? resolved : match
})
ctx.resolvedSecretTraceRegistry?.recordResolvedInputProjection(
inputPath,
result,
projectedReferenceResult
)
return result
}
apps/sim/executor/variables/resolvers/block.test.ts
+54
View File
@@ -1,7 +1,9 @@
import { describe, expect, it, vi } from 'vitest'
import { isLargeValueRef } from '@/lib/execution/payloads/large-value-ref'
import { compactExecutionPayload } from '@/lib/execution/payloads/serializer'
import { ExecutionState } from '@/executor/execution/state'
import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
import { navigatePathAsync } from '@/executor/variables/resolvers/reference-async.server'
import { BlockResolver } from './block'
import { RESOLVED_EMPTY, type ResolutionContext } from './reference'
@@ -640,6 +642,58 @@ describe('BlockResolver', () => {
{ plaintext: 'secret-value', replacement: '{{API_KEY}}' },
])
})
it('filters compacted block candidates against the exact selected leaf', async () => {
const workflow = createTestWorkflow([{ id: 'source' }])
const resolver = new BlockResolver(workflow, navigatePathAsync)
const compacted = await compactExecutionPayload(
{
result: {
huge: 'p'.repeat(9 * 1024 * 1024),
public: 'ok',
secret: 'secret-value',
},
},
{
preserveRoot: true,
workspaceId: 'workspace-1',
workflowId: 'workflow-1',
executionId: 'execution-1',
}
)
expect(isLargeValueRef(compacted.result.huge)).toBe(true)
const candidateProvenance = {
version: 1 as const,
complete: true,
entries: [{ name: 'API_KEY', encryptedValue: 'secret-value' }],
}
const publicRegistry = new ResolvedSecretTraceRegistry()
const publicContext = createTestContext('current')
publicContext.inputPath = ['prompt']
publicContext.executionContext.resolvedSecretTraceRegistry = publicRegistry
publicContext.executionState.setBlockOutput('source', compacted, 0, candidateProvenance)
await expect(resolver.resolveAsync('<source.result.public>', publicContext)).resolves.toBe(
'ok'
)
expect(publicRegistry.isComplete()).toBe(true)
expect(publicRegistry.getActiveMatches()).toEqual([])
const secretRegistry = new ResolvedSecretTraceRegistry()
const secretContext = createTestContext('current')
secretContext.inputPath = ['prompt']
secretContext.executionContext.resolvedSecretTraceRegistry = secretRegistry
secretContext.executionState.setBlockOutput('source', compacted, 0, candidateProvenance)
await expect(resolver.resolveAsync('<source.result.secret>', secretContext)).resolves.toBe(
'secret-value'
)
expect(secretRegistry.isComplete()).toBe(true)
expect(secretRegistry.getActiveMatches()).toEqual([
{ plaintext: 'secret-value', replacement: '{{API_KEY}}' },
])
})
})
describe('formatValueForBlock', () => {
apps/sim/executor/variables/resolvers/block.ts
+8 -5
View File
@@ -332,11 +332,14 @@ export class BlockResolver implements Resolver {
) {
return value
}
await context.executionContext.resolvedSecretTraceRegistry.importProvenanceForValue(
state.resolvedSecretTraceProvenance,
value,
{ trusted: true }
)
const imported =
await context.executionContext.resolvedSecretTraceRegistry.importProvenanceForValueAtInputPath(
state.resolvedSecretTraceProvenance,
value,
context.inputPath,
{ trusted: true }
)
if (imported.matched) context.onResolvedSecretReference?.()
return value
}
apps/sim/executor/variables/resolvers/env.ts
+5 -1
View File
@@ -17,7 +17,11 @@ export class EnvResolver implements Resolver {
return reference
}
if (Object.hasOwn(context.executionContext.environmentVariables, varName)) {
context.executionContext.resolvedSecretTraceRegistry?.recordResolved(varName, value)
context.executionContext.resolvedSecretTraceRegistry?.recordResolvedAtInputPath(
varName,
value,
context.inputPath
)
}
return value
}
apps/sim/executor/variables/resolvers/loop.ts
+7 -1
View File
@@ -279,7 +279,13 @@ export class LoopResolver implements Resolver {
const resolvedValue = await value
const registry = context.executionContext.resolvedSecretTraceRegistry
if (!registry || !provenance) return resolvedValue
await registry.importProvenanceForValue(provenance, resolvedValue, { trusted: true })
const imported = await registry.importProvenanceForValueAtInputPath(
provenance,
resolvedValue,
context.inputPath,
{ trusted: true }
)
if (imported.matched) context.onResolvedSecretReference?.()
return resolvedValue
}
apps/sim/executor/variables/resolvers/parallel.ts
+7 -1
View File
@@ -386,7 +386,13 @@ export class ParallelResolver implements Resolver {
const resolvedValue = await value
const registry = context.executionContext.resolvedSecretTraceRegistry
if (!registry || !provenance) return resolvedValue
await registry.importProvenanceForValue(provenance, resolvedValue, { trusted: true })
const imported = await registry.importProvenanceForValueAtInputPath(
provenance,
resolvedValue,
context.inputPath,
{ trusted: true }
)
if (imported.matched) context.onResolvedSecretReference?.()
return resolvedValue
}
apps/sim/executor/variables/resolvers/reference.ts
+2
View File
@@ -34,6 +34,8 @@ export interface ResolutionContext {
currentNodeId: string
loopScope?: LoopScope
allowLargeValueRefs?: boolean
inputPath?: readonly string[]
onResolvedSecretReference?: () => void
}
export interface Resolver {
apps/sim/executor/variables/resolvers/workflow.ts
+7 -1
View File
@@ -129,7 +129,13 @@ export class WorkflowResolver implements Resolver {
context.executionContext.workflowVariableResolvedSecretTraceProvenance?.[variableId]
if (!registry || !provenance) return value
await registry.importProvenanceForValue(provenance, value, { trusted: true })
const imported = await registry.importProvenanceForValueAtInputPath(
provenance,
value,
context.inputPath,
{ trusted: true }
)
if (imported.matched) context.onResolvedSecretReference?.()
return value
}
apps/sim/lib/copilot/chat/post.ts
+3 -41
View File
@@ -4,7 +4,6 @@ import { copilotChats } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { getErrorMessage } from '@sim/utils/errors'
import { generateId } from '@sim/utils/id'
import { isPlainRecord } from '@sim/utils/object'
import { eq } from 'drizzle-orm'
import { type NextRequest, NextResponse } from 'next/server'
import { z } from 'zod'
@@ -66,8 +65,6 @@ import {
isWorkspaceAccessDeniedError,
type PermissionType,
} from '@/lib/workspaces/permissions/utils'
import { projectResolvedSecretModelContent } from '@/executor/utils/resolved-secret-content-projection'
import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
import type { ChatContext } from '@/stores/panel'
export const maxDuration = 3600
@@ -522,36 +519,6 @@ async function resolveAgentContexts(params: {
return agentContexts
}
function projectAgentContextInputs(
message: string,
contexts: UnifiedChatRequest['contexts'],
registry: ResolvedSecretTraceRegistry | undefined
): { message: string; contexts: UnifiedChatRequest['contexts'] } {
const labels = (contexts ?? []).map((context) => context.label ?? null)
const projection = projectResolvedSecretModelContent({ message, labels }, registry)
if (!projection.safe || !isPlainRecord(projection.value)) {
throw new Error('Agent context input could not be safely projected')
}
const projectedMessage = projection.value.message
const projectedLabels = projection.value.labels
if (
typeof projectedMessage !== 'string' ||
!Array.isArray(projectedLabels) ||
projectedLabels.length !== labels.length ||
!projectedLabels.every((label) => label === null || typeof label === 'string')
) {
throw new Error('Agent context input could not be safely projected')
}
return {
message: projectedMessage,
contexts: contexts?.map((context, index) => ({
...context,
...(projectedLabels[index] === null ? {} : { label: projectedLabels[index] }),
})),
}
}
async function persistUserMessage(params: {
chatId?: string
userMessageId: string
@@ -1211,12 +1178,7 @@ export async function handleUnifiedChatPost(req: NextRequest) {
}),
activeOtelRoot.context
)
const agentContextsPromise = executionContextPromise.then((executionContext) => {
const projected = projectAgentContextInputs(
body.message,
normalizedContexts,
executionContext.resolvedSecretTraceRegistry
)
const agentContextsPromise = executionContextPromise.then(() => {
return withCopilotSpan(
TraceSpan.CopilotChatResolveAgentContexts,
{
@@ -1225,10 +1187,10 @@ export async function handleUnifiedChatPost(req: NextRequest) {
},
() =>
resolveAgentContexts({
contexts: projected.contexts,
contexts: normalizedContexts,
resourceAttachments: body.resourceAttachments,
userId: authenticatedUserId,
message: projected.message,
message: body.message,
workspaceId,
chatId: actualChatId,
requestId,
apps/sim/lib/copilot/mcp-tools.test.ts
+22 -113
View File
@@ -50,130 +50,39 @@ describe('mothership MCP tool schemas', () => {
])
})
it('reports tagged-server discovery provenance without placing it in tool schemas', async () => {
const provenance = {
version: 1,
complete: true,
entries: [{ name: 'MCP_TOKEN', encryptedValue: 'encrypted-token' }],
scope: { userId: 'user-1', workspaceId: 'ws-1' },
}
discoverServerTools.mockImplementationOnce(
async (
_userId: string,
_serverId: string,
_workspaceId: string,
_forceRefresh: boolean,
report: (value: unknown) => void
) => {
report(provenance)
return [
{
serverId: 'mcp-server-1',
name: 'search',
description: 'Search docs',
inputSchema: { type: 'object' },
},
]
}
)
const recordProvenance = vi.fn()
const tools = await buildTaggedMcpToolSchemas(
'user-1',
'ws-1',
['mcp-server-1'],
recordProvenance
)
expect(discoverServerTools).toHaveBeenCalledWith(
'user-1',
'mcp-server-1',
'ws-1',
false,
expect.any(Function)
)
expect(recordProvenance).toHaveBeenCalledWith(provenance)
expect(JSON.stringify(tools)).not.toContain('encrypted-token')
expect(JSON.stringify(tools)).not.toContain('resolvedSecretTraceProvenance')
})
it('reports incomplete provenance when tagged-server discovery returns no report', async () => {
discoverServerTools.mockResolvedValue([])
const recordProvenance = vi.fn()
await buildTaggedMcpToolSchemas('user-1', 'ws-1', ['mcp-server-1'], recordProvenance)
expect(recordProvenance).toHaveBeenCalledWith({
version: 1,
complete: false,
entries: [],
scope: { userId: 'user-1', workspaceId: 'ws-1' },
})
})
it('uses a selected block tool cached schema without discovering the server', async () => {
const recordProvenance = vi.fn()
const tools = await buildSelectedMcpToolSchemas(
'user-1',
'ws-1',
[
{
type: 'mcp',
params: { serverId: 'mcp-server-1', toolName: 'search', serverName: 'Docs' },
schema: { type: 'object', properties: { query: { type: 'string' } } },
},
],
recordProvenance
)
const tools = await buildSelectedMcpToolSchemas('user-1', 'ws-1', [
{
type: 'mcp',
params: { serverId: 'mcp-server-1', toolName: 'search', serverName: 'Docs' },
schema: { type: 'object', properties: { query: { type: 'string' } } },
},
])
expect(discoverServerTools).not.toHaveBeenCalled()
expect(recordProvenance).not.toHaveBeenCalled()
expect(tools[0]).toMatchObject({
name: 'mcp-server-1-search',
input_schema: { type: 'object', properties: { query: { type: 'string' } } },
})
})
it('reports provenance from selected tools that require server discovery', async () => {
const provenance = {
version: 1,
complete: true,
entries: [{ name: 'MCP_TOKEN', encryptedValue: 'encrypted-token' }],
scope: { userId: 'user-1', workspaceId: 'ws-1' },
}
discoverServerTools.mockImplementationOnce(
async (
_userId: string,
_serverId: string,
_workspaceId: string,
_forceRefresh: boolean,
report: (value: unknown) => void
) => {
report(provenance)
return [
{
serverId: 'mcp-server-1',
name: 'search',
inputSchema: { type: 'object' },
},
]
}
)
const recordProvenance = vi.fn()
it('discovers a selected legacy tool without a cached schema', async () => {
discoverServerTools.mockResolvedValueOnce([
{
serverId: 'mcp-server-1',
name: 'search',
inputSchema: { type: 'object' },
},
])
const tools = await buildSelectedMcpToolSchemas(
'user-1',
'ws-1',
[
{
type: 'mcp',
params: { serverId: 'mcp-server-1', toolName: 'search' },
},
],
recordProvenance
)
const tools = await buildSelectedMcpToolSchemas('user-1', 'ws-1', [
{
type: 'mcp',
params: { serverId: 'mcp-server-1', toolName: 'search' },
},
])
expect(recordProvenance).toHaveBeenCalledWith(provenance)
expect(discoverServerTools).toHaveBeenCalledWith('user-1', 'mcp-server-1', 'ws-1')
expect(tools[0]).toMatchObject({ name: 'mcp-server-1-search' })
})
})
apps/sim/lib/copilot/mcp-tools.ts
+6 -41
View File
@@ -8,8 +8,6 @@ import type { ToolInput } from '@/executor/handlers/agent/types'
const logger = createLogger('CopilotMcpTools')
type ResolvedSecretTraceProvenanceCallback = (provenance: unknown) => void
function toMothershipMcpTool(tool: {
serverId: string
serverName?: string
@@ -53,26 +51,11 @@ function dedupeMcpTools(tools: ToolSchema[]): ToolSchema[] {
async function discoverServerTools(
userId: string,
workspaceId: string,
serverId: string,
onResolvedSecretTraceProvenance?: ResolvedSecretTraceProvenanceCallback
serverId: string
): Promise<McpTool[]> {
let provenanceReported = false
try {
const { mcpService } = await import('@/lib/mcp/service')
if (!onResolvedSecretTraceProvenance) {
return await mcpService.discoverServerTools(userId, serverId, workspaceId)
}
return await mcpService.discoverServerTools(
userId,
serverId,
workspaceId,
false,
(provenance) => {
provenanceReported = true
onResolvedSecretTraceProvenance(provenance)
}
)
return await mcpService.discoverServerTools(userId, serverId, workspaceId)
} catch (error) {
logger.warn('Failed to resolve tagged MCP server tools', {
serverId,
@@ -80,15 +63,6 @@ async function discoverServerTools(
error: toError(error).message,
})
return []
} finally {
if (onResolvedSecretTraceProvenance && !provenanceReported) {
onResolvedSecretTraceProvenance({
version: 1,
complete: false,
entries: [],
scope: { userId, workspaceId },
})
}
}
}
@@ -99,17 +73,14 @@ async function discoverServerTools(
export async function buildTaggedMcpToolSchemas(
userId: string,
workspaceId: string,
serverIds: string[],
onResolvedSecretTraceProvenance?: ResolvedSecretTraceProvenanceCallback
serverIds: string[]
): Promise<ToolSchema[]> {
const uniqueServerIds = [...new Set(serverIds.filter(Boolean))]
if (uniqueServerIds.length === 0) return []
await validateMcpToolsAllowed(userId, workspaceId)
const discovered = await Promise.all(
uniqueServerIds.map((serverId) =>
discoverServerTools(userId, workspaceId, serverId, onResolvedSecretTraceProvenance)
)
uniqueServerIds.map((serverId) => discoverServerTools(userId, workspaceId, serverId))
)
return dedupeMcpTools(discovered.flat().map(toMothershipMcpTool))
}
@@ -122,8 +93,7 @@ export async function buildTaggedMcpToolSchemas(
export async function buildSelectedMcpToolSchemas(
userId: string,
workspaceId: string,
selections: ToolInput[],
onResolvedSecretTraceProvenance?: ResolvedSecretTraceProvenanceCallback
selections: ToolInput[]
): Promise<ToolSchema[]> {
const selected = selections.filter(
(tool) =>
@@ -160,12 +130,7 @@ export async function buildSelectedMcpToolSchemas(
let discovery = discoveredByServer.get(serverId)
if (!discovery) {
discovery = discoverServerTools(
userId,
workspaceId,
serverId,
onResolvedSecretTraceProvenance
)
discovery = discoverServerTools(userId, workspaceId, serverId)
discoveredByServer.set(serverId, discovery)
}
const match = (await discovery).find((tool) => tool.name === toolName)
apps/sim/lib/copilot/model-visible-schema.test.ts
-149
View File
@@ -1,149 +0,0 @@
/**
* @vitest-environment node
*/
import { describe, expect, it } from 'vitest'
import {
collectModelVisibleSchemaContent,
restoreModelVisibleSchemaValues,
} from '@/lib/copilot/model-visible-schema'
describe('model-visible schema classification', () => {
it('projects display text, preserves public grammar, and guards dynamic semantics', () => {
const schema = {
type: 'object',
properties: {
tokenField: {
type: 'string',
title: 'Visible title',
description: 'Visible description',
enum: ['semantic-value'],
default: 'semantic-default',
},
},
required: ['tokenField'],
}
const content = collectModelVisibleSchemaContent(schema)
expect(content.projectedValues).toEqual(['Visible title', 'Visible description'])
expect(content.guardedValues).toEqual(
expect.arrayContaining(['tokenField', ['semantic-value'], 'semantic-default', ['tokenField']])
)
expect(
restoreModelVisibleSchemaValues(schema, ['Projected title', 'Projected description'])
).toEqual({
type: 'object',
properties: {
tokenField: {
type: 'string',
title: 'Projected title',
description: 'Projected description',
enum: ['semantic-value'],
default: 'semantic-default',
},
},
required: ['tokenField'],
})
})
it('preserves validated controls while guarding arbitrary or invalid semantic values', () => {
const schema = {
type: ['object', 'null'],
nullable: true,
readOnly: false,
format: 'secret-format',
$schema: 'secret-schema-uri',
contentEncoding: 'secret-encoding',
contentMediaType: 'secret-media-type',
properties: {
invalidType: { type: 'secret-type' },
invalidBoolean: { deprecated: 'secret-deprecated' },
},
}
expect(collectModelVisibleSchemaContent(schema).guardedValues).toEqual(
expect.arrayContaining([
'secret-format',
'secret-schema-uri',
'secret-encoding',
'secret-media-type',
'invalidType',
'secret-type',
'invalidBoolean',
'secret-deprecated',
])
)
})
it.each([
['string', { type: 'string' }],
['true', { nullable: true }],
])('preserves the validated public control %s outside secret matching', (_secret, schema) => {
expect(collectModelVisibleSchemaContent(schema).guardedValues).toEqual([])
})
it.each([
{ items: 'not-a-schema' },
{ allOf: ['not-a-schema'] },
{ properties: { field: 'not-a-schema' } },
])('rejects malformed child schemas instead of silently preserving them', (schema) => {
expect(() => collectModelVisibleSchemaContent(schema)).toThrow(
'Model-visible schema content could not be safely projected'
)
expect(() => restoreModelVisibleSchemaValues(schema, [])).toThrow(
'Model-visible schema content could not be safely projected'
)
})
it('accepts boolean schemas at every schema-child position', () => {
const schema = {
additionalProperties: false,
allOf: [true],
properties: { field: false },
}
expect(collectModelVisibleSchemaContent(schema)).toEqual({
projectedValues: [],
guardedValues: ['field'],
})
expect(restoreModelVisibleSchemaValues(schema, [])).toEqual(schema)
})
it('guards arbitrary keys at the root and within child schemas', () => {
const schema = {
'root-semantic-key': true,
properties: {
field: {
'child-semantic-key': 'value',
},
},
}
expect(collectModelVisibleSchemaContent(schema).guardedValues).toEqual(
expect.arrayContaining(['root-semantic-key', 'child-semantic-key'])
)
})
it('restores safe canonical controls byte-for-byte', () => {
const schema = {
type: ['object', 'null'],
nullable: true,
readOnly: false,
properties: { value: { type: 'string' } },
}
expect(restoreModelVisibleSchemaValues(schema, [])).toEqual(schema)
})
it('rejects oversized schema collections before duplicating them', () => {
const oversized = new Array(100_001)
const schema = { allOf: oversized }
expect(() => collectModelVisibleSchemaContent(schema)).toThrow(
'Model-visible schema content could not be safely projected'
)
expect(() => restoreModelVisibleSchemaValues(schema, [])).toThrow(
'Model-visible schema content could not be safely projected'
)
})
})
apps/sim/lib/copilot/model-visible-schema.ts
-361
View File
@@ -1,361 +0,0 @@
import { isPlainRecord } from '@sim/utils/object'
const MAX_SCHEMA_NODES = 100_000
const MAX_SCHEMA_DEPTH = 100
const SCHEMA_DISPLAY_KEYS = new Set(['description', 'title', '$comment', 'example', 'examples'])
const SCHEMA_SINGLE_CHILD_KEYS = new Set([
'additionalProperties',
'contains',
'contentSchema',
'else',
'if',
'items',
'not',
'propertyNames',
'then',
'unevaluatedItems',
'unevaluatedProperties',
])
const SCHEMA_ARRAY_CHILD_KEYS = new Set(['allOf', 'anyOf', 'oneOf', 'prefixItems'])
const SCHEMA_MAP_CHILD_KEYS = new Set([
'$defs',
'definitions',
'dependentSchemas',
'patternProperties',
'properties',
])
const SCHEMA_TYPE_NAMES = new Set([
'array',
'boolean',
'integer',
'null',
'number',
'object',
'string',
])
const SCHEMA_BOOLEAN_CONTROL_KEYS = new Set([
'deprecated',
'nullable',
'readOnly',
'uniqueItems',
'writeOnly',
])
const SCHEMA_NUMBER_CONTROL_KEYS = new Set([
'exclusiveMaximum',
'exclusiveMinimum',
'maximum',
'minimum',
'multipleOf',
])
const SCHEMA_NONNEGATIVE_INTEGER_CONTROL_KEYS = new Set([
'maxContains',
'maxItems',
'maxLength',
'maxProperties',
'minContains',
'minItems',
'minLength',
'minProperties',
])
const SCHEMA_KNOWN_KEYS = new Set([
'$anchor',
'$comment',
'$defs',
'$dynamicAnchor',
'$dynamicRef',
'$id',
'$ref',
'$schema',
'$vocabulary',
'additionalProperties',
'allOf',
'anyOf',
'const',
'contains',
'contentEncoding',
'contentMediaType',
'contentSchema',
'default',
'definitions',
'deprecated',
'dependentRequired',
'dependentSchemas',
'description',
'else',
'enum',
'example',
'examples',
'exclusiveMaximum',
'exclusiveMinimum',
'format',
'if',
'items',
'maxContains',
'maxItems',
'maxLength',
'maxProperties',
'maximum',
'minContains',
'minItems',
'minLength',
'minProperties',
'minimum',
'multipleOf',
'not',
'nullable',
'oneOf',
'pattern',
'patternProperties',
'prefixItems',
'properties',
'propertyNames',
'readOnly',
'required',
'then',
'title',
'type',
'unevaluatedItems',
'unevaluatedProperties',
'uniqueItems',
'writeOnly',
])
export type ModelVisibleSchemaAction =
| 'preserve'
| 'project'
| 'traverse'
| 'verify'
| 'traverse-verify-key'
| 'verify-key-value'
export class ModelVisibleSchemaError extends Error {
constructor() {
super('Model-visible schema content could not be safely projected')
this.name = 'ModelVisibleSchemaError'
}
}
export function getModelVisibleSchemaAction(
parentKey: string | undefined,
key: string,
value?: unknown
): ModelVisibleSchemaAction {
if (parentKey !== undefined && SCHEMA_MAP_CHILD_KEYS.has(parentKey)) {
return isSchemaNode(value) ? 'traverse-verify-key' : 'verify-key-value'
}
if (SCHEMA_DISPLAY_KEYS.has(key)) return 'project'
if (SCHEMA_SINGLE_CHILD_KEYS.has(key)) return isSchemaNode(value) ? 'traverse' : 'verify'
if (SCHEMA_ARRAY_CHILD_KEYS.has(key)) return Array.isArray(value) ? 'traverse' : 'verify'
if (SCHEMA_MAP_CHILD_KEYS.has(key)) return isPlainRecord(value) ? 'traverse' : 'verify'
if (isPublicSchemaControl(key, value)) return 'preserve'
return 'verify'
}
interface SchemaTraversalState {
nodes: number
ancestors: WeakSet<object>
}
function visitSchemaNode(state: SchemaTraversalState, depth: number): void {
state.nodes += 1
if (state.nodes > MAX_SCHEMA_NODES || depth > MAX_SCHEMA_DEPTH) {
throw new ModelVisibleSchemaError()
}
}
function isSchemaNode(value: unknown): value is boolean | Record<string, unknown> {
return typeof value === 'boolean' || isPlainRecord(value)
}
function isPublicSchemaControl(key: string, value: unknown): boolean {
if (key === 'type') {
if (typeof value === 'string') return SCHEMA_TYPE_NAMES.has(value)
return (
Array.isArray(value) &&
value.length > 0 &&
value.every((item) => typeof item === 'string' && SCHEMA_TYPE_NAMES.has(item))
)
}
if (SCHEMA_BOOLEAN_CONTROL_KEYS.has(key)) return typeof value === 'boolean'
if (SCHEMA_NUMBER_CONTROL_KEYS.has(key)) {
return typeof value === 'number' && Number.isFinite(value)
}
if (SCHEMA_NONNEGATIVE_INTEGER_CONTROL_KEYS.has(key)) {
return typeof value === 'number' && Number.isInteger(value) && value >= 0
}
return false
}
function schemaRecordEntries(value: Record<string, unknown>): Array<[string, unknown]> {
const keys = Reflect.ownKeys(value)
if (keys.length > MAX_SCHEMA_NODES) throw new ModelVisibleSchemaError()
const entries: Array<[string, unknown]> = []
for (const key of keys) {
if (typeof key !== 'string') throw new ModelVisibleSchemaError()
const descriptor = Object.getOwnPropertyDescriptor(value, key)
if (!descriptor?.enumerable || !('value' in descriptor)) {
throw new ModelVisibleSchemaError()
}
entries.push([key, descriptor.value])
}
return entries
}
function schemaArrayValues(value: unknown[]): unknown[] {
if (Object.getPrototypeOf(value) !== Array.prototype) throw new ModelVisibleSchemaError()
if (value.length > MAX_SCHEMA_NODES) throw new ModelVisibleSchemaError()
const values = new Array<unknown>(value.length)
let entries = 0
for (const key of Reflect.ownKeys(value)) {
if (key === 'length') continue
if (typeof key !== 'string') throw new ModelVisibleSchemaError()
const index = Number(key)
if (!Number.isInteger(index) || index < 0 || index >= value.length || String(index) !== key) {
throw new ModelVisibleSchemaError()
}
const descriptor = Object.getOwnPropertyDescriptor(value, key)
if (!descriptor?.enumerable || !('value' in descriptor)) {
throw new ModelVisibleSchemaError()
}
values[index] = descriptor.value
entries += 1
}
if (entries !== value.length) throw new ModelVisibleSchemaError()
return values
}
function schemaChildren(key: string, value: unknown): unknown[] {
if (SCHEMA_SINGLE_CHILD_KEYS.has(key)) return [value]
if (SCHEMA_ARRAY_CHILD_KEYS.has(key)) {
if (!Array.isArray(value)) throw new ModelVisibleSchemaError()
return schemaArrayValues(value)
}
if (SCHEMA_MAP_CHILD_KEYS.has(key)) {
if (!isPlainRecord(value)) throw new ModelVisibleSchemaError()
return schemaRecordEntries(value).map(([, child]) => child)
}
return []
}
export interface ModelVisibleSchemaContent {
projectedValues: unknown[]
guardedValues: unknown[]
}
/** Splits schema display text from semantic fields whose exact bytes must remain unchanged. */
export function collectModelVisibleSchemaContent(schema: unknown): ModelVisibleSchemaContent {
const projectedValues: unknown[] = []
const guardedValues: unknown[] = []
const state: SchemaTraversalState = { nodes: 0, ancestors: new WeakSet<object>() }
const visit = (candidate: unknown, depth: number): void => {
visitSchemaNode(state, depth)
if (typeof candidate === 'boolean') return
if (!isPlainRecord(candidate)) throw new ModelVisibleSchemaError()
if (state.ancestors.has(candidate)) throw new ModelVisibleSchemaError()
state.ancestors.add(candidate)
try {
for (const [key, value] of schemaRecordEntries(candidate)) {
const action = getModelVisibleSchemaAction(undefined, key, value)
if (action === 'project') {
projectedValues.push(value)
continue
}
if (action === 'preserve') continue
if (action === 'verify') {
if (
SCHEMA_SINGLE_CHILD_KEYS.has(key) ||
SCHEMA_ARRAY_CHILD_KEYS.has(key) ||
SCHEMA_MAP_CHILD_KEYS.has(key)
) {
throw new ModelVisibleSchemaError()
}
if (!SCHEMA_KNOWN_KEYS.has(key)) guardedValues.push(key)
guardedValues.push(value)
continue
}
if (action !== 'traverse') continue
if (SCHEMA_MAP_CHILD_KEYS.has(key)) {
if (!isPlainRecord(value)) throw new ModelVisibleSchemaError()
for (const [childKey, child] of schemaRecordEntries(value)) {
guardedValues.push(childKey)
visit(child, depth + 1)
}
continue
}
for (const child of schemaChildren(key, value)) visit(child, depth + 1)
}
} finally {
state.ancestors.delete(candidate)
}
}
visit(schema, 0)
return { projectedValues, guardedValues }
}
export function collectModelVisibleSchemaValues(schema: unknown): unknown[] {
return collectModelVisibleSchemaContent(schema).projectedValues
}
export function restoreModelVisibleSchemaValues(schema: unknown, projected: unknown): unknown {
if (!Array.isArray(projected)) throw new ModelVisibleSchemaError()
const state: SchemaTraversalState = { nodes: 0, ancestors: new WeakSet<object>() }
let cursor = 0
const visit = (candidate: unknown, depth: number): unknown => {
visitSchemaNode(state, depth)
if (typeof candidate === 'boolean') return candidate
if (!isPlainRecord(candidate)) throw new ModelVisibleSchemaError()
if (state.ancestors.has(candidate)) throw new ModelVisibleSchemaError()
state.ancestors.add(candidate)
try {
let restored = candidate
for (const [key, value] of schemaRecordEntries(candidate)) {
const action = getModelVisibleSchemaAction(undefined, key, value)
let nextValue = value
if (action === 'project') {
if (cursor >= projected.length) throw new ModelVisibleSchemaError()
nextValue = projected[cursor]
cursor += 1
} else if (
action === 'verify' &&
(SCHEMA_SINGLE_CHILD_KEYS.has(key) ||
SCHEMA_ARRAY_CHILD_KEYS.has(key) ||
SCHEMA_MAP_CHILD_KEYS.has(key))
) {
throw new ModelVisibleSchemaError()
} else if (action === 'traverse' || action === 'traverse-verify-key') {
if (SCHEMA_SINGLE_CHILD_KEYS.has(key)) {
nextValue = visit(value, depth + 1)
} else if (SCHEMA_ARRAY_CHILD_KEYS.has(key)) {
if (!Array.isArray(value)) throw new ModelVisibleSchemaError()
nextValue = schemaArrayValues(value).map((child) => visit(child, depth + 1))
} else if (SCHEMA_MAP_CHILD_KEYS.has(key)) {
if (!isPlainRecord(value)) throw new ModelVisibleSchemaError()
nextValue = Object.fromEntries(
schemaRecordEntries(value).map(([childKey, child]) => [
childKey,
visit(child, depth + 1),
])
)
}
}
if (nextValue !== value) {
if (restored === candidate) restored = { ...candidate }
restored[key] = nextValue
}
}
return restored
} finally {
state.ancestors.delete(candidate)
}
}
const restored = visit(schema, 0)
if (cursor !== projected.length) throw new ModelVisibleSchemaError()
return restored
}
apps/sim/lib/copilot/request/handlers/handlers.test.ts
+57 -4
View File
@@ -473,7 +473,7 @@ describe('sse-handlers tool lifecycle', () => {
expect(updated?.result?.output).toBe('done')
})
it('projects resolved Function secrets before every Copilot-visible result sink', async () => {
it('projects resolved Function output while leaving resource metadata unchanged', async () => {
const registry = new ResolvedSecretTraceRegistry([
{
name: 'SECRET',
@@ -485,7 +485,9 @@ describe('sse-handlers tool lifecycle', () => {
execContext.resolvedSecretTraceRegistry = registry
execContext.chatId = 'chat-1'
executeTool.mockImplementationOnce(async (_name, _params, toolContext) => {
toolContext.resolvedSecretTraceRegistry?.recordResolved('SECRET', 'secret-value')
toolContext.resolvedSecretTraceRegistry?.recordResolved('SECRET', 'secret-value', {
propagated: true,
})
return {
success: true,
output: {
@@ -542,12 +544,11 @@ describe('sse-handlers tool lifecycle', () => {
resource: {
type: 'file',
id: 'file-1',
title: '{{SECRET}}.txt',
title: 'secret-value.txt',
},
},
})
expect(JSON.stringify(completeAsyncToolCall.mock.calls)).not.toContain('secret-value')
expect(JSON.stringify(onEvent.mock.calls)).not.toContain('secret-value')
})
it('emits a structural result for a detached background workflow tool', async () => {
@@ -1661,6 +1662,58 @@ describe('sse-handlers tool lifecycle', () => {
)
})
it('forwards workspace secret references unchanged to the resolved integration operation', async () => {
isSimExecuted.mockReturnValue(false)
executeTool.mockResolvedValueOnce({ success: true, output: { searchResults: [] } })
await sseHandlers.tool(
{
type: MothershipStreamV1EventType.tool,
payload: {
toolCallId: 'gateway-serper',
toolName: 'call_integration_tool',
executor: MothershipStreamV1ToolExecutor.go,
mode: MothershipStreamV1ToolMode.sync,
phase: MothershipStreamV1ToolPhase.call,
status: 'generating',
partial: true,
},
} satisfies StreamEvent,
context,
execContext,
{ interactive: false, timeout: 1000 }
)
await sseHandlers.tool(
{
type: MothershipStreamV1EventType.tool,
payload: {
toolCallId: 'gateway-serper',
toolName: 'serper_search',
arguments: {
query: 'invoice',
apiKey: '{{SERPER_API_KEY}}',
},
executor: MothershipStreamV1ToolExecutor.sim,
mode: MothershipStreamV1ToolMode.async,
phase: MothershipStreamV1ToolPhase.call,
},
} satisfies StreamEvent,
context,
execContext,
{ interactive: false, timeout: 1000 }
)
await sleep(0)
expect(executeTool).toHaveBeenCalledOnce()
expect(executeTool).toHaveBeenCalledWith(
'serper_search',
{ query: 'invoice', apiKey: '{{SERPER_API_KEY}}' },
expect.any(Object)
)
})
it('clears pending continuation state when a run resumes', async () => {
context.awaitingAsyncContinuation = {
checkpointId: 'cp-1',
apps/sim/lib/copilot/request/lifecycle/run.test.ts
+248 -416
View File
@@ -142,7 +142,7 @@ import { runCopilotLifecycle } from '@/lib/copilot/request/lifecycle/run'
afterAll(resetEnvFlagsMock)
const ARBITRARY_SCHEMA_CONTROL_KEYS = [
const SCHEMA_CONTROL_KEYS = [
'$schema',
'format',
'contentEncoding',
@@ -279,76 +279,56 @@ describe('runCopilotLifecycle', () => {
})
})
it('projects secrets in every model-visible initial Go payload field without rewriting foreign aliases', async () => {
it('preserves ordinary initial Go payload fields that collide with a configured secret', async () => {
const secret = 'mothership-secret'
const registry = new ResolvedSecretTraceRegistry([
{ name: 'TOKEN', plaintext: secret, encryptedValue: 'ciphertext' },
])
registry.recordResolved('TOKEN', secret)
const payload = {
message: `message ${secret} __var_FOREIGN`,
messages: [{ role: 'user', content: secret }],
context: [{ type: 'resource', content: secret }],
contexts: [{ type: 'mcp', content: secret }],
workspaceContext: `workspace ${secret}`,
integrationTools: [{ name: 'tool', description: secret }],
mothershipTools: [{ name: 'mcp', description: '__sim_code_2_binding_0' }],
fileAttachments: [
{
name: `${secret}.txt`,
key: 'raw-storage-key',
source: { type: 'base64', data: 'c2FmZQ==' },
},
],
workspaceId: 'ws-1',
messageId: 'stream-model-projection',
}
let capturedRequestBody = ''
mockRunStreamLoop.mockImplementationOnce(async (_url: string, request: RequestInit) => {
capturedRequestBody = String(request.body)
})
await runCopilotLifecycle(
{
message: `message ${secret} __var_FOREIGN`,
messages: [{ role: 'user', content: secret }],
context: [{ type: 'resource', content: secret }],
contexts: [{ type: 'mcp', content: secret }],
workspaceContext: `workspace ${secret}`,
integrationTools: [{ name: 'tool', description: secret }],
mothershipTools: [{ name: 'mcp', description: '__sim_code_2_binding_0' }],
fileAttachments: [
{
name: `${secret}.txt`,
key: 'raw-storage-key',
source: { type: 'base64', data: 'c2FmZQ==' },
},
],
workspaceId: 'ws-1',
messageId: 'stream-model-projection',
},
{
userId: 'user-1',
workspaceId: 'ws-1',
executionContext: {
userId: 'user-1',
workflowId: '',
workspaceId: 'ws-1',
},
resolvedSecretTraceRegistry: registry,
}
)
expect(capturedRequestBody).not.toContain(secret)
expect(capturedRequestBody).toContain('__var_FOREIGN')
expect(capturedRequestBody).toContain('__sim_code_2_binding_0')
expect(JSON.parse(capturedRequestBody)).toMatchObject({
message: 'message {{TOKEN}} __var_FOREIGN',
messages: [{ role: 'user', content: '{{TOKEN}}' }],
workspaceContext: 'workspace {{TOKEN}}',
mothershipTools: [{ name: 'mcp', description: '__sim_code_2_binding_0' }],
await runCopilotLifecycle(payload, {
userId: 'user-1',
workspaceId: 'ws-1',
fileAttachments: [
{
name: '{{TOKEN}}.txt',
key: 'raw-storage-key',
source: { type: 'base64', data: 'c2FmZQ==' },
},
],
executionContext: {
userId: 'user-1',
workflowId: '',
workspaceId: 'ws-1',
},
resolvedSecretTraceRegistry: registry,
})
const { enterpriseByokEligible, ...sent } = JSON.parse(capturedRequestBody)
expect(enterpriseByokEligible).toBe(false)
expect(sent).toEqual(payload)
})
it('projects large tool catalogs at the tool-definition boundary', async () => {
it('preserves large ordinary tool catalogs without scanning configured secret values', async () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'TOKEN', plaintext: 'catalog-secret', encryptedValue: 'ciphertext' },
])
registry.recordResolved('TOKEN', 'catalog-secret')
const toolCount = 4_000
const propertiesPerTool = 8
// These definitions are individually small, but flattening their semantic fields into one
// synthetic projection value creates 108,001 traversal nodes and crosses the per-value budget.
const integrationTools = Array.from({ length: toolCount }, (_, toolIndex) => {
const properties = Object.fromEntries(
Array.from({ length: propertiesPerTool }, (_, propertyIndex) => [
@@ -392,17 +372,13 @@ describe('runCopilotLifecycle', () => {
expect(result.success).toBe(true)
expect(mockRunStreamLoop).toHaveBeenCalledOnce()
const sent = JSON.parse(capturedRequestBody)
expect(sent.integrationTools).toHaveLength(integrationTools.length)
expect(sent.integrationTools[0].description).toBe('Tool 0 uses {{TOKEN}}')
expect(sent.integrationTools.at(-1).name).toBe(`tool_${toolCount - 1}`)
expect(capturedRequestBody).not.toContain('catalog-secret')
expect(sent.integrationTools).toEqual(integrationTools)
})
it('projects selected JSON and attachment fields exactly once when plaintext overlaps its alias', async () => {
it('preserves ordinary JSON and attachment fields when plaintext overlaps its alias', async () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'TOKEN', plaintext: 'TOKEN', encryptedValue: 'ciphertext' },
])
registry.recordResolved('TOKEN', 'TOKEN')
let capturedRequestBody = ''
mockRunStreamLoop.mockImplementationOnce(async (_url: string, request: RequestInit) => {
capturedRequestBody = String(request.body)
@@ -449,25 +425,23 @@ describe('runCopilotLifecycle', () => {
)
const sent = JSON.parse(capturedRequestBody)
expect(sent.message).toBe('{{TOKEN}}')
expect(sent.message).toBe('TOKEN')
expect(sent.messages[0]).toMatchObject({
content: '{{TOKEN}}',
function_call: { arguments: JSON.stringify({ value: '{{TOKEN}}' }) },
content: 'TOKEN',
function_call: { arguments: JSON.stringify({ value: 'TOKEN' }) },
tool_calls: [
{
function: { arguments: JSON.stringify({ value: '{{TOKEN}}' }) },
function: { arguments: JSON.stringify({ value: 'TOKEN' }) },
},
],
files: [
{
name: '{{TOKEN}}.txt',
context: 'Context {{TOKEN}}',
name: 'TOKEN.txt',
context: 'Context TOKEN',
},
],
})
expect(sent.fileAttachments).toEqual([{ name: '{{TOKEN}}.txt', key: 'safe-key' }])
expect(capturedRequestBody.replaceAll('{{TOKEN}}', '')).not.toContain('TOKEN')
expect(capturedRequestBody).not.toContain('{{{{TOKEN}}}}')
expect(sent.fileAttachments).toEqual([{ name: 'TOKEN.txt', key: 'safe-key' }])
})
it('omits only unsafe durable attachments before the initial Go request', async () => {
@@ -521,375 +495,242 @@ describe('runCopilotLifecycle', () => {
})
it.each(['123', 'true'])(
'keeps low-entropy Copilot JSON valid while separating content from controls (%s)',
'preserves low-entropy configured-secret collisions across Copilot JSON (%s)',
async (secret) => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'TOKEN', plaintext: secret, encryptedValue: 'ciphertext' },
])
registry.recordResolved('TOKEN', secret)
const converted = secret === '123' ? 123 : true
let capturedRequestBody = ''
mockRunStreamLoop.mockImplementationOnce(async (_url: string, request: RequestInit) => {
capturedRequestBody = String(request.body)
})
await runCopilotLifecycle(
{
message: `Message ${secret}`,
messages: [
{
id: secret,
role: secret,
name: 'assistant-safe',
content: `Transcript ${secret}`,
function_call: {
name: 'legacy-safe',
arguments: JSON.stringify({ value: secret, converted }),
const payload = {
message: `Message ${secret}`,
messages: [
{
id: secret,
role: secret,
name: 'assistant-safe',
content: `Transcript ${secret}`,
function_call: {
name: 'legacy-safe',
arguments: JSON.stringify({ value: secret, converted }),
},
tool_calls: [
{
id: secret,
type: secret,
function: {
name: 'tool-safe',
arguments: JSON.stringify({ value: secret, converted }),
},
},
tool_calls: [
{
],
fileAttachments: [
{
id: secret,
key: secret,
filename: `${secret}.txt`,
media_type: secret,
},
],
contexts: [{ kind: secret, label: `Label ${secret}`, serverId: secret }],
contentBlocks: [
{
type: secret,
content: `Block ${secret}`,
toolCall: {
id: secret,
type: secret,
function: {
name: 'tool-safe',
arguments: JSON.stringify({ value: secret, converted }),
},
name: 'nested-tool-safe',
state: secret,
params: { value: secret },
result: { success: true, output: { value: secret, converted } },
display: { title: `Title ${secret}` },
},
],
fileAttachments: [
{
id: secret,
key: secret,
filename: `${secret}.txt`,
media_type: secret,
},
],
contexts: [{ kind: secret, label: `Label ${secret}`, serverId: secret }],
contentBlocks: [
{
type: secret,
content: `Block ${secret}`,
toolCall: {
id: secret,
name: 'nested-tool-safe',
state: secret,
params: { value: secret },
result: { success: true, output: { value: secret, converted } },
display: { title: `Title ${secret}` },
},
},
],
},
],
context: [
{ type: secret, tag: secret, path: secret, content: `Unsafe ${secret}` },
{
type: secret,
tag: secret,
path: 'files/safe.txt',
content: `Context ${secret}`,
},
],
contexts: [
{
kind: secret,
serverId: secret,
label: `Context label ${secret}`,
},
],
integrationTools: [
{
name: 'safe_tool',
description: `Description ${secret}`,
input_schema: {
type: 'object',
properties: {
value: {
type: 'string',
title: `Title ${secret}`,
description: `Field ${secret}`,
enum: ['public'],
},
},
required: ['value'],
},
params: { runtimeControl: secret },
service: secret,
operation: secret,
oauth: { required: true, provider: secret },
},
{
name: 'unsafe_schema_tool',
description: 'Unsafe schema',
input_schema: {
type: 'object',
properties: { [secret]: { type: 'string' } },
required: [secret],
},
},
{
name: secret,
description: 'Unsafe name',
input_schema: { type: 'object', properties: {}, required: [] },
},
],
responseFormat: {
name: 'safe_response',
schema: {
],
},
],
context: [
{ type: secret, tag: secret, path: secret, content: `Unsafe ${secret}` },
{
type: secret,
tag: secret,
path: 'files/safe.txt',
content: `Context ${secret}`,
},
],
contexts: [
{
kind: secret,
serverId: secret,
label: `Context label ${secret}`,
},
],
integrationTools: [
{
name: 'safe_tool',
description: `Description ${secret}`,
input_schema: {
type: 'object',
properties: {
value: {
type: 'string',
description: `Result ${secret}`,
title: `Title ${secret}`,
description: `Field ${secret}`,
enum: ['public'],
},
},
required: ['value'],
},
params: { runtimeControl: secret },
service: secret,
operation: secret,
oauth: { required: true, provider: secret },
},
fileAttachments: [
{
id: secret,
name: `${secret}.txt`,
key: secret,
mimeType: secret,
},
],
vfs: {
workspace: { id: secret, ownerId: secret, name: `Workspace ${secret}` },
files: [
{
id: secret,
path: secret,
folderPath: secret,
type: secret,
name: `File ${secret}`,
},
{
id: 'safe-file-id',
path: 'files/safe.txt',
folderPath: 'files',
type: 'text/plain',
name: `Safe ${secret}`,
},
],
mcpServers: [
{ id: secret, name: `Unsafe ${secret}`, url: `https://${secret}.example` },
{
id: 'safe-mcp-id',
name: `Safe MCP ${secret}`,
url: 'https://mcp.example',
},
],
},
userTimezone: secret,
userMetadata: {
name: `User ${secret}`,
email: `owner+${secret}@example.com`,
timezone: secret,
},
desktopCapabilities: {
terminal: true,
terminals: [
{
id: secret,
cwd: `/workspace/${secret}`,
running: `command ${secret}`,
active: true,
},
{
id: 'safe-terminal-id',
cwd: '/workspace/safe',
running: `safe command ${secret}`,
active: true,
},
],
browser: true,
browserSessions: [
{ hostname: secret, evidence: 'cookies', lastObservedAt: '2026-01-01T00:00:00.000Z' },
{
hostname: 'safe.example',
evidence: 'sign-in-completed',
lastObservedAt: '2026-02-01T00:00:00.000Z',
},
],
},
workspaceId: 'ws-1',
messageId: `stream-low-entropy-${secret}`,
},
{
userId: 'user-1',
workspaceId: 'ws-1',
executionContext: {
userId: 'user-1',
workflowId: '',
workspaceId: 'ws-1',
},
resolvedSecretTraceRegistry: registry,
}
)
const sent = JSON.parse(capturedRequestBody)
expect(sent.messages[0]).toMatchObject({
id: secret,
role: secret,
name: 'assistant-safe',
content: 'Transcript {{TOKEN}}',
function_call: {
name: 'legacy-safe',
},
tool_calls: [
{
id: secret,
type: secret,
function: {
name: 'tool-safe',
name: 'unsafe_schema_tool',
description: 'Unsafe schema',
input_schema: {
type: 'object',
properties: { [secret]: { type: 'string' } },
required: [secret],
},
},
{
name: secret,
description: 'Unsafe name',
input_schema: { type: 'object', properties: {}, required: [] },
},
],
responseFormat: {
name: 'safe_response',
schema: {
type: 'object',
properties: {
value: {
type: 'string',
description: `Result ${secret}`,
enum: ['public'],
},
},
required: ['value'],
},
},
fileAttachments: [
{
id: secret,
name: `${secret}.txt`,
key: secret,
filename: '{{TOKEN}}.txt',
media_type: secret,
mimeType: secret,
},
],
contexts: [{ kind: secret, label: 'Label {{TOKEN}}', serverId: secret }],
contentBlocks: [
{
type: secret,
content: 'Block {{TOKEN}}',
toolCall: {
vfs: {
workspace: { id: secret, ownerId: secret, name: `Workspace ${secret}` },
files: [
{
id: secret,
name: 'nested-tool-safe',
state: secret,
params: { value: '{{TOKEN}}' },
result: {
success: true,
output: { value: '{{TOKEN}}', converted: '{{TOKEN}}' },
},
display: { title: 'Title {{TOKEN}}' },
path: secret,
folderPath: secret,
type: secret,
name: `File ${secret}`,
},
},
],
})
expect(JSON.parse(sent.messages[0].function_call.arguments)).toEqual({
value: '{{TOKEN}}',
converted: '{{TOKEN}}',
})
expect(JSON.parse(sent.messages[0].tool_calls[0].function.arguments)).toEqual({
value: '{{TOKEN}}',
converted: '{{TOKEN}}',
})
expect(sent.context).toEqual([
{
type: secret,
tag: '{{TOKEN}}',
path: 'files/safe.txt',
content: 'Context {{TOKEN}}',
},
])
expect(sent.contexts).toEqual([
{
kind: secret,
serverId: secret,
label: 'Context label {{TOKEN}}',
},
])
expect(sent.integrationTools).toHaveLength(1)
expect(sent.integrationTools[0]).toMatchObject({
name: 'safe_tool',
description: 'Description {{TOKEN}}',
input_schema: {
properties: {
value: {
title: 'Title {{TOKEN}}',
description: 'Field {{TOKEN}}',
enum: ['public'],
{
id: 'safe-file-id',
path: 'files/safe.txt',
folderPath: 'files',
type: 'text/plain',
name: `Safe ${secret}`,
},
},
required: ['value'],
],
mcpServers: [
{ id: secret, name: `Unsafe ${secret}`, url: `https://${secret}.example` },
{
id: 'safe-mcp-id',
name: `Safe MCP ${secret}`,
url: 'https://mcp.example',
},
],
},
params: { runtimeControl: secret },
service: secret,
operation: secret,
oauth: { required: true, provider: secret },
})
expect(sent.responseFormat).toMatchObject({
name: 'safe_response',
schema: {
properties: {
value: { description: 'Result {{TOKEN}}', enum: ['public'] },
},
required: ['value'],
userTimezone: secret,
userMetadata: {
name: `User ${secret}`,
email: `owner+${secret}@example.com`,
timezone: secret,
},
desktopCapabilities: {
terminal: true,
terminals: [
{
id: secret,
cwd: `/workspace/${secret}`,
running: `command ${secret}`,
active: true,
},
{
id: 'safe-terminal-id',
cwd: '/workspace/safe',
running: `safe command ${secret}`,
active: true,
},
],
browser: true,
browserSessions: [
{ hostname: secret, evidence: 'cookies', lastObservedAt: '2026-01-01T00:00:00.000Z' },
{
hostname: 'safe.example',
evidence: 'sign-in-completed',
lastObservedAt: '2026-02-01T00:00:00.000Z',
},
],
},
workspaceId: 'ws-1',
messageId: `stream-low-entropy-${secret}`,
}
await runCopilotLifecycle(payload, {
userId: 'user-1',
workspaceId: 'ws-1',
executionContext: {
userId: 'user-1',
workflowId: '',
workspaceId: 'ws-1',
},
resolvedSecretTraceRegistry: registry,
})
expect(sent.fileAttachments[0]).toEqual({
id: secret,
name: '{{TOKEN}}.txt',
key: secret,
mimeType: secret,
})
expect(sent.vfs).toEqual({
workspace: { id: secret, ownerId: secret, name: 'Workspace {{TOKEN}}' },
files: [
{
id: 'safe-file-id',
path: 'files/safe.txt',
folderPath: 'files',
type: 'text/plain',
name: 'Safe {{TOKEN}}',
},
],
mcpServers: [
{
id: 'safe-mcp-id',
name: 'Safe MCP {{TOKEN}}',
url: 'https://mcp.example',
},
],
})
expect(sent).not.toHaveProperty('userTimezone')
expect(sent.userMetadata).toEqual({
name: 'User {{TOKEN}}',
email: 'owner+{{TOKEN}}@example.com',
})
expect(sent.desktopCapabilities).toEqual({
terminal: true,
terminals: [
{
id: 'safe-terminal-id',
cwd: '/workspace/safe',
running: 'safe command {{TOKEN}}',
active: true,
},
],
browser: true,
browserSessions: [
{
hostname: 'safe.example',
evidence: 'sign-in-completed',
lastObservedAt: '2026-02-01T00:00:00.000Z',
},
],
})
const { enterpriseByokEligible, ...sent } = JSON.parse(capturedRequestBody)
expect(enterpriseByokEligible).toBe(false)
expect(sent).toEqual(payload)
}
)
it.each(ARBITRARY_SCHEMA_CONTROL_KEYS)(
'guards arbitrary %s schema controls before initial Copilot model egress',
it.each(SCHEMA_CONTROL_KEYS)(
'preserves ordinary %s schema controls without scanning configured secret values',
async (controlKey) => {
const secret = `copilot-schema-control-secret-${controlKey}`
const registry = new ResolvedSecretTraceRegistry([
{ name: 'TOKEN', plaintext: secret, encryptedValue: 'ciphertext' },
])
registry.recordResolved('TOKEN', secret)
const unsafeSchema = {
const schema = {
type: 'object',
properties: {},
[controlKey]: secret,
}
const integrationTools = [
{
name: 'schema_tool',
description: 'Schema with an ordinary configured-secret collision',
input_schema: schema,
},
{
name: 'safe_tool',
description: 'Safe schema',
input_schema: { type: 'object', properties: {} },
},
]
let capturedRequestBody = ''
mockRunStreamLoop.mockImplementationOnce(async (_url: string, request: RequestInit) => {
capturedRequestBody = String(request.body)
@@ -899,18 +740,7 @@ describe('runCopilotLifecycle', () => {
{
message: 'Use a safe tool',
messageId: `stream-schema-tool-${controlKey}`,
integrationTools: [
{
name: 'unsafe_tool',
description: 'Unsafe schema control',
input_schema: unsafeSchema,
},
{
name: 'safe_tool',
description: 'Safe schema',
input_schema: { type: 'object', properties: {} },
},
],
integrationTools,
},
{
userId: 'user-1',
@@ -920,9 +750,7 @@ describe('runCopilotLifecycle', () => {
}
)
expect(JSON.parse(capturedRequestBody).integrationTools).toEqual([
expect.objectContaining({ name: 'safe_tool' }),
])
expect(JSON.parse(capturedRequestBody).integrationTools).toEqual(integrationTools)
mockRunStreamLoop.mockClear()
mockRunStreamLoop.mockImplementationOnce(async (_url: string, request: RequestInit) => {
@@ -932,7 +760,7 @@ describe('runCopilotLifecycle', () => {
{
message: 'Use a response schema',
messageId: `stream-schema-response-${controlKey}`,
responseFormat: { name: 'unsafe_response', schema: unsafeSchema },
responseFormat: { name: 'ordinary_response', schema },
},
{
userId: 'user-1',
@@ -943,13 +771,15 @@ describe('runCopilotLifecycle', () => {
)
expect(result.success).toBe(true)
expect(JSON.parse(capturedRequestBody)).not.toHaveProperty('responseFormat')
expect(capturedRequestBody).not.toContain(secret)
expect(JSON.parse(capturedRequestBody).responseFormat).toEqual({
name: 'ordinary_response',
schema,
})
expect(mockRunStreamLoop).toHaveBeenCalledOnce()
}
)
it('omits malformed and oversized optional Copilot response schemas', async () => {
it('does not couple optional Copilot response schemas to secret provenance', async () => {
const registry = new ResolvedSecretTraceRegistry()
for (const [index, schema] of [
{ properties: { field: 'not-a-schema' } },
@@ -976,7 +806,9 @@ describe('runCopilotLifecycle', () => {
)
expect(result.success).toBe(true)
expect(JSON.parse(capturedRequestBody)).not.toHaveProperty('responseFormat')
expect(JSON.stringify(JSON.parse(capturedRequestBody).responseFormat)).toBe(
JSON.stringify({ name: 'unsafe_response', schema })
)
expect(mockRunStreamLoop).toHaveBeenCalledOnce()
}
})
@@ -1082,7 +914,7 @@ describe('runCopilotLifecycle', () => {
expect(JSON.parse(capturedRequestBody).responseFormat.schema).toEqual(schema)
})
it('fails before the initial Go request when model projection is incomplete', async () => {
it('does not block ordinary initial Go payloads on unrelated incomplete provenance', async () => {
const registry = new ResolvedSecretTraceRegistry()
registry.markIncomplete()
@@ -1100,11 +932,12 @@ describe('runCopilotLifecycle', () => {
}
)
expect(result).toMatchObject({
success: false,
error: 'Copilot model input could not be safely projected',
expect(result.success).toBe(true)
expect(JSON.parse(String(mockRunStreamLoop.mock.calls[0]?.[1].body))).toMatchObject({
message: 'possibly secret',
messageId: 'stream-incomplete-projection',
})
expect(mockRunStreamLoop).not.toHaveBeenCalled()
expect(mockRunStreamLoop).toHaveBeenCalledOnce()
})
describe('tool permission feature flag', () => {
@@ -1526,11 +1359,10 @@ describe('runCopilotLifecycle', () => {
}
})
it('fails closed instead of sending a secret-bearing tool name on resume', async () => {
it('preserves a resume tool name that collides with a configured secret', async () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'TOKEN', plaintext: 'unsafe-tool', encryptedValue: 'ciphertext' },
])
registry.recordResolved('TOKEN', 'unsafe-tool')
mockRunStreamLoop.mockImplementationOnce(
async (
_fetchUrl: string,
@@ -1559,11 +1391,11 @@ describe('runCopilotLifecycle', () => {
}
)
expect(result).toMatchObject({
success: false,
error: 'Copilot model input could not be safely projected',
expect(result.success).toBe(true)
expect(mockRunStreamLoop).toHaveBeenCalledTimes(2)
expect(JSON.parse(String(mockRunStreamLoop.mock.calls[1]?.[1].body))).toMatchObject({
results: [{ callId: 'tool-1', name: 'unsafe-tool', success: true }],
})
expect(mockRunStreamLoop).toHaveBeenCalledTimes(1)
})
it('runs legacy-v0 during Sim-first deployment without guessed billing aliases', async () => {
apps/sim/lib/copilot/request/lifecycle/run.ts
+7 -626
View File
@@ -4,7 +4,7 @@ import type { PermissionType } from '@sim/platform-authz/workspace'
import { toError } from '@sim/utils/errors'
import { sleep } from '@sim/utils/helpers'
import { generateId } from '@sim/utils/id'
import { isPlainRecord, omit } from '@sim/utils/object'
import { omit } from '@sim/utils/object'
import {
type AttributedBillingRequestEnvelope,
assertBillingAttributionSnapshot,
@@ -28,20 +28,6 @@ import {
MothershipStreamV1RunKind,
MothershipStreamV1ToolOutcome,
} from '@/lib/copilot/generated/mothership-stream-v1'
import {
COPILOT_CONTEXT_MODEL_TEXT_KEYS,
COPILOT_CONTEXT_ROUTING_KEYS,
COPILOT_DESKTOP_MODEL_TEXT_KEYS,
COPILOT_MESSAGE_DISPLAY_KEYS,
COPILOT_USER_METADATA_MODEL_TEXT_KEYS,
COPILOT_VFS_MODEL_TEXT_KEYS,
COPILOT_VFS_ROUTING_KEYS,
isCopilotModelTextKey,
} from '@/lib/copilot/model-visible-content'
import {
collectModelVisibleSchemaContent,
getModelVisibleSchemaAction,
} from '@/lib/copilot/model-visible-schema'
import { getAutoAllowedTools } from '@/lib/copilot/persistence/tool-permission/auto-allow'
import { createStreamingContext } from '@/lib/copilot/request/context/request-context'
import { buildToolCallSummaries } from '@/lib/copilot/request/context/result'
@@ -83,63 +69,17 @@ import {
isHosted,
} from '@/lib/core/config/env-flags'
import { filterModelSafeWorkspaceFileAttachments } from '@/lib/uploads/contexts/workspace/workspace-file-secret-provenance'
import {
isResolvedSecretModelContentUnchanged,
projectResolvedSecretModelContent,
projectResolvedSecretModelJsonStrings,
} from '@/executor/utils/resolved-secret-content-projection'
import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
const logger = createLogger('CopilotLifecycle')
const MAX_RESUME_ATTEMPTS = 3
const RESUME_BACKOFF_MS = [250, 500, 1000] as const
const MAX_SELECTED_CONTENT_NODES = 100_000
const MAX_SELECTED_CONTENT_DEPTH = 100
const SIMPLE_MODEL_CONTENT_KEYS = [
'message',
'systemPrompt',
'workspaceContext',
'commands',
'implicitFeedback',
'workflowName',
] as const
const TOOL_PAYLOAD_KEYS = ['tools', 'integrationTools', 'mothershipTools'] as const
const TOOL_SCHEMA_KEYS = new Set(['input_schema', 'parameters', 'outputs'])
const MOTHERSHIP_CODE_TOOL_ROUTES = new Set([
'/api/copilot',
'/api/mothership',
'/api/mothership/execute',
])
const MESSAGE_CONTAINER_KEYS = new Set([
'contentBlocks',
'contexts',
'display',
'fileAttachments',
'files',
'function',
'function_call',
'result',
'toolCall',
'tool_calls',
])
const MESSAGE_OPAQUE_CONTENT_KEYS = new Set(['error', 'output', 'params'])
const ATTACHMENT_PARENT_KEYS = new Set(['attachments', 'fileAttachments', 'files'])
const MESSAGE_HANDLE_PARENT_KEYS = new Set(['function', 'function_call', 'toolCall'])
type SelectedContentAction =
| 'preserve'
| 'project'
| 'project-json'
| 'traverse'
| 'traverse-verify-key'
| 'verify-key-value'
| 'verify'
type SelectedContentSelector = (
path: readonly string[],
key: string,
value: unknown
) => SelectedContentAction
class CopilotModelContentProjectionError extends Error {
constructor() {
@@ -148,451 +88,6 @@ class CopilotModelContentProjectionError extends Error {
}
}
interface SelectedContentTraversalState {
nodes: number
ancestors: WeakSet<object>
}
interface SelectedContentBuckets {
projected: unknown[]
jsonStrings: string[]
guarded: unknown[]
}
function visitSelectedContentNode(state: SelectedContentTraversalState, depth: number): void {
state.nodes += 1
if (state.nodes > MAX_SELECTED_CONTENT_NODES || depth > MAX_SELECTED_CONTENT_DEPTH) {
throw new CopilotModelContentProjectionError()
}
}
function projectModelContent(value: unknown, registry: ResolvedSecretTraceRegistry): unknown {
const projection = projectResolvedSecretModelContent(value, registry)
if (!projection.safe) throw new CopilotModelContentProjectionError()
return projection.value
}
function collectSelectedContent(
value: unknown,
selector: SelectedContentSelector,
selected: SelectedContentBuckets,
path: readonly string[] = [],
state: SelectedContentTraversalState = { nodes: 0, ancestors: new WeakSet<object>() },
depth = 0
): void {
visitSelectedContentNode(state, depth)
if (value === null || typeof value !== 'object') return
if (state.ancestors.has(value)) throw new CopilotModelContentProjectionError()
state.ancestors.add(value)
try {
if (Array.isArray(value)) {
for (const item of value) {
collectSelectedContent(item, selector, selected, [...path, '*'], state, depth + 1)
}
return
}
if (!isPlainRecord(value)) throw new CopilotModelContentProjectionError()
for (const [key, item] of Object.entries(value)) {
const action = selector(path, key, item)
if (action === 'project') {
selected.projected.push(item)
} else if (action === 'project-json') {
if (typeof item !== 'string') throw new CopilotModelContentProjectionError()
selected.jsonStrings.push(item)
} else if (action === 'verify') {
selected.guarded.push(item)
} else if (action === 'verify-key-value') {
selected.guarded.push(key, item)
} else if (action === 'traverse' || action === 'traverse-verify-key') {
if (action === 'traverse-verify-key') selected.guarded.push(key)
collectSelectedContent(item, selector, selected, [...path, key], state, depth + 1)
}
}
} finally {
state.ancestors.delete(value)
}
}
function restoreSelectedContent(
value: unknown,
selector: SelectedContentSelector,
projected: readonly unknown[],
projectedJsonStrings: readonly string[],
cursor: { projected: number; jsonStrings: number },
path: readonly string[] = [],
state: SelectedContentTraversalState = { nodes: 0, ancestors: new WeakSet<object>() },
depth = 0
): unknown {
visitSelectedContentNode(state, depth)
if (value === null || typeof value !== 'object') return value
if (state.ancestors.has(value)) throw new CopilotModelContentProjectionError()
state.ancestors.add(value)
try {
if (Array.isArray(value)) {
return value.map((item) =>
restoreSelectedContent(
item,
selector,
projected,
projectedJsonStrings,
cursor,
[...path, '*'],
state,
depth + 1
)
)
}
if (!isPlainRecord(value)) throw new CopilotModelContentProjectionError()
const restored: Record<string, unknown> = { ...value }
for (const [key, item] of Object.entries(value)) {
const action = selector(path, key, item)
if (action === 'project') {
if (cursor.projected >= projected.length) throw new CopilotModelContentProjectionError()
restored[key] = projected[cursor.projected]
cursor.projected += 1
} else if (action === 'project-json') {
if (cursor.jsonStrings >= projectedJsonStrings.length) {
throw new CopilotModelContentProjectionError()
}
restored[key] = projectedJsonStrings[cursor.jsonStrings]
cursor.jsonStrings += 1
} else if (action === 'traverse' || action === 'traverse-verify-key') {
restored[key] = restoreSelectedContent(
item,
selector,
projected,
projectedJsonStrings,
cursor,
[...path, key],
state,
depth + 1
)
}
}
return restored
} finally {
state.ancestors.delete(value)
}
}
function projectSelectedContent(
value: unknown,
registry: ResolvedSecretTraceRegistry,
selector: SelectedContentSelector
): unknown {
const selected: SelectedContentBuckets = { projected: [], jsonStrings: [], guarded: [] }
collectSelectedContent(value, selector, selected)
if (!isResolvedSecretModelContentUnchanged(selected.guarded, registry)) {
throw new CopilotModelContentProjectionError()
}
const projected = projectModelContent(selected.projected, registry)
if (!Array.isArray(projected) || projected.length !== selected.projected.length) {
throw new CopilotModelContentProjectionError()
}
const jsonProjection = projectResolvedSecretModelJsonStrings(selected.jsonStrings, registry)
if (
!jsonProjection.safe ||
!Array.isArray(jsonProjection.value) ||
!jsonProjection.value.every((item) => typeof item === 'string') ||
jsonProjection.value.length !== selected.jsonStrings.length
) {
throw new CopilotModelContentProjectionError()
}
const cursor = { projected: 0, jsonStrings: 0 }
const restored = restoreSelectedContent(value, selector, projected, jsonProjection.value, cursor)
if (cursor.projected !== projected.length || cursor.jsonStrings !== jsonProjection.value.length) {
throw new CopilotModelContentProjectionError()
}
return restored
}
function projectStructuredContent(
value: unknown,
registry: ResolvedSecretTraceRegistry,
selector: SelectedContentSelector,
shape: 'array' | 'record' | 'record-or-array'
): unknown {
if (
(shape === 'array' && !Array.isArray(value)) ||
(shape === 'record' && !isPlainRecord(value)) ||
(shape === 'record-or-array' && !Array.isArray(value) && !isPlainRecord(value))
) {
throw new CopilotModelContentProjectionError()
}
return projectSelectedContent(value, registry, selector)
}
function schemaContentAction(
path: readonly string[],
key: string,
value: unknown
): SelectedContentAction {
return getModelVisibleSchemaAction(path.at(-1), key, value)
}
const toolContentSelector: SelectedContentSelector = (path, key, value) => {
if (path.length === 0 && key === 'description') return 'project'
if (path.length === 0 && key === 'name') return 'verify'
if (path.length === 0 && TOOL_SCHEMA_KEYS.has(key)) return 'traverse'
const schemaRootIndex = path.findIndex((segment) => TOOL_SCHEMA_KEYS.has(segment))
if (schemaRootIndex >= 0) {
return schemaContentAction(path.slice(schemaRootIndex + 1), key, value)
}
return 'preserve'
}
const contextContentSelector: SelectedContentSelector = (_path, key, value) => {
if (isCopilotModelTextKey(COPILOT_CONTEXT_MODEL_TEXT_KEYS, key)) return 'project'
return value !== null && typeof value === 'object' ? 'traverse' : 'preserve'
}
function nearestPathContainer(path: readonly string[]): string | undefined {
for (let index = path.length - 1; index >= 0; index -= 1) {
if (path[index] !== '*') return path[index]
}
return undefined
}
const messageContentSelector: SelectedContentSelector = (path, key, value) => {
if (key === 'content') {
return value !== null && typeof value === 'object' ? 'traverse' : 'project'
}
if (MESSAGE_OPAQUE_CONTENT_KEYS.has(key)) return 'project'
if (key === 'arguments') return 'project-json'
if (isCopilotModelTextKey(COPILOT_MESSAGE_DISPLAY_KEYS, key)) return 'project'
if (
(key === 'name' || key === 'filename' || key === 'fileName') &&
ATTACHMENT_PARENT_KEYS.has(nearestPathContainer(path) ?? '')
) {
return 'project'
}
if (
key === 'name' &&
(path.length === 1 || MESSAGE_HANDLE_PARENT_KEYS.has(nearestPathContainer(path) ?? ''))
) {
return 'verify'
}
if (key === 'name') return 'project'
if (key === 'context' && nearestPathContainer(path) === 'files') return 'project'
if (MESSAGE_CONTAINER_KEYS.has(key)) return 'traverse'
return 'preserve'
}
const responseFormatContentSelector: SelectedContentSelector = (path, key, value) => {
if (path.length === 0 && (key === 'description' || key === 'instructions')) return 'project'
if (path.length === 0 && key === 'name') return 'verify'
if (path.length === 0 && key === 'schema') return 'traverse'
if (path.length === 0) return schemaContentAction(path, key, value)
if (path[0] === 'schema') return schemaContentAction(path.slice(1), key, value)
return 'preserve'
}
const vfsContentSelector: SelectedContentSelector = (_path, key, value) => {
if (isCopilotModelTextKey(COPILOT_VFS_MODEL_TEXT_KEYS, key)) return 'project'
return value !== null && typeof value === 'object' ? 'traverse' : 'preserve'
}
const userMetadataContentSelector: SelectedContentSelector = (_path, key) =>
isCopilotModelTextKey(COPILOT_USER_METADATA_MODEL_TEXT_KEYS, key) ? 'project' : 'preserve'
const desktopContentSelector: SelectedContentSelector = (_path, key, value) => {
if (isCopilotModelTextKey(COPILOT_DESKTOP_MODEL_TEXT_KEYS, key)) return 'project'
return value !== null && typeof value === 'object' ? 'traverse' : 'preserve'
}
function projectModelSafeToolPayloads(
value: unknown,
registry: ResolvedSecretTraceRegistry
): unknown[] {
if (!Array.isArray(value)) throw new CopilotModelContentProjectionError()
const projected: unknown[] = []
for (const candidate of value) {
if (!isPlainRecord(candidate) || typeof candidate.name !== 'string') continue
try {
for (const schemaKey of TOOL_SCHEMA_KEYS) {
if (Object.hasOwn(candidate, schemaKey)) {
collectModelVisibleSchemaContent(candidate[schemaKey])
}
}
projected.push(projectStructuredContent(candidate, registry, toolContentSelector, 'record'))
} catch {
// Tool definitions are independent protocol entities. Reject an unsafe definition without
// turning the entire catalog into one synthetic projection value or failing safe siblings.
}
}
// Projection completeness is a request-level invariant, even when every candidate was rejected.
projectModelContent([], registry)
return projected
}
function hasModelSafeRoutingFields(
value: Record<string, unknown>,
routingKeys: readonly string[],
registry: ResolvedSecretTraceRegistry
): boolean {
for (const routingKey of routingKeys) {
if (!Object.hasOwn(value, routingKey)) continue
const routingValue = value[routingKey]
if (
typeof routingValue !== 'string' ||
!isResolvedSecretModelContentUnchanged(routingValue, registry)
) {
return false
}
}
return true
}
function filterModelSafeContextPayload(
value: unknown,
registry: ResolvedSecretTraceRegistry
): Record<string, unknown> | unknown[] | undefined {
if (Array.isArray(value)) {
return value.filter(
(candidate) =>
isPlainRecord(candidate) &&
hasModelSafeRoutingFields(candidate, COPILOT_CONTEXT_ROUTING_KEYS, registry)
)
}
if (!isPlainRecord(value)) throw new CopilotModelContentProjectionError()
return hasModelSafeRoutingFields(value, COPILOT_CONTEXT_ROUTING_KEYS, registry)
? value
: undefined
}
function filterModelSafeVfsPayload(
value: unknown,
registry: ResolvedSecretTraceRegistry
): Record<string, unknown> {
if (!isPlainRecord(value)) throw new CopilotModelContentProjectionError()
const filtered: Record<string, unknown> = { ...value }
for (const [collectionKey, collection] of Object.entries(value)) {
if (collectionKey === 'envVars') {
if (!Array.isArray(collection) || !collection.every((item) => typeof item === 'string')) {
throw new CopilotModelContentProjectionError()
}
filtered[collectionKey] = collection.filter((name) =>
isResolvedSecretModelContentUnchanged(name, registry)
)
continue
}
if (!Array.isArray(collection)) continue
filtered[collectionKey] = collection.filter((candidate) => {
if (!isPlainRecord(candidate)) return false
return hasModelSafeRoutingFields(candidate, COPILOT_VFS_ROUTING_KEYS, registry)
})
}
return filtered
}
function filterModelSafeUserMetadata(
value: unknown,
registry: ResolvedSecretTraceRegistry
): Record<string, unknown> {
if (!isPlainRecord(value)) throw new CopilotModelContentProjectionError()
const filtered = { ...value }
if (
Object.hasOwn(filtered, 'timezone') &&
(typeof filtered.timezone !== 'string' ||
!isResolvedSecretModelContentUnchanged(filtered.timezone, registry))
) {
return omit(filtered, ['timezone'])
}
return filtered
}
function filterModelSafeDesktopCapabilities(
value: unknown,
registry: ResolvedSecretTraceRegistry
): Record<string, unknown> {
if (!isPlainRecord(value)) throw new CopilotModelContentProjectionError()
const filtered: Record<string, unknown> = { ...value }
if (Object.hasOwn(value, 'terminals')) {
if (!Array.isArray(value.terminals)) throw new CopilotModelContentProjectionError()
filtered.terminals = value.terminals.filter((terminal) => {
if (!isPlainRecord(terminal)) return false
return (
!Object.hasOwn(terminal, 'cwd') ||
(typeof terminal.cwd === 'string' &&
isResolvedSecretModelContentUnchanged(terminal.cwd, registry))
)
})
}
if (Object.hasOwn(value, 'browserSessions')) {
if (!Array.isArray(value.browserSessions)) throw new CopilotModelContentProjectionError()
filtered.browserSessions = value.browserSessions.filter(
(session) =>
isPlainRecord(session) &&
typeof session.hostname === 'string' &&
isResolvedSecretModelContentUnchanged(session.hostname, registry)
)
}
return filtered
}
function projectAttachmentDisplayNames(
payload: Record<string, unknown>,
registry: ResolvedSecretTraceRegistry
): Partial<Record<'attachments' | 'fileAttachments', unknown>> {
const projected: Partial<Record<'attachments' | 'fileAttachments', unknown>> = {}
for (const key of ['attachments', 'fileAttachments'] as const) {
if (!Object.hasOwn(payload, key)) continue
const attachments = payload[key]
if (!Array.isArray(attachments)) throw new CopilotModelContentProjectionError()
const displayNames = attachments.map((attachment) => {
if (!isPlainRecord(attachment)) throw new CopilotModelContentProjectionError()
return {
...(Object.hasOwn(attachment, 'name') ? { name: attachment.name } : {}),
...(Object.hasOwn(attachment, 'filename') ? { filename: attachment.filename } : {}),
}
})
const projection = projectResolvedSecretModelContent(displayNames, registry)
if (
!projection.safe ||
!Array.isArray(projection.value) ||
projection.value.length !== attachments.length
) {
throw new CopilotModelContentProjectionError()
}
const projectedDisplayNames = projection.value
projected[key] = attachments.map((attachment, index) => {
const displayName = projectedDisplayNames[index]
if (!isPlainRecord(attachment) || !isPlainRecord(displayName)) {
throw new CopilotModelContentProjectionError()
}
const name = displayName.name
const filename = displayName.filename
if (name !== undefined && typeof name !== 'string') {
throw new CopilotModelContentProjectionError()
}
if (filename !== undefined && typeof filename !== 'string') {
throw new CopilotModelContentProjectionError()
}
return {
...attachment,
...(name !== undefined ? { name } : {}),
...(filename !== undefined ? { filename } : {}),
}
})
}
return projected
}
async function omitUnsafeInitialCopilotAttachments(
payload: Record<string, unknown>,
workspaceId?: string
@@ -625,115 +120,11 @@ async function omitUnsafeInitialCopilotAttachments(
return projected
}
async function projectInitialCopilotPayload(
async function filterInitialCopilotAttachmentsForModel(
payload: Record<string, unknown>,
registry: ResolvedSecretTraceRegistry,
workspaceId?: string
): Promise<Record<string, unknown>> {
projectModelContent([], registry)
let projectedPayload = { ...payload }
const simpleContent: Record<string, unknown> = {}
for (const key of SIMPLE_MODEL_CONTENT_KEYS) {
if (Object.hasOwn(payload, key)) simpleContent[key] = payload[key]
}
const projectedSimpleContent = projectModelContent(simpleContent, registry)
if (!isPlainRecord(projectedSimpleContent)) throw new CopilotModelContentProjectionError()
for (const key of SIMPLE_MODEL_CONTENT_KEYS) {
if (Object.hasOwn(payload, key) && Object.hasOwn(projectedSimpleContent, key)) {
projectedPayload[key] = projectedSimpleContent[key]
}
}
if (Object.hasOwn(payload, 'userTimezone')) {
if (
typeof payload.userTimezone === 'string' &&
isResolvedSecretModelContentUnchanged(payload.userTimezone, registry)
) {
projectedPayload.userTimezone = payload.userTimezone
} else {
projectedPayload = omit(projectedPayload, ['userTimezone'])
}
}
if (Object.hasOwn(payload, 'messages')) {
projectedPayload.messages = projectStructuredContent(
payload.messages,
registry,
messageContentSelector,
'array'
)
}
for (const key of ['context', 'contexts'] as const) {
if (Object.hasOwn(payload, key)) {
if (typeof payload[key] === 'string') {
projectedPayload[key] = projectModelContent(payload[key], registry)
continue
}
const safeContexts = filterModelSafeContextPayload(payload[key], registry)
if (safeContexts === undefined) {
projectedPayload = omit(projectedPayload, [key])
} else {
projectedPayload[key] = projectStructuredContent(
safeContexts,
registry,
contextContentSelector,
'record-or-array'
)
}
}
}
for (const key of TOOL_PAYLOAD_KEYS) {
if (Object.hasOwn(payload, key)) {
projectedPayload[key] = projectModelSafeToolPayloads(payload[key], registry)
}
}
if (Object.hasOwn(payload, 'responseFormat')) {
try {
if (
isPlainRecord(payload.responseFormat) &&
Object.hasOwn(payload.responseFormat, 'schema')
) {
collectModelVisibleSchemaContent(payload.responseFormat.schema)
}
projectedPayload.responseFormat =
typeof payload.responseFormat === 'string'
? projectModelContent(payload.responseFormat, registry)
: projectStructuredContent(
payload.responseFormat,
registry,
responseFormatContentSelector,
'record'
)
} catch {
logger.warn('Omitting a Copilot response format with unsafe model-input provenance')
projectedPayload = omit(projectedPayload, ['responseFormat'])
}
}
if (Object.hasOwn(payload, 'vfs')) {
projectedPayload.vfs = projectStructuredContent(
filterModelSafeVfsPayload(payload.vfs, registry),
registry,
vfsContentSelector,
'record'
)
}
if (Object.hasOwn(payload, 'userMetadata')) {
projectedPayload.userMetadata = projectStructuredContent(
filterModelSafeUserMetadata(payload.userMetadata, registry),
registry,
userMetadataContentSelector,
'record'
)
}
if (Object.hasOwn(payload, 'desktopCapabilities')) {
projectedPayload.desktopCapabilities = projectStructuredContent(
filterModelSafeDesktopCapabilities(payload.desktopCapabilities, registry),
registry,
desktopContentSelector,
'record'
)
}
Object.assign(projectedPayload, projectAttachmentDisplayNames(payload, registry))
return omitUnsafeInitialCopilotAttachments(projectedPayload, workspaceId)
return omitUnsafeInitialCopilotAttachments(payload, workspaceId)
}
async function ensureModelEgressRegistry(
@@ -919,10 +310,9 @@ export async function runCopilotLifecycle(
let onCompleteStarted = false
try {
const modelEgressRegistry = await ensureModelEgressRegistry(execContext, lifecycleOptions)
const modelSafeRequestPayload = await projectInitialCopilotPayload(
await ensureModelEgressRegistry(execContext, lifecycleOptions)
const modelSafeRequestPayload = await filterInitialCopilotAttachmentsForModel(
requestPayload,
modelEgressRegistry,
lifecycleOptions.workspaceId
)
await runCheckpointLoop(
@@ -1135,8 +525,7 @@ async function waitForToolIds(context: StreamingContext, toolIds: string[]): Pro
function collectResultsForToolIds(
context: StreamingContext,
toolIds: string[],
checkpointId: string,
registry: ResolvedSecretTraceRegistry
checkpointId: string
): Array<{ callId: string; name: string; data: unknown; success: boolean }> {
return toolIds.map((toolCallId) => {
const tool = context.toolCalls.get(toolCallId)
@@ -1146,9 +535,6 @@ function collectResultsForToolIds(
)
}
const name = tool.name || ''
if (!isResolvedSecretModelContentUnchanged(name, registry)) {
throw new CopilotModelContentProjectionError()
}
return {
callId: toolCallId,
name,
@@ -1238,9 +624,7 @@ async function driveOneChildChain(
if (isAborted(options, context)) return null
await waitForToolIds(context, toolIds)
const registry = execContext.resolvedSecretTraceRegistry
if (!registry) throw new CopilotModelContentProjectionError()
const results = collectResultsForToolIds(context, toolIds, checkpointId, registry)
const results = collectResultsForToolIds(context, toolIds, checkpointId)
const leg = makeResumeLegContext(context)
await runResumeLegWithRetry(
@@ -1651,9 +1035,6 @@ async function runCheckpointLoop(
throw new Error(`Cannot resume: missing result for pending tool call ${toolCallId}`)
}
const name = tool.name || ''
if (!isResolvedSecretModelContentUnchanged(name, execContext.resolvedSecretTraceRegistry)) {
throw new CopilotModelContentProjectionError()
}
results.push({
callId: toolCallId,
name,
apps/sim/lib/copilot/request/lifecycle/start.test.ts
+2 -2
View File
@@ -325,7 +325,7 @@ describe('createSSEStream terminal error handling', () => {
expect(lifecycleTraceparent).toMatch(/^00-[0-9a-f]{32}-[0-9a-f]{16}-0[0-9a-f]$/)
})
it('projects title input using the execution-context secret registry', async () => {
it('does not scan manually authored title input against unrelated active secrets', async () => {
runCopilotLifecycle.mockResolvedValue({
success: true,
content: 'OK',
@@ -362,7 +362,7 @@ describe('createSSEStream terminal error handling', () => {
await vi.waitFor(() => expect(fetchGo).toHaveBeenCalled())
const [, request] = fetchGo.mock.calls.at(-1) ?? []
expect(JSON.parse(request.body)).toEqual(
expect.objectContaining({ message: 'hello {{TOKEN}}' })
expect.objectContaining({ message: 'hello secret-value' })
)
})
})
apps/sim/lib/copilot/request/lifecycle/start.ts
+1 -14
View File
@@ -53,8 +53,6 @@ import { TraceCollector } from '@/lib/copilot/request/trace'
import { getMothershipBaseURL, getMothershipSourceEnvHeaders } from '@/lib/copilot/server/agent-url'
import { env } from '@/lib/core/config/env'
import { isCopilotBillingAttributionV1Enabled, isHosted } from '@/lib/core/config/env-flags'
import { projectResolvedSecretModelContent } from '@/executor/utils/resolved-secret-content-projection'
import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
export { SSE_RESPONSE_HEADERS }
@@ -251,9 +249,6 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS
requestId,
publisher,
otelContext,
resolvedSecretTraceRegistry:
orchestrateOptions.resolvedSecretTraceRegistry ??
orchestrateOptions.executionContext?.resolvedSecretTraceRegistry,
})
try {
@@ -442,7 +437,6 @@ function fireTitleGeneration(params: {
requestId: string
publisher: StreamWriter
otelContext?: Context
resolvedSecretTraceRegistry?: ResolvedSecretTraceRegistry
}): void {
const {
chatId,
@@ -457,18 +451,11 @@ function fireTitleGeneration(params: {
requestId,
publisher,
otelContext,
resolvedSecretTraceRegistry,
} = params
if (!chatId || currentChat?.title || !isNewChat) return
const projectedMessage = projectResolvedSecretModelContent(message, resolvedSecretTraceRegistry)
if (!projectedMessage.safe || typeof projectedMessage.value !== 'string') {
logger.warn(`[${requestId}] Skipping title generation because its input was not safe`)
return
}
requestChatTitle({
message: projectedMessage.value,
message,
model: titleModel,
provider: titleProvider,
userId,
apps/sim/lib/copilot/request/tools/client.ts
+2 -2
View File
@@ -84,7 +84,7 @@ export async function waitForClientToolCompletion({
const completion = await waitForToolCompletion(toolCallId, timeoutMs, abortSignal)
if (!completion) return null
const toolRegistry = registry?.forkForToolInput(undefined)
const toolRegistry = registry?.forkForInputPaths([])
const genericMessage = getGenericCompletionMessage(completion.status)
const binding = runId ? { toolCallId, runId, userId } : undefined
const registryCanImport = toolRegistry !== undefined && !toolRegistry.isPermanentlyIncomplete()
@@ -232,7 +232,7 @@ export async function waitForWorkflowToolCompletion({
abortSignal,
registry,
}: WaitForWorkflowToolCompletionOptions): Promise<AsyncTerminalCompletionSnapshot | null> {
const toolRegistry = registry?.forkForToolInput(undefined)
const toolRegistry = registry?.forkForInputPaths([])
const finishPendingActivation = toolRegistry?.beginPendingActivation()
let completion: AsyncTerminalCompletionSnapshot | null = null
let trustedExecution: Awaited<ReturnType<typeof getTrustedWorkflowToolExecution>> = null
apps/sim/lib/copilot/request/tools/executor.test.ts
+3 -1
View File
@@ -202,7 +202,9 @@ describe('executeToolAndReport provenance isolation', () => {
_params: Record<string, unknown>,
toolContext: ExecutionContext
) => {
toolContext.resolvedSecretTraceRegistry?.recordResolved('TOKEN', 'secret-value')
toolContext.resolvedSecretTraceRegistry?.recordResolved('TOKEN', 'secret-value', {
propagated: true,
})
return { success: true, output: { value: 'secret-value' } }
}
)
apps/sim/lib/copilot/request/tools/executor.ts
+1 -3
View File
@@ -277,9 +277,7 @@ export function buildToolExecutionContext(
return {
...execContext,
toolCallId: toolCall.id,
resolvedSecretTraceRegistry: execContext.resolvedSecretTraceRegistry?.forkForToolInput(
toolCall.params
),
resolvedSecretTraceRegistry: execContext.resolvedSecretTraceRegistry?.forkForInputPaths([]),
...(toolCall.parentToolCallId ? { parentToolCallId: toolCall.parentToolCallId } : {}),
}
}
apps/sim/lib/copilot/request/tools/files.test.ts
+388 -33
View File
@@ -3,10 +3,15 @@
*/
import { beforeEach, describe, expect, it, vi } from 'vitest'
const { mockWriteWorkspaceFileByPath } = vi.hoisted(() => ({
const { mockEncryptSecret, mockWriteWorkspaceFileByPath } = vi.hoisted(() => ({
mockEncryptSecret: vi.fn(),
mockWriteWorkspaceFileByPath: vi.fn(),
}))
vi.mock('@/lib/core/security/encryption', () => ({
encryptSecret: mockEncryptSecret,
}))
vi.mock('@/lib/copilot/vfs/resource-writer', () => ({
writeWorkspaceFileByPath: mockWriteWorkspaceFileByPath,
}))
@@ -27,8 +32,8 @@ import {
serializeOutputForFile,
unwrapFunctionExecuteOutput,
} from '@/lib/copilot/request/tools/files'
import { projectToolResultForCopilot } from '@/lib/copilot/request/tools/resolved-secret-result'
import type { ExecutionContext } from '@/lib/copilot/request/types'
import { MAX_INLINE_MATERIALIZATION_BYTES } from '@/lib/execution/payloads/limits'
import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
describe('unwrapFunctionExecuteOutput', () => {
@@ -121,6 +126,7 @@ describe('maybeWriteOutputToFile', () => {
beforeEach(() => {
vi.clearAllMocks()
mockEncryptSecret.mockResolvedValue({ encrypted: 'encrypted-csv-representation', iv: 'iv' })
mockWriteWorkspaceFileByPath.mockResolvedValue({
id: 'file-1',
name: 'report.csv',
@@ -164,23 +170,104 @@ describe('maybeWriteOutputToFile', () => {
expect(result.success).toBe(true)
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledTimes(1)
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledWith(
expect.objectContaining({ secretProvenance: { status: 'exact', entries: [] } })
)
})
it('persists canonical aliases and leaves unrelated low-entropy public values unchanged', async () => {
const parentRegistry = new ResolvedSecretTraceRegistry([
{
name: 'OUTPUT_SECRET',
plaintext: 'secret-value',
encryptedValue: 'encrypted-output-secret',
},
{
name: 'UNRELATED',
plaintext: 'true',
encryptedValue: 'encrypted-unrelated',
},
it('classifies large structured output from its serialized bytes instead of its object count', async () => {
const registry = new ResolvedSecretTraceRegistry(
[
{ name: 'TOKEN', plaintext: 'secret-value', encryptedValue: 'encrypted-token' },
{
name: 'TOKEN_ALIAS',
plaintext: 'secret-value',
encryptedValue: 'encrypted-token-alias',
},
],
{ userId: 'user-1', workspaceId: 'workspace-1' }
)
registry.recordResolved('TOKEN', 'secret-value')
registry.recordResolved('TOKEN_ALIAS', 'secret-value')
const rows = Array.from({ length: 10_000 }, (_, index) => ({
id: index,
name: `row-${index}`,
status: 'ready',
enabled: true,
token: index === 9_999 ? 'secret-value' : 'public-value',
}))
const result = await maybeWriteOutputToFile(
FunctionExecute.id,
{ outputs: { files: [{ path: 'files/report.json', mode: 'overwrite' }] } },
{ success: true, output: { result: rows, stdout: '' } },
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledWith(
expect.objectContaining({
secretProvenance: {
status: 'exact',
entries: [
{
name: 'TOKEN',
encryptedValue: 'encrypted-token',
sourceUserId: 'user-1',
sourceWorkspaceId: 'workspace-1',
},
{
name: 'TOKEN_ALIAS',
encryptedValue: 'encrypted-token-alias',
sourceUserId: 'user-1',
sourceWorkspaceId: 'workspace-1',
},
],
},
})
)
})
it('writes raw output with unknown provenance when serialized output exceeds the scan budget', async () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'TOKEN', plaintext: 'secret-value', encryptedValue: 'encrypted-token' },
])
registry.recordResolved('TOKEN', 'secret-value')
const result = await maybeWriteOutputToFile(
FunctionExecute.id,
{ outputs: { files: [{ path: 'files/report.txt', mode: 'overwrite' }] } },
{
success: true,
output: { result: 'x'.repeat(MAX_INLINE_MATERIALIZATION_BYTES + 1), stdout: '' },
},
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledWith(
expect.objectContaining({ secretProvenance: { status: 'unknown' } })
)
})
it('persists raw bytes with exact provenance and leaves sibling literals unclassified', async () => {
const parentRegistry = new ResolvedSecretTraceRegistry(
[
{
name: 'OUTPUT_SECRET',
plaintext: 'secret-value',
encryptedValue: 'encrypted-output-secret',
},
{
name: 'UNRELATED',
plaintext: 'true',
encryptedValue: 'encrypted-unrelated',
},
],
{ userId: 'user-1', workspaceId: 'workspace-1' }
)
parentRegistry.recordResolved('UNRELATED', 'true')
const toolRegistry = parentRegistry.forkForToolInput({ code: 'return {{OUTPUT_SECRET}}' })
const toolRegistry = parentRegistry.forkForInputPaths([])
toolRegistry.recordResolved('OUTPUT_SECRET', 'secret-value')
const runtimeOutput = {
result: { token: 'secret-value', publicLabel: 'true', enabled: true },
@@ -195,26 +282,295 @@ describe('maybeWriteOutputToFile', () => {
)
expect(result.success).toBe(true)
const persisted = mockWriteWorkspaceFileByPath.mock.calls[0][0].buffer.toString('utf8')
const write = mockWriteWorkspaceFileByPath.mock.calls[0][0]
const persisted = write.buffer.toString('utf8')
expect(JSON.parse(persisted)).toEqual({
token: '{{OUTPUT_SECRET}}',
token: 'secret-value',
publicLabel: 'true',
enabled: true,
})
expect(write.secretProvenance).toEqual({
status: 'exact',
entries: [
{
name: 'OUTPUT_SECRET',
encryptedValue: 'encrypted-output-secret',
sourceUserId: 'user-1',
sourceWorkspaceId: 'workspace-1',
},
],
})
expect(runtimeOutput.result.token).toBe('secret-value')
const laterRead = projectToolResultForCopilot(
{ success: true, output: { content: persisted } },
new ResolvedSecretTraceRegistry()
)
expect(JSON.parse((laterRead.output as { content: string }).content)).toEqual({
token: '{{OUTPUT_SECRET}}',
publicLabel: 'true',
enabled: true,
})
})
it('does not write when exact persistence provenance is unavailable', async () => {
it('tracks both logical and quote-escaped CSV representations', async () => {
const secret = 'a"b\\c\nline'
const registry = new ResolvedSecretTraceRegistry(
[{ name: 'CSV_SECRET', plaintext: secret, encryptedValue: 'encrypted-csv-secret' }],
{ userId: 'user-1', workspaceId: 'workspace-1' }
)
registry.recordResolved('CSV_SECRET', secret)
const result = await maybeWriteOutputToFile(
FunctionExecute.id,
{ outputs: { files: [{ path: 'files/report.csv', mode: 'overwrite' }] } },
{ success: true, output: { result: [{ value: secret }], stdout: '' } },
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
const write = mockWriteWorkspaceFileByPath.mock.calls[0][0]
expect(write.buffer.toString('utf8')).toBe('value\n"a""b\\c\nline"')
expect(write.secretProvenance).toEqual({
status: 'exact',
entries: [
{
name: 'CSV_SECRET',
encryptedValue: 'encrypted-csv-representation',
sourceUserId: 'user-1',
sourceWorkspaceId: 'workspace-1',
},
{
name: 'CSV_SECRET',
encryptedValue: 'encrypted-csv-secret',
sourceUserId: 'user-1',
sourceWorkspaceId: 'workspace-1',
},
],
})
expect(mockEncryptSecret).toHaveBeenCalledWith('a""b\\c\nline')
})
it('deduplicates repeated CSV quote transformations across the table', async () => {
const secret = 'secret"value'
const registry = new ResolvedSecretTraceRegistry(
[{ name: 'CSV_SECRET', plaintext: secret, encryptedValue: 'encrypted-csv-secret' }],
{ userId: 'user-1', workspaceId: 'workspace-1' }
)
registry.recordResolved('CSV_SECRET', secret)
const rows = Array.from({ length: 10_000 }, () => ({ first: secret, second: secret }))
const result = await maybeWriteOutputToFile(
FunctionExecute.id,
{ outputs: { files: [{ path: 'files/report.csv', mode: 'overwrite' }] } },
{ success: true, output: { result: rows, stdout: '' } },
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
expect(mockEncryptSecret).toHaveBeenCalledTimes(1)
expect(mockEncryptSecret).toHaveBeenCalledWith('secret""value')
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledTimes(1)
})
it('reuses serialized provenance for output files with the same format', async () => {
const secret = 'a"b'
const registry = new ResolvedSecretTraceRegistry(
[{ name: 'CSV_SECRET', plaintext: secret, encryptedValue: 'encrypted-csv-secret' }],
{ userId: 'user-1', workspaceId: 'workspace-1' }
)
registry.recordResolved('CSV_SECRET', secret)
const result = await maybeWriteOutputToFile(
FunctionExecute.id,
{
outputs: {
files: [
{ path: 'files/one.csv', mode: 'overwrite' },
{ path: 'files/two.csv', mode: 'overwrite' },
],
},
},
{ success: true, output: { result: [{ value: secret }], stdout: '' } },
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
expect(mockEncryptSecret).toHaveBeenCalledTimes(1)
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledTimes(2)
})
it('preserves existing multi-file outputs beyond twenty declarations', async () => {
const files = Array.from({ length: 21 }, (_, index) => ({
path: `files/report-${index}.txt`,
mode: 'overwrite' as const,
}))
const result = await maybeWriteOutputToFile(
FunctionExecute.id,
{ outputs: { files } },
{ success: true, output: { result: 'content', stdout: '' } },
buildContext()
)
expect(result.success).toBe(true)
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledTimes(21)
})
it('tracks a persisted legacy runtime alias', async () => {
const registry = new ResolvedSecretTraceRegistry(
[{ name: 'API_KEY', plaintext: 'secret-value', encryptedValue: 'encrypted-api-key' }],
{ userId: 'user-1', workspaceId: 'workspace-1' }
)
registry.recordResolved('API_KEY', 'secret-value')
const result = await maybeWriteOutputToFile(
FunctionExecute.id,
{ outputs: { files: [{ path: 'files/report.txt', mode: 'overwrite' }] } },
{ success: true, output: { result: '__var_API_KEY', stdout: '' } },
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledWith(
expect.objectContaining({
buffer: Buffer.from('__var_API_KEY'),
secretProvenance: {
status: 'exact',
entries: [
{
name: 'API_KEY',
encryptedValue: 'encrypted-api-key',
sourceUserId: 'user-1',
sourceWorkspaceId: 'workspace-1',
},
],
},
})
)
})
it('preserves anonymous provenance without inventing a secret name', async () => {
const registry = {
exportCommittedProvenanceForValue: vi.fn().mockReturnValue({
version: 1,
complete: true,
entries: [{ encryptedValue: 'encrypted-anonymous' }],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
}),
} as unknown as ResolvedSecretTraceRegistry
const result = await maybeWriteOutputToFile(
FunctionExecute.id,
{ outputs: { files: [{ path: 'files/report.txt', mode: 'overwrite' }] } },
{ success: true, output: { result: 'anonymous-secret', stdout: '' } },
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledWith(
expect.objectContaining({
secretProvenance: {
status: 'exact',
entries: [
{
encryptedValue: 'encrypted-anonymous',
sourceUserId: 'user-1',
sourceWorkspaceId: 'workspace-1',
},
],
},
})
)
})
it('preserves the secret source when a different actor writes within the same workspace', async () => {
const registry = new ResolvedSecretTraceRegistry(
[
{
name: 'OUTPUT_SECRET',
plaintext: 'secret-value',
encryptedValue: 'encrypted-output-secret',
},
],
{ userId: 'workflow-owner', workspaceId: 'workspace-1' }
)
registry.recordResolved('OUTPUT_SECRET', 'secret-value')
const result = await maybeWriteOutputToFile(
FunctionExecute.id,
{ outputs: { files: [{ path: 'files/report.txt', mode: 'overwrite' }] } },
{ success: true, output: { result: 'secret-value', stdout: '' } },
buildContext({ userId: 'billing-actor', resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledWith(
expect.objectContaining({
userId: 'billing-actor',
secretProvenance: {
status: 'exact',
entries: [
{
name: 'OUTPUT_SECRET',
encryptedValue: 'encrypted-output-secret',
sourceUserId: 'workflow-owner',
sourceWorkspaceId: 'workspace-1',
},
],
},
})
)
})
it('writes raw output with unknown provenance when the source scope differs', async () => {
const registry = new ResolvedSecretTraceRegistry(
[
{
name: 'OUTPUT_SECRET',
plaintext: 'secret-value',
encryptedValue: 'encrypted-output-secret',
},
],
{ userId: 'workflow-owner', workspaceId: 'workspace-2' }
)
registry.recordResolved('OUTPUT_SECRET', 'secret-value')
const result = await maybeWriteOutputToFile(
FunctionExecute.id,
{ outputs: { files: [{ path: 'files/report.txt', mode: 'overwrite' }] } },
{ success: true, output: { result: 'secret-value', stdout: '' } },
buildContext({ userId: 'billing-actor', resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledWith(
expect.objectContaining({ secretProvenance: { status: 'unknown' } })
)
})
it('prepares every file provenance and marks unavailable lineage unknown', async () => {
const exportCommittedProvenanceForValue = vi
.fn()
.mockReturnValueOnce({ version: 1, complete: true, entries: [] })
.mockReturnValueOnce({ version: 1, complete: false, entries: [] })
const registry = {
exportCommittedProvenanceForValue,
} as unknown as ResolvedSecretTraceRegistry
const result = await maybeWriteOutputToFile(
FunctionExecute.id,
{
outputs: {
files: [
{ path: 'files/report.json', mode: 'overwrite' },
{ path: 'files/report.txt', mode: 'overwrite' },
],
},
},
{ success: true, output: { result: { token: 'value' }, stdout: '' } },
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
expect(exportCommittedProvenanceForValue).toHaveBeenCalledTimes(2)
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledTimes(2)
expect(mockWriteWorkspaceFileByPath.mock.calls[1]?.[0]).toEqual(
expect.objectContaining({ secretProvenance: { status: 'unknown' } })
)
})
it('preserves legacy writes without a registry and marks their provenance unknown', async () => {
const result = await maybeWriteOutputToFile(
FunctionExecute.id,
{ outputs: { files: [{ path: 'files/report.json', mode: 'overwrite' }] } },
@@ -222,11 +578,10 @@ describe('maybeWriteOutputToFile', () => {
buildContext({ resolvedSecretTraceRegistry: undefined })
)
expect(result).toEqual({
success: false,
error: 'Tool output could not be persisted safely because secret provenance was unavailable.',
})
expect(mockWriteWorkspaceFileByPath).not.toHaveBeenCalled()
expect(result.success).toBe(true)
expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledWith(
expect.objectContaining({ secretProvenance: { status: 'unknown' } })
)
})
it('fails loudly instead of silently skipping declared outputs when workspace context is missing', async () => {
apps/sim/lib/copilot/request/tools/files.ts
+191 -37
View File
@@ -7,15 +7,24 @@ import { TraceEvent } from '@/lib/copilot/generated/trace-events-v1'
import { TraceSpan } from '@/lib/copilot/generated/trace-spans-v1'
import { withCopilotSpan } from '@/lib/copilot/request/otel'
import { denyOutputWriteWithoutWritePermission } from '@/lib/copilot/request/tools/permissions'
import {
projectToolErrorMessageForCopilot,
projectToolOutputForPersistence,
} from '@/lib/copilot/request/tools/resolved-secret-result'
import { projectToolErrorMessageForCopilot } from '@/lib/copilot/request/tools/resolved-secret-result'
import type { ExecutionContext, ToolCallResult } from '@/lib/copilot/request/types'
import { decodeVfsPathSegments } from '@/lib/copilot/vfs/path-utils'
import { writeWorkspaceFileByPath } from '@/lib/copilot/vfs/resource-writer'
import {
createWorkspaceFileSecretProvenanceFromRegistry,
type WorkspaceFileSecretProvenance,
type WorkspaceFileSecretProvenanceRepresentation,
} from '@/lib/uploads/contexts/workspace/workspace-file-secret-provenance'
import type { ResolvedSecretMatcher } from '@/executor/utils/resolved-secret-matcher'
import {
createResolvedSecretMatcher,
scanResolvedSecretString,
} from '@/executor/utils/resolved-secret-matcher'
import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
const logger = createLogger('CopilotToolResultFiles')
const MAX_OUTPUT_FILE_PROVENANCE_REPRESENTATIONS = 10_000
export const OUTPUT_PATH_TOOLS: Set<string> = new Set([FunctionExecute.id, UserTable.id])
@@ -99,21 +108,7 @@ export function escapeCsvValue(value: unknown): string {
}
export function convertRowsToCsv(rows: Record<string, unknown>[]): string {
if (rows.length === 0) return ''
const headerSet = new Set<string>()
for (const row of rows) {
for (const key of Object.keys(row)) {
headerSet.add(key)
}
}
const headers = [...headerSet]
const lines = [headers.map(escapeCsvValue).join(',')]
for (const row of rows) {
lines.push(headers.map((h) => escapeCsvValue(row[h])).join(','))
}
return lines.join('\n')
return convertRowsToCsvWithProvenance(rows).content
}
export function normalizeOutputWorkspaceFileName(outputPath: string): string {
@@ -131,19 +126,149 @@ export function resolveOutputFormat(fileName: string, explicit?: string): Output
return EXT_TO_FORMAT[ext] ?? 'json'
}
export function serializeOutputForFile(output: unknown, format: OutputFormat): string {
interface SerializedOutputFile {
content: string
provenanceValue: unknown
provenanceRepresentations?: readonly WorkspaceFileSecretProvenanceRepresentation[]
provenanceRepresentationsComplete?: boolean
}
function convertRowsToCsvWithProvenance(
rows: Record<string, unknown>[],
registry?: ResolvedSecretTraceRegistry
): {
content: string
representations: readonly WorkspaceFileSecretProvenanceRepresentation[]
representationSourceValues: readonly string[]
representationsComplete: boolean
} {
if (rows.length === 0) {
return {
content: '',
representations: [],
representationSourceValues: [],
representationsComplete: true,
}
}
const headerSet = new Set<string>()
for (const row of rows) {
for (const key of Object.keys(row)) {
headerSet.add(key)
}
}
const headers = [...headerSet]
const representations = new Map<
string,
{
representation: WorkspaceFileSecretProvenanceRepresentation
sourceValue: string
}
>()
let representationsComplete = true
let csvQuoteTransformMatcher: ResolvedSecretMatcher | undefined
if (registry) {
try {
const scanLiterals = new Set<string>()
for (const { plaintext } of registry.getActiveMatches()) {
const jsonEncoded = JSON.stringify(plaintext).slice(1, -1)
if (plaintext.includes('"')) scanLiterals.add(plaintext)
if (jsonEncoded.includes('"')) scanLiterals.add(jsonEncoded)
}
csvQuoteTransformMatcher = createResolvedSecretMatcher(
[...scanLiterals].map((plaintext) => ({ plaintext, replacement: '' }))
)
} catch {
representationsComplete = false
}
}
const serializeCell = (sourceValue: unknown): string => {
const persistedValue = escapeCsvValue(sourceValue)
const serializedSource =
sourceValue === null || sourceValue === undefined
? ''
: typeof sourceValue === 'object'
? JSON.stringify(sourceValue)
: String(sourceValue)
if (
registry &&
csvQuoteTransformMatcher &&
representationsComplete &&
serializedSource.includes('"')
) {
try {
scanResolvedSecretString(serializedSource, csvQuoteTransformMatcher, (scanLiteral) => {
if (!representationsComplete) return
const transformedLiteral = scanLiteral.replace(/"/g, '""')
const sourceProvenance = registry.exportCommittedProvenanceForValue(scanLiteral)
if (!sourceProvenance.complete) {
representationsComplete = false
return
}
if (sourceProvenance.entries.length === 0) return
const representationKey = `${transformedLiteral}\u0000${sourceProvenance.entries
.map((entry) => `${entry.name ?? ''}\u0000${entry.encryptedValue}`)
.join('\u0001')}`
if (representations.has(representationKey)) return
if (representations.size >= MAX_OUTPUT_FILE_PROVENANCE_REPRESENTATIONS) {
representationsComplete = false
return
}
representations.set(representationKey, {
representation: { sourceProvenance, persistedValue: transformedLiteral },
sourceValue: scanLiteral,
})
})
} catch {
representationsComplete = false
}
}
return persistedValue
}
const lines = [headers.map(serializeCell).join(',')]
for (const row of rows) {
lines.push(headers.map((header) => serializeCell(row[header])).join(','))
}
return {
content: lines.join('\n'),
representations: [...representations.values()].map(({ representation }) => representation),
representationSourceValues: [...representations.values()].map(({ sourceValue }) => sourceValue),
representationsComplete,
}
}
function prepareOutputForFile(
output: unknown,
format: OutputFormat,
registry?: ResolvedSecretTraceRegistry
): SerializedOutputFile {
const unwrapped = unwrapFunctionExecuteOutput(output)
if (typeof unwrapped === 'string') return unwrapped
if (typeof unwrapped === 'string') {
return { content: unwrapped, provenanceValue: unwrapped }
}
if (format === 'csv') {
const rows = extractTabularData(unwrapped)
if (rows && rows.length > 0) {
return convertRowsToCsv(rows)
const { content, representations, representationSourceValues, representationsComplete } =
convertRowsToCsvWithProvenance(rows, registry)
return {
content,
provenanceValue: [content, ...representationSourceValues],
provenanceRepresentations: representations,
provenanceRepresentationsComplete: representationsComplete,
}
}
}
return JSON.stringify(unwrapped, null, 2)
const content = JSON.stringify(unwrapped, null, 2)
return { content, provenanceValue: content }
}
export function serializeOutputForFile(output: unknown, format: OutputFormat): string {
return prepareOutputForFile(output, format).content
}
export interface OutputFileDeclaration {
@@ -213,7 +338,6 @@ export async function maybeWriteOutputToFile(
const outputFiles = getOutputFileDeclarations(params).filter((file) => !file.sandboxPath)
if (outputFiles.length === 0) return result
// The tool declared workspace file outputs; passing the successful result
// through without writing them would be a silent no-op the model reads as
// "file written", so fail loudly instead — but keep the computed output so
@@ -230,6 +354,7 @@ export async function maybeWriteOutputToFile(
output: result.output,
}
}
const { userId, workspaceId } = context
const outputObject =
result.output && typeof result.output === 'object' && !Array.isArray(result.output)
@@ -252,13 +377,7 @@ export async function maybeWriteOutputToFile(
const denied = denyOutputWriteWithoutWritePermission(context)
if (denied) return denied
const persistedOutput = projectToolOutputForPersistence(
unwrapFunctionExecuteOutput(result.output),
context.resolvedSecretTraceRegistry
)
if (!persistedOutput.safe) {
return { success: false, error: persistedOutput.error }
}
const registry = context.resolvedSecretTraceRegistry
// Only span the actual write path (where we upload to storage). Fast
// no-op returns above don't need a span — they'd just pad the trace
@@ -267,27 +386,61 @@ export async function maybeWriteOutputToFile(
TraceSpan.CopilotToolsWriteOutputFile,
{
[TraceAttr.ToolName]: toolName,
[TraceAttr.WorkspaceId]: context.workspaceId,
[TraceAttr.WorkspaceId]: workspaceId,
},
async (span) => {
try {
const writtenFiles = []
const preparedByFormat = new Map<
OutputFormat,
Promise<{
buffer: Buffer
secretProvenance: WorkspaceFileSecretProvenance
}>
>()
const preparedFiles = []
for (const outputFile of outputFiles) {
const fileName = normalizeOutputWorkspaceFileName(
outputFile.formatPath ?? outputFile.path
)
const format = resolveOutputFormat(fileName, outputFile.format)
const content = serializeOutputForFile(persistedOutput.value, format)
const contentType = outputFile.mimeType || FORMAT_TO_CONTENT_TYPE[format]
const buffer = Buffer.from(content, 'utf-8')
let prepared = preparedByFormat.get(format)
if (!prepared) {
prepared = (async () => {
const {
content,
provenanceValue,
provenanceRepresentations,
provenanceRepresentationsComplete,
} = prepareOutputForFile(result.output, format, registry)
const decision = await createWorkspaceFileSecretProvenanceFromRegistry(
registry,
content,
{ userId, workspaceId },
provenanceValue,
provenanceRepresentations,
provenanceRepresentationsComplete
)
return {
buffer: Buffer.from(content, 'utf-8'),
secretProvenance: decision.safe ? decision.provenance : { status: 'unknown' },
}
})()
preparedByFormat.set(format, prepared)
}
const { buffer, secretProvenance } = await prepared
preparedFiles.push({ outputFile, format, contentType, buffer, secretProvenance })
}
const writtenFiles = []
for (const { outputFile, format, contentType, buffer, secretProvenance } of preparedFiles) {
if (context.abortSignal?.aborted) {
throw new Error('Request aborted before tool mutation could be applied')
}
const written = await writeWorkspaceFileByPath({
workspaceId: context.workspaceId!,
userId: context.userId!,
workspaceId,
userId,
target: {
path: outputFile.path,
mode: outputFile.mode ?? 'create',
@@ -295,6 +448,7 @@ export async function maybeWriteOutputToFile(
},
buffer,
inferredMimeType: contentType,
secretProvenance,
})
writtenFiles.push({
...written,
apps/sim/lib/copilot/request/tools/resolved-secret-result.test.ts
+40 -22
View File
@@ -24,7 +24,7 @@ describe('projectToolResultForCopilot', () => {
'projects active exact and embedded secrets for %s without mutating runtime output',
(toolName) => {
const registry = createRegistry()
registry.recordResolved('SECRET', 'secret-value')
registry.recordResolved('SECRET', 'secret-value', { propagated: true })
const runtimeResult = {
success: true,
output: {
@@ -51,7 +51,7 @@ describe('projectToolResultForCopilot', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'Test', plaintext: 'Test', encryptedValue: 'ciphertext' },
])
registry.recordResolved('Test', 'Test')
registry.recordResolved('Test', 'Test', { propagated: true })
const runtimeResult = {
success: true,
output: {
@@ -81,7 +81,7 @@ describe('projectToolResultForCopilot', () => {
it('projects both output and error from a failed Function execution', () => {
const registry = createRegistry()
registry.recordResolved('SECRET', 'secret-value')
registry.recordResolved('SECRET', 'secret-value', { propagated: true })
expect(
projectToolResultForCopilot(
@@ -99,9 +99,9 @@ describe('projectToolResultForCopilot', () => {
})
})
it('projects secret-bearing object keys and omits content when replacement collides', () => {
it('projects the model-only copy without mutating raw structural object keys', () => {
const registry = createRegistry()
registry.recordResolved('SECRET', 'secret-value')
registry.recordResolved('SECRET', 'secret-value', { propagated: true })
expect(
projectToolResultForCopilot(
@@ -116,6 +116,13 @@ describe('projectToolResultForCopilot', () => {
output: { 'prefix-{{SECRET}}': 'safe' },
})
const raw = {
success: true,
output: { 'prefix-secret-value': 'safe' },
}
projectToolResultForCopilot(raw, registry)
expect(raw.output).toEqual({ 'prefix-secret-value': 'safe' })
expect(
projectToolResultForCopilot(
{
@@ -133,9 +140,9 @@ describe('projectToolResultForCopilot', () => {
{ name: 'BRACE', plaintext: '{', encryptedValue: 'encrypted-brace' },
{ name: 'JOINED', plaintext: 'ac', encryptedValue: 'encrypted-ac' },
])
registry.recordResolved('MIDDLE', 'B')
registry.recordResolved('BRACE', '{')
registry.recordResolved('JOINED', 'ac')
registry.recordResolved('MIDDLE', 'B', { propagated: true })
registry.recordResolved('BRACE', '{', { propagated: true })
registry.recordResolved('JOINED', 'ac', { propagated: true })
expect(projectToolResultForCopilot({ success: true, output: 'aBc' }, registry)).toEqual({
success: true,
@@ -147,7 +154,7 @@ describe('projectToolResultForCopilot', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'F_SECRET', plaintext: 'F', encryptedValue: 'encrypted-f' },
])
registry.recordResolved('F_SECRET', 'F')
registry.recordResolved('F_SECRET', 'F', { propagated: true })
const projected = projectToolResultForCopilot(
{
@@ -165,9 +172,21 @@ describe('projectToolResultForCopilot', () => {
})
})
it('emits the fixed missing-error message without projecting it as runtime content', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'T_SECRET', plaintext: 'T', encryptedValue: 'encrypted-t' },
])
registry.recordResolved('T_SECRET', 'T', { propagated: true })
expect(projectToolResultForCopilot({ success: false }, registry)).toEqual({
success: false,
error: TOOL_RESULT_UNAVAILABLE_ERROR,
})
})
it('does not project transformed values', () => {
const registry = createRegistry()
registry.recordResolved('SECRET', 'secret-value')
registry.recordResolved('SECRET', 'secret-value', { propagated: true })
const encoded = Buffer.from('secret-value').toString('base64')
expect(
@@ -181,9 +200,9 @@ describe('projectToolResultForCopilot', () => {
{ name: 'BOOLEAN', plaintext: 'true', encryptedValue: 'boolean-ciphertext' },
{ name: 'NULL', plaintext: 'null', encryptedValue: 'null-ciphertext' },
])
registry.recordResolved('NUMBER', '123')
registry.recordResolved('BOOLEAN', 'true')
registry.recordResolved('NULL', 'null')
registry.recordResolved('NUMBER', '123', { propagated: true })
registry.recordResolved('BOOLEAN', 'true', { propagated: true })
registry.recordResolved('NULL', 'null', { propagated: true })
expect(
projectToolResultForCopilot(
@@ -301,9 +320,9 @@ describe('projectToolResultForCopilot', () => {
).toEqual({ success: false, error: TOOL_RESULT_UNAVAILABLE_ERROR })
})
it('projects Copilot-visible resource metadata without changing the runtime result', () => {
it('leaves resource metadata outside plaintext result projection', () => {
const registry = createRegistry()
registry.recordResolved('SECRET', 'secret-value')
registry.recordResolved('SECRET', 'secret-value', { propagated: true })
const result = {
success: true,
resources: [
@@ -322,7 +341,7 @@ describe('projectToolResultForCopilot', () => {
{
type: 'file',
id: 'file-1',
title: '{{SECRET}}.txt',
title: 'secret-value.txt',
path: '/workspace/report.txt',
},
],
@@ -343,9 +362,9 @@ describe('projectToolResultForCopilot', () => {
title: 'report.txt',
path: '/workspace/secret-value/report.txt',
},
])('omits resources whose routing controls contain a secret', (resource) => {
])('leaves resource routing controls outside plaintext result projection', (resource) => {
const registry = createRegistry()
registry.recordResolved('SECRET', 'secret-value')
registry.recordResolved('SECRET', 'secret-value', { propagated: true })
const projected = projectToolResultForCopilot(
{
success: true,
@@ -355,18 +374,17 @@ describe('projectToolResultForCopilot', () => {
registry
)
expect(projected).toEqual({ success: true, output: {}, resources: [] })
expect(JSON.stringify(projected)).not.toContain('secret-value')
expect(projected).toEqual({ success: true, output: {}, resources: [resource] })
})
it('projects every tool result once provenance is active', () => {
it('does not project a tool result from merely active input provenance', () => {
const registry = createRegistry()
registry.recordResolved('SECRET', 'secret-value')
const result = { success: true, output: 'secret-value' }
expect(projectToolResultForCopilot(result, registry)).toEqual({
success: true,
output: '{{SECRET}}',
output: 'secret-value',
})
expect(result).toEqual({ success: true, output: 'secret-value' })
})
apps/sim/lib/copilot/request/tools/resolved-secret-result.ts
+11 -109
View File
@@ -1,113 +1,23 @@
import { isPlainRecord } from '@sim/utils/object'
import type { MothershipResource } from '@/lib/copilot/resources/types'
import type { ToolExecutionResult } from '@/lib/copilot/tool-executor/types'
import {
isResolvedSecretModelContentUnchanged,
projectResolvedSecretModelContent,
projectResolvedSecretModelControlMessage,
projectResolvedSecretModelJsonContent,
} from '@/executor/utils/resolved-secret-content-projection'
import { projectResolvedSecretModelJsonContent } from '@/executor/utils/resolved-secret-content-projection'
import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
export const TOOL_RESULT_UNAVAILABLE_ERROR =
'Tool execution settled, but its result could not be returned safely. Do not retry a mutation automatically.'
export const TOOL_OUTPUT_PERSISTENCE_UNAVAILABLE_ERROR =
'Tool output could not be persisted safely because secret provenance was unavailable.'
function structuralResult(result: ToolExecutionResult): ToolExecutionResult {
return { success: result.success === true }
}
function resourceContent(
resources: MothershipResource[]
): Array<{ type: string; id: string; title: string; path?: string }> {
return resources.map((resource) => ({
type: resource.type,
id: resource.id,
title: resource.title,
...(resource.path !== undefined ? { path: resource.path } : {}),
}))
}
function modelSafeResources(
resources: MothershipResource[],
registry: ResolvedSecretTraceRegistry | undefined
): MothershipResource[] {
return resources.filter((resource) =>
isResolvedSecretModelContentUnchanged([resource.type, resource.id, resource.path], registry)
)
}
function restoreProjectedResources(
resources: MothershipResource[],
projectedContent: unknown
): MothershipResource[] | undefined {
if (!Array.isArray(projectedContent) || projectedContent.length !== resources.length) {
return undefined
}
const projectedResources: MothershipResource[] = []
for (let index = 0; index < resources.length; index += 1) {
const resource = resources[index]
const content = projectedContent[index]
if (
!isPlainRecord(content) ||
typeof content.type !== 'string' ||
typeof content.id !== 'string' ||
typeof content.title !== 'string' ||
(content.path !== undefined && typeof content.path !== 'string')
) {
return undefined
}
if (content.type !== resource.type || content.id !== resource.id) continue
projectedResources.push({
type: resource.type,
id: resource.id,
title: content.title,
...(content.path !== undefined ? { path: content.path } : {}),
})
}
return projectedResources
}
function omittedResult(
result: ToolExecutionResult,
registry: ResolvedSecretTraceRegistry | undefined
): ToolExecutionResult {
function omittedResult(result: ToolExecutionResult): ToolExecutionResult {
if (result.success) return { success: true }
const error =
projectResolvedSecretModelControlMessage(TOOL_RESULT_UNAVAILABLE_ERROR, registry) ??
TOOL_RESULT_UNAVAILABLE_ERROR
return { success: false, error }
return { success: false, error: TOOL_RESULT_UNAVAILABLE_ERROR }
}
export type CopilotToolResultProjection =
| { safe: true; result: ToolExecutionResult }
| { safe: false; result: ToolExecutionResult }
export type CopilotPersistedOutputProjection =
| { safe: true; value: unknown }
| { safe: false; error: string }
/**
* Projects only the exact value a Copilot tool is about to persist. The isolated per-tool registry
* makes this causal: values activated by this tool become canonical `{{NAME}}` aliases, while a
* public low-entropy value cannot be rewritten merely because an earlier sibling activated the
* same bytes. Persisting the alias makes later reads safe without a second provenance store.
*/
export function projectToolOutputForPersistence(
value: unknown,
registry: ResolvedSecretTraceRegistry | undefined
): CopilotPersistedOutputProjection {
const projection = projectResolvedSecretModelContent(value, registry)
return projection.safe
? { safe: true, value: projection.value }
: { safe: false, error: TOOL_OUTPUT_PERSISTENCE_UNAVAILABLE_ERROR }
}
/**
* Projects terminal tool content and reports whether the complete content was safe to cross.
* Callers that isolate provenance per tool call may merge that child registry only when `safe`
@@ -119,15 +29,14 @@ export function inspectToolResultForCopilot(
registry: ResolvedSecretTraceRegistry | undefined
): CopilotToolResultProjection {
try {
const resultRegistry = registry?.forkForPropagatedEntries()
const content: Record<string, unknown> = {}
const resources =
result.resources !== undefined ? modelSafeResources(result.resources, registry) : undefined
const resources = result.resources
if (Object.hasOwn(result, 'output')) content.output = result.output
if (Object.hasOwn(result, 'error')) content.error = result.error
if (resources !== undefined) content.resources = resourceContent(resources)
const projection = projectResolvedSecretModelJsonContent(content, registry)
const projection = projectResolvedSecretModelJsonContent(content, resultRegistry)
if (!projection.safe || !projection.value || typeof projection.value !== 'object') {
return { safe: false, result: omittedResult(result, registry) }
return { safe: false, result: omittedResult(result) }
}
const projectedContent = projection.value as Record<string, unknown>
@@ -135,26 +44,19 @@ export function inspectToolResultForCopilot(
if (Object.hasOwn(projectedContent, 'output')) projected.output = projectedContent.output
if (Object.hasOwn(projectedContent, 'error')) {
if (typeof projectedContent.error !== 'string') {
return { safe: false, result: omittedResult(result, registry) }
return { safe: false, result: omittedResult(result) }
}
projected.error = projectedContent.error
}
if (resources !== undefined) {
const projectedResources = restoreProjectedResources(resources, projectedContent.resources)
if (!projectedResources) {
return { safe: false, result: omittedResult(result, registry) }
}
projected.resources = projectedResources
projected.resources = resources
}
if (!projected.success && !projected.error) {
projected.error = projectResolvedSecretModelControlMessage(
TOOL_RESULT_UNAVAILABLE_ERROR,
registry
)
projected.error = TOOL_RESULT_UNAVAILABLE_ERROR
}
return { safe: true, result: projected }
} catch {
return { safe: false, result: omittedResult(result, registry) }
return { safe: false, result: omittedResult(result) }
}
}
apps/sim/lib/copilot/request/tools/tables.test.ts
+348 -65
View File
@@ -2,7 +2,7 @@
* @vitest-environment node
*/
import { loggerMock } from '@sim/testing'
import { encryptionMock, encryptionMockFns, loggerMock } from '@sim/testing'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import type { TableDefinition } from '@/lib/table'
@@ -20,6 +20,8 @@ vi.mock('@/lib/table/rows/service', () => ({
replaceTableRows: mockReplaceTableRows,
}))
vi.mock('@/lib/core/security/encryption', () => encryptionMock)
vi.mock('@/lib/copilot/request/otel', () => ({
withCopilotSpan: (
_name: string,
@@ -147,22 +149,25 @@ describe('maybeWriteOutputToTable', () => {
expect(table.id).toBe('tbl_1')
})
it('projects activated secrets before persistence without rewriting sibling literals', async () => {
const parentRegistry = new ResolvedSecretTraceRegistry([
{
name: 'OUTPUT_SECRET',
plaintext: 'secret-value',
encryptedValue: 'encrypted-output-secret',
},
{
name: 'UNRELATED',
plaintext: 'true',
encryptedValue: 'encrypted-unrelated',
},
])
it('persists raw values with per-cell provenance without rewriting sibling literals', async () => {
const parentRegistry = new ResolvedSecretTraceRegistry(
[
{
name: 'OUTPUT_SECRET',
plaintext: 'secret-value',
encryptedValue: 'encrypted-output-secret',
},
{
name: 'UNRELATED',
plaintext: 'true',
encryptedValue: 'encrypted-unrelated',
},
],
{ userId: 'user-1', workspaceId: 'workspace-1' }
)
parentRegistry.recordResolved('UNRELATED', 'true')
const toolRegistry = parentRegistry.forkForToolInput({ code: 'return {{OUTPUT_SECRET}}' })
toolRegistry.recordResolved('OUTPUT_SECRET', 'secret-value')
const toolRegistry = parentRegistry.forkForInputPaths([])
toolRegistry.recordResolved('OUTPUT_SECRET', 'secret-value', { propagated: true })
const runtimeRows = [{ name: 'secret-value', age: '123', status: 'true' }]
const result = await maybeWriteOutputToTable(
@@ -173,9 +178,35 @@ describe('maybeWriteOutputToTable', () => {
)
expect(result.success).toBe(true)
const persistedRows = mockReplaceTableRows.mock.calls[0][0].rows
const persistedWrite = mockReplaceTableRows.mock.calls[0][0]
const persistedRows = persistedWrite.rows
expect(persistedRows).toEqual([
{ col_name: '{{OUTPUT_SECRET}}', col_age: '123', col_status: 'true' },
{ col_name: 'secret-value', col_age: '123', col_status: 'true' },
])
expect(persistedWrite.secretProvenance).toEqual([
{
complete: true,
columns: {
col_name: {
version: 1,
complete: true,
entries: [{ name: 'OUTPUT_SECRET', encryptedValue: 'encrypted-output-secret' }],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
},
col_age: {
version: 1,
complete: true,
entries: [],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
},
col_status: {
version: 1,
complete: true,
entries: [],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
},
},
},
])
expect(runtimeRows).toEqual([{ name: 'secret-value', age: '123', status: 'true' }])
@@ -189,14 +220,164 @@ describe('maybeWriteOutputToTable', () => {
},
})
encryptionMockFns.mockDecryptSecret.mockResolvedValue({ decrypted: 'secret-value' })
const readRegistry = new ResolvedSecretTraceRegistry([], {
userId: 'user-1',
workspaceId: 'workspace-1',
})
expect(
await readRegistry.importCrossingProvenance(
persistedWrite.secretProvenance[0].columns.col_name,
persistedRows,
{ trusted: true }
)
).toBe(true)
const laterRead = projectToolResultForCopilot(
{ success: true, output: { data: { rows: persistedRows } } },
new ResolvedSecretTraceRegistry()
readRegistry
)
expect(laterRead.output).toEqual({ data: { rows: persistedRows } })
expect(laterRead.output).toEqual({
data: {
rows: [{ col_name: '{{OUTPUT_SECRET}}', col_age: '123', col_status: 'true' }],
},
})
})
it('does not write when table persistence provenance is incomplete', async () => {
it('never rewrites stored public text when an active secret has a common value', async () => {
const registry = new ResolvedSecretTraceRegistry(
[{ name: 'SHORT_SECRET', plaintext: 'x', encryptedValue: 'encrypted-short-secret' }],
{ userId: 'user-1', workspaceId: 'workspace-1' }
)
registry.recordResolved('SHORT_SECRET', 'x')
const rows = [{ name: 'Box eSign' }, { name: 'Brex' }, { name: 'hex' }]
const result = await maybeWriteOutputToTable(
FunctionExecute.id,
{ outputTable: 'tbl_1' },
{ success: true, output: { result: rows } },
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
expect(mockReplaceTableRows.mock.calls[0][0].rows).toEqual([
{ col_name: 'Box eSign' },
{ col_name: 'Brex' },
{ col_name: 'hex' },
])
})
it('binds provenance to the values produced by table coercion', async () => {
const registry = new ResolvedSecretTraceRegistry(
[
{ name: 'NUMBER_SECRET', plaintext: '123', encryptedValue: 'encrypted-number' },
{ name: 'INVALID_SECRET', plaintext: 'not-a-number', encryptedValue: 'encrypted-invalid' },
],
{ userId: 'user-1', workspaceId: 'workspace-1' }
)
registry.recordResolved('NUMBER_SECRET', '123')
registry.recordResolved('INVALID_SECRET', 'not-a-number')
const result = await maybeWriteOutputToTable(
FunctionExecute.id,
{ outputTable: 'tbl_1' },
{
success: true,
output: { result: [{ age: '123' }, { age: 'not-a-number' }] },
},
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
expect(mockReplaceTableRows.mock.calls[0][0].secretProvenance).toEqual([
{
complete: true,
columns: {
col_age: {
version: 1,
complete: true,
entries: [{ name: 'NUMBER_SECRET', encryptedValue: 'encrypted-number' }],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
},
},
},
{
complete: true,
columns: {
col_age: {
version: 1,
complete: true,
entries: [],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
},
},
},
])
})
it('accepts same-workspace provenance from a different actor and preserves its source', async () => {
const registry = new ResolvedSecretTraceRegistry(
[{ name: 'OUTPUT_SECRET', plaintext: 'secret-value', encryptedValue: 'encrypted-secret' }],
{ userId: 'workflow-owner', workspaceId: 'workspace-1' }
)
registry.recordResolved('OUTPUT_SECRET', 'secret-value')
const result = await maybeWriteOutputToTable(
FunctionExecute.id,
{ outputTable: 'tbl_1' },
{ success: true, output: { result: [{ name: 'secret-value' }] } },
buildContext({ userId: 'billing-actor', resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
expect(
mockReplaceTableRows.mock.calls[0][0].secretProvenance[0].columns.col_name.scope
).toEqual({ userId: 'workflow-owner', workspaceId: 'workspace-1' })
})
it('persists raw rows with unknown provenance when the source workspace differs', async () => {
const registry = new ResolvedSecretTraceRegistry(
[{ name: 'OUTPUT_SECRET', plaintext: 'secret-value', encryptedValue: 'encrypted-secret' }],
{ userId: 'user-1', workspaceId: 'workspace-2' }
)
registry.recordResolved('OUTPUT_SECRET', 'secret-value')
const result = await maybeWriteOutputToTable(
FunctionExecute.id,
{ outputTable: 'tbl_1' },
{ success: true, output: { result: [{ name: 'secret-value' }] } },
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
expect(mockReplaceTableRows).toHaveBeenCalledWith(
expect.objectContaining({ secretProvenance: [{ complete: false, columns: {} }] }),
expect.anything(),
expect.any(String)
)
})
it('persists raw rows with unknown provenance when the source scope is unavailable', async () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'OUTPUT_SECRET', plaintext: 'secret-value', encryptedValue: 'encrypted-secret' },
])
registry.recordResolved('OUTPUT_SECRET', 'secret-value')
const result = await maybeWriteOutputToTable(
FunctionExecute.id,
{ outputTable: 'tbl_1' },
{ success: true, output: { result: [{ name: 'secret-value' }] } },
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
expect(mockReplaceTableRows).toHaveBeenCalledWith(
expect.objectContaining({ secretProvenance: [{ complete: false, columns: {} }] }),
expect.anything(),
expect.any(String)
)
})
it('persists raw rows with unknown provenance when lineage is incomplete', async () => {
const registry = new ResolvedSecretTraceRegistry()
registry.markIncomplete()
@@ -207,14 +388,15 @@ describe('maybeWriteOutputToTable', () => {
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result).toEqual({
success: false,
error: 'Tool output could not be persisted safely because secret provenance was unavailable.',
})
expect(mockReplaceTableRows).not.toHaveBeenCalled()
expect(result.success).toBe(true)
expect(mockReplaceTableRows).toHaveBeenCalledWith(
expect.objectContaining({ secretProvenance: [{ complete: false, columns: {} }] }),
expect.anything(),
expect.any(String)
)
})
it('preserves legacy table writes when execution provenance is unavailable', async () => {
it('preserves legacy table writes without certifying unavailable provenance', async () => {
const result = await maybeWriteOutputToTable(
FunctionExecute.id,
{ outputTable: 'tbl_1' },
@@ -224,7 +406,10 @@ describe('maybeWriteOutputToTable', () => {
expect(result.success).toBe(true)
expect(mockReplaceTableRows).toHaveBeenCalledWith(
expect.objectContaining({ rows: [{ col_name: 'unknown' }] }),
expect.objectContaining({
rows: [{ col_name: 'unknown' }],
secretProvenance: [{ complete: false, columns: {} }],
}),
expect.anything(),
expect.any(String)
)
@@ -271,10 +456,11 @@ describe('maybeWriteOutputToTable', () => {
})
it('keeps raw errors for terminal projection but projects application logs and OTel events', async () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'SECRET', plaintext: 'secret-value', encryptedValue: 'encrypted-secret-value' },
])
registry.recordResolved('SECRET', 'secret-value')
const registry = new ResolvedSecretTraceRegistry(
[{ name: 'SECRET', plaintext: 'secret-value', encryptedValue: 'encrypted-secret-value' }],
{ userId: 'user-1', workspaceId: 'workspace-1' }
)
registry.recordResolved('SECRET', 'secret-value', { propagated: true })
mockReplaceTableRows.mockRejectedValue(new Error('Duplicate value "secret-value"'))
const result = await maybeWriteOutputToTable(
@@ -343,11 +529,14 @@ describe('maybeWriteReadCsvToTable', () => {
])
})
it('projects active secret literals into string-compatible CSV columns', async () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'NUMBER', plaintext: '123', encryptedValue: 'encrypted-number' },
{ name: 'BOOLEAN', plaintext: 'true', encryptedValue: 'encrypted-boolean' },
])
it('persists raw CSV cells with per-cell secret provenance', async () => {
const registry = new ResolvedSecretTraceRegistry(
[
{ name: 'NUMBER', plaintext: '123', encryptedValue: 'encrypted-number' },
{ name: 'BOOLEAN', plaintext: 'true', encryptedValue: 'encrypted-boolean' },
],
{ userId: 'user-1', workspaceId: 'workspace-1' }
)
registry.recordResolved('NUMBER', '123')
registry.recordResolved('BOOLEAN', 'true')
@@ -363,8 +552,27 @@ describe('maybeWriteReadCsvToTable', () => {
expect.objectContaining({
rows: [
{
col_name: '{{NUMBER}}',
col_status: '{{BOOLEAN}}',
col_name: '123',
col_status: 'true',
},
],
secretProvenance: [
{
complete: true,
columns: {
col_name: {
version: 1,
complete: true,
entries: [{ name: 'NUMBER', encryptedValue: 'encrypted-number' }],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
},
col_status: {
version: 1,
complete: true,
entries: [{ name: 'BOOLEAN', encryptedValue: 'encrypted-boolean' }],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
},
},
},
],
}),
@@ -373,32 +581,104 @@ describe('maybeWriteReadCsvToTable', () => {
)
})
it('rejects active secret literals in number and boolean columns before mutation', async () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'NUMBER', plaintext: '123', encryptedValue: 'encrypted-number' },
{ name: 'BOOLEAN', plaintext: 'true', encryptedValue: 'encrypted-boolean' },
])
it('retains original provenance after a quote-escaped CSV round trip', async () => {
encryptionMockFns.mockDecryptSecret.mockImplementation(async (encryptedValue: string) => ({
decrypted: encryptedValue === 'encrypted-original' ? 'a"b' : '"a""b"',
}))
const registry = new ResolvedSecretTraceRegistry([], {
userId: 'user-1',
workspaceId: 'workspace-1',
})
await expect(
registry.importProvenance(
{
version: 1,
complete: true,
entries: [
{ name: 'CSV_SECRET', encryptedValue: 'encrypted-original' },
{ name: 'CSV_SECRET', encryptedValue: 'encrypted-representation' },
],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
},
{ trusted: true }
)
).resolves.toBe(true)
const result = await maybeWriteReadCsvToTable(
ReadTool.id,
{ outputTable: 'tbl_1', path: 'files/people.csv' },
{ success: true, output: { content: 'name\n"a""b"' } },
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result.success).toBe(true)
expect(mockReplaceTableRows).toHaveBeenCalledWith(
expect.objectContaining({
rows: [{ col_name: 'a"b' }],
secretProvenance: [
{
complete: true,
columns: {
col_name: {
version: 1,
complete: true,
entries: [{ name: 'CSV_SECRET', encryptedValue: 'encrypted-original' }],
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
},
},
},
],
}),
expect.anything(),
expect.any(String)
)
})
it('preserves numeric and boolean cells while recording their provenance', async () => {
const registry = new ResolvedSecretTraceRegistry(
[
{ name: 'NUMBER', plaintext: '123', encryptedValue: 'encrypted-number' },
{ name: 'BOOLEAN', plaintext: 'true', encryptedValue: 'encrypted-boolean' },
],
{ userId: 'user-1', workspaceId: 'workspace-1' }
)
registry.recordResolved('NUMBER', '123')
registry.recordResolved('BOOLEAN', 'true')
const result = await maybeWriteReadCsvToTable(
ReadTool.id,
{ outputTable: 'tbl_1', path: 'files/people.csv' },
{ success: true, output: { content: 'name,age,active\nAlice,123,true' } },
{ outputTable: 'tbl_1', path: 'files/people.json' },
{
success: true,
output: { content: '[{"name":"Alice","age":123,"active":true}]' },
},
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result).toEqual({
success: false,
error:
'Tool output could not be persisted safely because a resolved secret is incompatible with the target column type.',
})
expect(mockReplaceTableRows).not.toHaveBeenCalled()
expect(JSON.stringify(result)).not.toContain('123')
expect(JSON.stringify(result)).not.toContain('true')
expect(result.success).toBe(true)
expect(mockReplaceTableRows).toHaveBeenCalledWith(
expect.objectContaining({
rows: [{ col_name: 'Alice', col_age: 123, col_active: true }],
secretProvenance: [
expect.objectContaining({
complete: true,
columns: expect.objectContaining({
col_age: expect.objectContaining({
entries: [{ name: 'NUMBER', encryptedValue: 'encrypted-number' }],
}),
col_active: expect.objectContaining({
entries: [{ name: 'BOOLEAN', encryptedValue: 'encrypted-boolean' }],
}),
}),
}),
],
}),
expect.anything(),
expect.any(String)
)
})
it('does not import CSV rows when persistence provenance is incomplete', async () => {
it('imports raw CSV rows with unknown provenance when lineage is incomplete', async () => {
const registry = new ResolvedSecretTraceRegistry()
registry.markIncomplete()
@@ -409,14 +689,15 @@ describe('maybeWriteReadCsvToTable', () => {
buildContext({ resolvedSecretTraceRegistry: registry })
)
expect(result).toEqual({
success: false,
error: 'Tool output could not be persisted safely because secret provenance was unavailable.',
})
expect(mockReplaceTableRows).not.toHaveBeenCalled()
expect(result.success).toBe(true)
expect(mockReplaceTableRows).toHaveBeenCalledWith(
expect.objectContaining({ secretProvenance: [{ complete: false, columns: {} }] }),
expect.anything(),
expect.any(String)
)
})
it('preserves legacy CSV imports when execution provenance is unavailable', async () => {
it('preserves legacy CSV imports without certifying unavailable provenance', async () => {
const result = await maybeWriteReadCsvToTable(
ReadTool.id,
{ outputTable: 'tbl_1', path: 'files/people.csv' },
@@ -428,6 +709,7 @@ describe('maybeWriteReadCsvToTable', () => {
expect(mockReplaceTableRows).toHaveBeenCalledWith(
expect.objectContaining({
rows: [{ col_name: 'legacy-value', col_age: '123', col_active: 'true' }],
secretProvenance: [{ complete: false, columns: {} }],
}),
expect.anything(),
expect.any(String)
@@ -462,10 +744,11 @@ describe('maybeWriteReadCsvToTable', () => {
})
it('projects active secret literals in CSV-import log and OTel errors', async () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'SECRET', plaintext: 'secret-value', encryptedValue: 'encrypted-secret-value' },
])
registry.recordResolved('SECRET', 'secret-value')
const registry = new ResolvedSecretTraceRegistry(
[{ name: 'SECRET', plaintext: 'secret-value', encryptedValue: 'encrypted-secret-value' }],
{ userId: 'user-1', workspaceId: 'workspace-1' }
)
registry.recordResolved('SECRET', 'secret-value', { propagated: true })
mockReplaceTableRows.mockRejectedValue(new Error('Duplicate value "secret-value"'))
const result = await maybeWriteReadCsvToTable(
apps/sim/lib/copilot/request/tools/tables.ts
+28 -41
View File
@@ -1,4 +1,3 @@
import { isDeepStrictEqual } from 'node:util'
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { generateId } from '@sim/utils/id'
@@ -11,40 +10,22 @@ import { TraceEvent } from '@/lib/copilot/generated/trace-events-v1'
import { TraceSpan } from '@/lib/copilot/generated/trace-spans-v1'
import { withCopilotSpan } from '@/lib/copilot/request/otel'
import { denyOutputWriteWithoutWritePermission } from '@/lib/copilot/request/tools/permissions'
import {
projectToolErrorMessageForCopilot,
projectToolOutputForPersistence,
} from '@/lib/copilot/request/tools/resolved-secret-result'
import { projectToolErrorMessageForCopilot } from '@/lib/copilot/request/tools/resolved-secret-result'
import type { ExecutionContext, ToolCallResult } from '@/lib/copilot/request/types'
import { isPrivateSecretProvenanceScopeCompatible } from '@/lib/execution/durable-secret-provenance'
import type { RowData, TableDefinition } from '@/lib/table'
import { buildIdByName, rowDataNameToId } from '@/lib/table/column-keys'
import { columnTypeOf } from '@/lib/table/column-types'
import { createExactEmptyTableRowSecretProvenance } from '@/lib/table/rows/secret-provenance'
import {
createTableRowSecretProvenanceFromRegistry,
createUnknownTableRowSecretProvenance,
} from '@/lib/table/rows/secret-provenance'
import { replaceTableRows } from '@/lib/table/rows/service'
import { getTableById } from '@/lib/table/service'
import { coerceRowValues } from '@/lib/table/validation'
const logger = createLogger('CopilotToolResultTables')
const MAX_OUTPUT_TABLE_ROWS = 10_000
const TABLE_SECRET_PROJECTION_UNSUPPORTED_ERROR =
'Tool output could not be persisted safely because a resolved secret is incompatible with the target column type.'
function hasUnsupportedProjectedCell(
table: TableDefinition,
sourceRows: Array<Record<string, unknown>>,
projectedRows: Array<Record<string, unknown>>
): boolean {
const columnsByName = new Map(table.schema.columns.map((column) => [column.name, column]))
for (let rowIndex = 0; rowIndex < projectedRows.length; rowIndex += 1) {
for (const [name, projectedValue] of Object.entries(projectedRows[rowIndex])) {
const column = columnsByName.get(name)
if (!column || isDeepStrictEqual(sourceRows[rowIndex]?.[name], projectedValue)) continue
const type = columnTypeOf(column).id
if (type !== 'string' && type !== 'json') return true
}
}
return false
}
/**
* Replaces a table's rows with wire rows keyed by column name. Translates the
@@ -57,37 +38,43 @@ async function replaceTableRowsFromWire(
rows: Array<Record<string, unknown>>,
context: ExecutionContext
): Promise<{ error?: string }> {
const persistenceProjection = context.resolvedSecretTraceRegistry
? projectToolOutputForPersistence(rows, context.resolvedSecretTraceRegistry)
: { safe: true as const, value: rows }
if (!persistenceProjection.safe) return { error: persistenceProjection.error }
if (
!Array.isArray(persistenceProjection.value) ||
!persistenceProjection.value.every(isPlainRecord)
) {
if (!rows.every(isPlainRecord)) {
return { error: 'Table rows could not be persisted safely' }
}
if (hasUnsupportedProjectedCell(table, rows, persistenceProjection.value)) {
return { error: TABLE_SECRET_PROJECTION_UNSUPPORTED_ERROR }
}
const idByName = buildIdByName(table.schema)
const idKeyedRows = persistenceProjection.value.map((row) =>
rowDataNameToId(row as RowData, idByName)
)
const idKeyedRows = rows.map((row) => rowDataNameToId(row as RowData, idByName))
const emptyIndex = idKeyedRows.findIndex((row) => Object.keys(row).length === 0)
if (emptyIndex !== -1) {
return {
error: `Row ${emptyIndex + 1} has no keys matching columns on table "${table.name}" (columns: ${table.schema.columns.map((c) => c.name).join(', ')})`,
}
}
const registry = context.resolvedSecretTraceRegistry
const persistedRows = idKeyedRows.map((row) => {
const persistedRow = { ...row }
coerceRowValues(persistedRow, table.schema)
return persistedRow
})
const destinationScope = { userId: context.userId, workspaceId: table.workspaceId }
const secretProvenance = persistedRows.map((row) => {
if (!registry) return createUnknownTableRowSecretProvenance()
const provenance = createTableRowSecretProvenanceFromRegistry(row, registry)
if (!provenance.complete) return createUnknownTableRowSecretProvenance()
const compatible = Object.values(provenance.columns).every(
(columnProvenance) =>
columnProvenance.entries.length === 0 ||
isPrivateSecretProvenanceScopeCompatible(columnProvenance.scope, destinationScope)
)
return compatible ? provenance : createUnknownTableRowSecretProvenance()
})
await replaceTableRows(
{
tableId: table.id,
rows: idKeyedRows,
workspaceId: table.workspaceId,
userId: context.userId,
secretProvenance: idKeyedRows.map(createExactEmptyTableRowSecretProvenance),
secretProvenance,
},
table,
generateId().slice(0, 8)
apps/sim/lib/copilot/tool-executor/executor.test.ts
+1 -1
View File
@@ -94,7 +94,7 @@ describe('copilot tool executor fallback', () => {
const registry = new ResolvedSecretTraceRegistry([
{ name: 'API_KEY', plaintext: secret, encryptedValue: 'encrypted-secret' },
])
registry.recordResolved('API_KEY', secret)
registry.recordResolved('API_KEY', secret, { propagated: true })
isKnownTool.mockReturnValue(true)
isSimExecuted.mockReturnValue(true)
isClientExecuted.mockReturnValue(false)
apps/sim/lib/copilot/tools/handlers/deployment/custom-block.test.ts
+3
View File
@@ -329,11 +329,14 @@ describe('executeDeployCustomBlock', () => {
it('ingests a workspace-file icon into public icon storage', async () => {
listWorkspaceFilesMock.mockResolvedValue([
{
id: 'file-1',
workspaceId: 'ws-1',
name: 'icon.png',
folderPath: null,
type: 'image/png',
size: 1024,
key: 'workspace/ws-1/123-abc-icon.png',
storageContext: 'workspace',
},
])
fetchWorkspaceFileBufferMock.mockResolvedValue(Buffer.from('png-bytes'))
apps/sim/lib/copilot/tools/handlers/deployment/custom-block.ts
-1
View File
@@ -68,7 +68,6 @@ async function resolveIconUrl(
if (record.size > MAX_ICON_BYTES) {
throw new CustomBlockValidationError('Icon file must be 5MB or smaller')
}
const buffer = await fetchWorkspaceFileBuffer(record)
const safeFileName = record.name.replace(/[^a-zA-Z0-9.-]/g, '_')
const uploaded = await uploadFile({

Some files were not shown because too many files have changed in this diff Show More