diff --git a/apps/sim/lib/api/contracts/hotspots.ts b/apps/sim/lib/api/contracts/hotspots.ts index c0f05e51587..82f24c81734 100644 --- a/apps/sim/lib/api/contracts/hotspots.ts +++ b/apps/sim/lib/api/contracts/hotspots.ts @@ -10,6 +10,7 @@ import { defineRouteContract } from '@/lib/api/contracts/types' import { DEFAULT_CODE_LANGUAGE } from '@/lib/execution/languages' import { PRIVATE_SECRET_PROVENANCE_FIELD } from '@/lib/execution/private-tool-metadata' import { MAX_BLOCK_MOUNTED_FILES } from '@/lib/execution/remote-sandbox/sandbox-paths' +import { MAX_FUNCTION_CODE_LENGTH } from '@/lib/function-execution/limits' import { MAX_PII_VALIDATION_DETECTED_ENTITIES, MAX_PII_VALIDATION_TEXT_CHARACTERS, @@ -165,8 +166,8 @@ const functionOutputFileSchema = z export const functionExecuteBodySchema = z .object({ - code: z.string().min(1, 'Code is required'), - sourceCode: z.string().optional(), + code: z.string().min(1, 'Code is required').max(MAX_FUNCTION_CODE_LENGTH), + sourceCode: z.string().max(MAX_FUNCTION_CODE_LENGTH).optional(), params: unknownRecordSchema.optional().default({}), timeout: z.coerce.number().int().positive().optional(), language: z.string().optional().default(DEFAULT_CODE_LANGUAGE), diff --git a/apps/sim/lib/api/contracts/mothership-chats.ts b/apps/sim/lib/api/contracts/mothership-chats.ts index b9006b462de..16b0b2a6a7b 100644 --- a/apps/sim/lib/api/contracts/mothership-chats.ts +++ b/apps/sim/lib/api/contracts/mothership-chats.ts @@ -9,6 +9,11 @@ import { } from '@/lib/api/contracts/secret-mount-policy' import { defineRouteContract } from '@/lib/api/contracts/types' import type { RESOLVED_SECRET_PROVENANCE_FIELD } from '@/lib/execution/private-tool-metadata' +import { + MAX_CHAT_CONTEXT_LABEL_LENGTH, + MAX_CHAT_CONTEXTS, + MAX_CHAT_MESSAGE_LENGTH, +} from '@/lib/mothership/chat/context-limits' import { ChatPayloadSchema } from '@/lib/mothership/generated/protocol' import type { AgentStreamEvent, TextDeltaClassification } from '@/providers/stream-events' @@ -71,7 +76,11 @@ export const markMothershipChatReadContract = defineRouteContract({ const mothershipExecuteMessageSchema = z.object({ role: z.enum(['system', 'user', 'assistant']), - content: z.string(), + content: z.string().max(MAX_CHAT_MESSAGE_LENGTH), +}) + +const mothershipContextInputSchema = scheduleContextSchema.extend({ + label: z.string().max(MAX_CHAT_CONTEXT_LABEL_LENGTH), }) const mothershipExecuteFileAttachmentSchema = z @@ -134,7 +143,7 @@ export const mothershipExecuteBodySchema = z.object({ * mirroring the interactive chat path. Headless executions use this to pass * captured contexts into the run without a live client. */ - contexts: z.array(scheduleContextSchema).optional(), + contexts: z.array(mothershipContextInputSchema).max(MAX_CHAT_CONTEXTS).optional(), mcpTools: z.array(mothershipExecuteMcpToolSchema).optional(), workflowId: z.string().optional(), executionId: z.string().optional(), @@ -162,7 +171,8 @@ export const mothershipChatGetQuerySchema = z export const mothershipChatPostEnvelopeSchema = z .object({ - message: z.string().optional(), + message: z.string().max(MAX_CHAT_MESSAGE_LENGTH).optional(), + contexts: z.array(mothershipContextInputSchema).max(MAX_CHAT_CONTEXTS).optional(), chatId: z.string().optional(), workflowId: z.string().optional(), workspaceId: z.string().optional(), diff --git a/apps/sim/lib/function-execution/execute-request.test.ts b/apps/sim/lib/function-execution/execute-request.test.ts index 868a2127c15..474c5688f7f 100644 --- a/apps/sim/lib/function-execution/execute-request.test.ts +++ b/apps/sim/lib/function-execution/execute-request.test.ts @@ -3157,6 +3157,53 @@ describe('Function execution request', () => { }) describe('Template Variable Resolution', () => { + it.each([ + { code: ' '.repeat(1024 * 1024 + 1), reason: 'source length' }, + { + code: 'return 1', + sourceCode: ' '.repeat(1024 * 1024 + 1), + reason: 'diagnostic source length', + }, + { code: `/* ${'< a.value>'.repeat(10_001)} */ return 1`, reason: 'block references' }, + { + code: `/* ${''.repeat(10_001)} */ return 1`, + reason: 'workflow references', + }, + ])('rejects excessive $reason before execution', async ({ code, sourceCode }) => { + const response = await POST( + createMockRequest('POST', { + code, + sourceCode, + workflowVariables: { total: { name: 'total', type: 'number', value: 1 } }, + }) + ) + expect(response.status).toBe(400) + }) + + it('preserves repeated values, legacy variable precedence, and skipped-block references', async () => { + mockExecuteInIsolatedVM.mockImplementationOnce(async (request) => ({ + result: await runInNewContext(`(async () => { ${request.code} })()`, { + ...request.contextVariables, + }), + stdout: '', + })) + const response = await POST( + createMockRequest('POST', { + code: 'return [.total, .total, , , typeof < skipped.value>, , ]', + blockNameMapping: { source: 'source-id', skipped: 'skipped-id' }, + blockData: { 'source-id': { result: '{"total":3}' } }, + workflowVariables: { + first: { name: 'Total amount', type: 'number', value: '7' }, + second: { name: 'totalamount', type: 'number', value: '99' }, + taxRate: { name: 'tax-rate', type: 'number', value: '11' }, + legacyTaxRate: { name: 'tax_rate', type: 'number', value: '12' }, + }, + }) + ) + expect(response.status).toBe(200) + expect((await response.json()).output.result).toEqual([3, 3, 7, 7, 'undefined', 11, 11]) + }) + it('keeps an exact-name/exact-value JavaScript secret out of source and returns its raw runtime value with private provenance', async () => { mockExecuteInIsolatedVM.mockResolvedValueOnce({ result: 'Test', stdout: '' }) diff --git a/apps/sim/lib/function-execution/execute-request.ts b/apps/sim/lib/function-execution/execute-request.ts index 00966573e4f..dd5b2f8e752 100644 --- a/apps/sim/lib/function-execution/execute-request.ts +++ b/apps/sim/lib/function-execution/execute-request.ts @@ -8,7 +8,6 @@ import { sha256Hex } from '@sim/security/hash' import { getErrorMessage } from '@sim/utils/errors' import { generateShortId } from '@sim/utils/id' import { toRecord } from '@sim/utils/object' -import { escapeRegExp } from '@sim/utils/string' import { NextResponse } from 'next/server' import type { ParsedFunctionExecuteBody } from '@/lib/api/contracts' import { isMothershipSandboxEnabled, isRemoteSandboxEnabled } from '@/lib/core/config/env-flags' @@ -87,6 +86,7 @@ import { } from '@/lib/execution/remote-sandbox/sandbox-paths' import type { SandboxCollectedFile, SandboxFile } from '@/lib/execution/remote-sandbox/types' import { isExecutionResourceLimitError } from '@/lib/execution/resource-errors' +import { MAX_FUNCTION_REFERENCES } from '@/lib/function-execution/limits' import type { SandboxExportedFile } from '@/lib/function-execution/output' import { planUserFileMounts, resolveUserFileMounts } from '@/lib/function-execution/sandbox-mounts' import { @@ -750,38 +750,33 @@ function scrubInternalIdentifiers(message: string, identifiers: readonly string[ function resolveWorkflowVariables( code: string, - workflowVariables: Record, - contextVariables: Record + workflowVariables: Record, + contextVariables: Record ): string { - let resolvedCode = code - - const regex = createWorkflowVariablePattern() - let match: RegExpExecArray | null - const replacements: Array<{ - match: string - index: number - variableName: string - variableValue: unknown - }> = [] - - while ((match = regex.exec(code)) !== null) { - const variableName = match[1].trim() - - const foundVariable = Object.entries(workflowVariables).find( - ([_, variable]) => normalizeName(variable.name || '') === variableName - ) - - if (!foundVariable) { - const availableVars = Object.values(workflowVariables) - .map((v) => v.name) - .filter(Boolean) + const variablesByName = new Map>() + for (const value of Object.values(workflowVariables)) { + const variable = toRecord(value) + if (typeof variable.name !== 'string') continue + const name = normalizeName(variable.name) + if (!variablesByName.has(name)) variablesByName.set(name, variable) + } + const replacements = new Map() + const boundNames = new Set() + + return code.replace(createWorkflowVariablePattern(), (_match, name: string) => { + const variableName = name.trim() + const cached = replacements.get(variableName) + if (cached !== undefined) return cached + + const variable = variablesByName.get(variableName) + if (!variable) { + const availableVars = [...variablesByName.values()].map((value) => value.name).filter(Boolean) throw new Error( `Variable "${variableName}" doesn't exist.` + (availableVars.length > 0 ? ` Available: ${availableVars.join(', ')}` : '') ) } - const variable = foundVariable[1] let variableValue: unknown = variable.value if (variable.value !== undefined && variable.value !== null) { @@ -805,24 +800,15 @@ function resolveWorkflowVariables( } } - replacements.push({ - match: match[0], - index: match.index, - variableName, - variableValue, - }) - } - - for (let i = replacements.length - 1; i >= 0; i--) { - const { match: matchStr, index, variableName, variableValue } = replacements[i] - const safeVarName = `__variable_${variableName.replace(/[^a-zA-Z0-9_]/g, '_')}` - contextVariables[safeVarName] = variableValue - resolvedCode = - resolvedCode.slice(0, index) + safeVarName + resolvedCode.slice(index + matchStr.length) - } - - return resolvedCode + // The original reverse rewrite gave the first reference precedence on binding-name collisions. + if (!boundNames.has(safeVarName)) { + contextVariables[safeVarName] = variableValue + boundNames.add(safeVarName) + } + replacements.set(variableName, safeVarName) + return safeVarName + }) } /** @@ -869,13 +855,12 @@ function resolveTagVariables( contextVariables: Record, language = 'javascript' ): string { - let resolvedCode = code const undefinedLiteral = language === 'python' ? 'None' : 'undefined' + const replacements = new Map() - const tagMatches = resolvedCode.match(TAG_PATTERN) || [] - - for (const match of tagMatches) { + return code.replace(TAG_PATTERN, (match) => { const tagName = match.slice(REFERENCE.START.length, -REFERENCE.END.length).trim() + if (replacements.has(tagName)) return replacements.get(tagName) ?? match const pathParts = tagName.split(REFERENCE.PATH_DELIMITER) const blockName = pathParts[0] const fieldPath = pathParts.slice(1) @@ -887,14 +872,15 @@ function resolveTagVariables( }) if (!result) { - continue + replacements.set(tagName, undefined) + return match } let tagValue = result.value if (tagValue === undefined) { - resolvedCode = resolvedCode.replace(new RegExp(escapeRegExp(match), 'g'), undefinedLiteral) - continue + replacements.set(tagName, undefinedLiteral) + return undefinedLiteral } if (typeof tagValue === 'string') { @@ -910,10 +896,9 @@ function resolveTagVariables( const safeVarName = `__tag_${tagName.replace(/_/g, '_1').replace(/\./g, '_0')}` contextVariables[safeVarName] = tagValue - resolvedCode = resolvedCode.replace(new RegExp(escapeRegExp(match), 'g'), safeVarName) - } - - return resolvedCode + replacements.set(tagName, safeVarName) + return safeVarName + }) } /** @@ -2281,6 +2266,23 @@ export async function executeFunctionRequest( ) includePrivateResolvedSecretNames = privateResolvedSecretNamesMetadataType !== undefined + let referenceCount = 0 + for (const _match of body.code.matchAll(TAG_PATTERN)) { + if (++referenceCount > MAX_FUNCTION_REFERENCES) { + return appendPrivateResolvedSecretNames( + NextResponse.json( + { + success: false, + error: `Function code exceeds the maximum of ${MAX_FUNCTION_REFERENCES} references`, + }, + { status: 400 } + ), + includePrivateResolvedSecretNames ? [] : null, + privateResolvedSecretNamesMetadataType + ) + } + } + const mountedWorkspaceFileProvenance = inspectMountedWorkspaceFileProvenance(req.headers, body) if (mountedWorkspaceFileProvenance.status === 'invalid') { return appendPrivateResolvedSecretNames( diff --git a/apps/sim/lib/function-execution/limits.ts b/apps/sim/lib/function-execution/limits.ts new file mode 100644 index 00000000000..ed34bf72f0c --- /dev/null +++ b/apps/sim/lib/function-execution/limits.ts @@ -0,0 +1,3 @@ +/** Bounds source scanning and the bindings created before sandbox execution. */ +export const MAX_FUNCTION_CODE_LENGTH = 1024 * 1024 +export const MAX_FUNCTION_REFERENCES = 10_000 diff --git a/apps/sim/lib/mothership/agent-cli/engines.test.ts b/apps/sim/lib/mothership/agent-cli/engines.test.ts index c033e4441c1..cab1eb118e6 100644 --- a/apps/sim/lib/mothership/agent-cli/engines.test.ts +++ b/apps/sim/lib/mothership/agent-cli/engines.test.ts @@ -136,6 +136,90 @@ const DEPS_STATE = { } describe('workflows deps', () => { + it('stops reading a wide object when its traversal budget is exhausted', async () => { + const value: Record = {} + for (let index = 0; index < 10_001; index++) { + Object.defineProperty(value, String(index), { + enumerable: true, + get() { + if (index === 10_000) throw new Error('Read beyond the traversal budget') + return '' + }, + }) + } + const state = { + ...DEPS_STATE, + blocks: { + ...DEPS_STATE.blocks, + target: { ...DEPS_STATE.blocks.target, subBlocks: { code: { value } } }, + }, + } + const result = await runEngine( + 'workflows deps', + ['wf-1', 'target'], + runtimeWith({ [STATE_PATH]: { data: state } }), + {} + ) + expect(result.exitCode).toBe(1) + expect(result.stderr).toMatch(/exceeds.*values/i) + expect(result.stdout).toBe('') + }) + + it.each([ + { reason: 'text size', value: 'x'.repeat(1024 * 1024 + 1) }, + { reason: 'reference count', value: ''.repeat(10_001) }, + { reason: 'nested value count', value: Array.from({ length: 10_001 }, () => '') }, + { + reason: 'path depth', + value: ` 'nested').join('.')}>`, + }, + ])( + 'refuses excessive $reason instead of returning an incomplete dependency report', + async ({ value }) => { + const state = { + ...DEPS_STATE, + blocks: { + ...DEPS_STATE.blocks, + target: { ...DEPS_STATE.blocks.target, subBlocks: { code: { value } } }, + }, + } + const result = await runEngine( + 'workflows deps', + ['wf-1', 'target'], + runtimeWith({ [STATE_PATH]: { data: state } }), + {} + ) + expect(result.exitCode).toBe(1) + expect(result.stderr).toMatch(/exceeds|maximum/i) + expect(result.stdout).toBe('') + } + ) + + it('groups block aliases and duplicate paths without changing first-reference order', async () => { + const state = structuredClone(DEPS_STATE) + state.blocks.target.subBlocks.code.value = + ' {{TOKEN}} {{TOKEN}}' + const result = await runEngine( + 'workflows deps', + ['wf-1', 'target'], + runtimeWith({ [STATE_PATH]: { data: state } }), + {} + ) + const report = JSON.parse(result.stdout) + expect(report.references).toEqual([ + { token: 'missing.value', kind: 'unknown' }, + { + token: 'fetchrows.result', + kind: 'block', + blockId: 'fetch', + blockName: 'Fetch rows', + paths: ['result', 'result.id'], + }, + ]) + expect(report.env).toEqual(['TOKEN']) + expect(report.mock['Fetch rows']).toEqual({ result: { id: null } }) + }) + it('builds indexed mocks that round-trip through the actual reference navigator', async () => { const state = structuredClone(DEPS_STATE) state.blocks.target.subBlocks.code.value = diff --git a/apps/sim/lib/mothership/agent-cli/engines/deps.ts b/apps/sim/lib/mothership/agent-cli/engines/deps.ts index 3889f2bb004..ca8e108a9d3 100644 --- a/apps/sim/lib/mothership/agent-cli/engines/deps.ts +++ b/apps/sim/lib/mothership/agent-cli/engines/deps.ts @@ -1,12 +1,9 @@ +import { isRecordLike } from '@sim/utils/object' import { fetchWorkflowState } from '@/lib/mothership/agent-cli/engines/workflow-state' import { type AgentCliEngine, agentCliFail, agentCliOk } from '@/lib/mothership/agent-cli/types' import { TriggerUtils } from '@/lib/workflows/triggers/triggers' import { normalizeName, SPECIAL_REFERENCE_PREFIXES } from '@/executor/constants' -import { - collectStringLeaves, - createEnvVarPattern, - createReferencePattern, -} from '@/executor/utils/reference-validation' +import { createEnvVarPattern, createReferencePattern } from '@/executor/utils/reference-validation' import { splitLeadingBracketPath } from '@/executor/variables/resolvers/reference' /** @@ -34,6 +31,10 @@ const CHILD_RETURNS_NOTE = "A child workflow's result is its actual final output. A Response block returns {data, status, headers}, with fields at result.data.; otherwise use the final block's output shape." const MOCK_NOTE = 'variableInputs: fill placeholders with representative upstream outputs' const MAX_MOCK_ARRAY_LENGTH = 128 +const MAX_DEPENDENCY_INPUT_CHARACTERS = 1024 * 1024 +const MAX_DEPENDENCY_INPUT_NODES = 10_000 +const MAX_DEPENDENCY_REFERENCES = 10_000 +const MAX_DEPENDENCY_PATH_DEPTH = 128 interface PathNode { children: Map @@ -81,6 +82,9 @@ function skeletonFromPaths(paths: readonly string[]): MockShape { const { property, pathParts } = splitLeadingBracketPath(part) return [property, ...pathParts] }) + if (segments.length > MAX_DEPENDENCY_PATH_DEPTH) { + throw new Error(`Dependency path exceeds the maximum depth of ${MAX_DEPENDENCY_PATH_DEPTH}`) + } let cursor = root for (const segment of segments) { const next = cursor.children.get(segment) ?? { children: new Map() } @@ -172,47 +176,98 @@ export const workflowDepsCommand: AgentCliEngine = { } const leaves: string[] = [] - collectStringLeaves(block.subBlocks ?? block, leaves) + const pending: unknown[] = [block.subBlocks ?? block] + let inputNodes = 0 + let inputCharacters = 0 + while (pending.length > 0) { + const value = pending.pop() + inputNodes++ + if (typeof value === 'string') { + inputCharacters += value.length + if (inputCharacters > MAX_DEPENDENCY_INPUT_CHARACTERS) { + return agentCliFail( + `Dependency input exceeds the maximum of ${MAX_DEPENDENCY_INPUT_CHARACTERS} characters` + ) + } + leaves.push(value) + } else if (Array.isArray(value) || isRecordLike(value)) { + const children: unknown[] = Array.isArray(value) ? value : [] + if (inputNodes + pending.length + children.length > MAX_DEPENDENCY_INPUT_NODES) { + return agentCliFail( + `Dependency input exceeds the maximum of ${MAX_DEPENDENCY_INPUT_NODES} values` + ) + } + if (!Array.isArray(value)) { + for (const key in value) { + if (!Object.hasOwn(value, key)) continue + if (inputNodes + pending.length + children.length >= MAX_DEPENDENCY_INPUT_NODES) { + return agentCliFail( + `Dependency input exceeds the maximum of ${MAX_DEPENDENCY_INPUT_NODES} values` + ) + } + children.push(value[key]) + } + } + for (let index = children.length - 1; index >= 0; index--) pending.push(children[index]) + } + } - const byToken = new Map() + const seenTokens = new Set() + const deps: DepView[] = [] + const byBlock = new Map }>() const envs = new Set() + let referenceCount = 0 for (const leaf of leaves) { for (const match of leaf.matchAll(TEMPLATE_REF)) { + if (++referenceCount > MAX_DEPENDENCY_REFERENCES) { + return agentCliFail( + `Dependency input exceeds the maximum of ${MAX_DEPENDENCY_REFERENCES} references` + ) + } const token = match[1] - if (!token || byToken.has(token)) continue + if (!token || seenTokens.has(token)) continue + seenTokens.add(token) const [head = '', ...pathParts] = token.split('.') const path = pathParts.join('.') const special = (SPECIAL_REFERENCE_PREFIXES as readonly string[]).includes(head) if (special) { - byToken.set(token, { token, kind: head as 'loop' | 'parallel' | 'variable' }) + deps.push({ token, kind: head as 'loop' | 'parallel' | 'variable' }) continue } const refBlockId = Object.hasOwn(blocks, head) ? head : nameToId.get(normalizeName(head)) if (refBlockId && refBlockId !== blockId) { - const existing = [...byToken.values()].find((d) => d.blockId === refBlockId) + const existing = byBlock.get(refBlockId) if (existing) { - if (path && !existing.paths?.includes(path)) existing.paths?.push(path) - byToken.set(token, existing) + if (path && !existing.paths.has(path)) { + existing.paths.add(path) + existing.dependency.paths?.push(path) + } } else { - byToken.set(token, { + const dependency: DepView = { token, kind: 'block', blockId: refBlockId, blockName: idToName.get(refBlockId), paths: path ? [path] : [], - }) + } + byBlock.set(refBlockId, { dependency, paths: new Set(dependency.paths) }) + deps.push(dependency) } } else if (!refBlockId) { - byToken.set(token, { token, kind: 'unknown' }) + deps.push({ token, kind: 'unknown' }) } } for (const match of leaf.matchAll(ENV_REF)) { + if (++referenceCount > MAX_DEPENDENCY_REFERENCES) { + return agentCliFail( + `Dependency input exceeds the maximum of ${MAX_DEPENDENCY_REFERENCES} references` + ) + } const key = match[1]?.trim() if (key) envs.add(key) } } - const deps = [...new Set(byToken.values())] const blockDeps = deps.filter((d) => d.kind === 'block') const predecessors = collectPredecessors(state, blocks, blockId, idToName) diff --git a/apps/sim/lib/mothership/chat/context-limits.ts b/apps/sim/lib/mothership/chat/context-limits.ts new file mode 100644 index 00000000000..102e9883027 --- /dev/null +++ b/apps/sim/lib/mothership/chat/context-limits.ts @@ -0,0 +1,4 @@ +/** Shared admission limits for interactive and scheduled context resolution. */ +export const MAX_CHAT_CONTEXTS = 100 +export const MAX_CHAT_CONTEXT_LABEL_LENGTH = 1000 +export const MAX_CHAT_MESSAGE_LENGTH = 1024 * 1024 diff --git a/apps/sim/lib/mothership/chat/post.ts b/apps/sim/lib/mothership/chat/post.ts index f7d22257fdb..d37992b01a7 100644 --- a/apps/sim/lib/mothership/chat/post.ts +++ b/apps/sim/lib/mothership/chat/post.ts @@ -37,6 +37,11 @@ import { prepareOrganizationChatAttachments, } from '@/lib/mothership/chat/assistant-images' import { buildOnComplete, buildOnError } from '@/lib/mothership/chat/completion' +import { + MAX_CHAT_CONTEXT_LABEL_LENGTH, + MAX_CHAT_CONTEXTS, + MAX_CHAT_MESSAGE_LENGTH, +} from '@/lib/mothership/chat/context-limits' import { DESKTOP_TERMINAL_HINT_ID_MAX_LENGTH, DESKTOP_TERMINAL_HINT_TEXT_MAX_LENGTH, @@ -218,7 +223,7 @@ const ChatContextSchema = z 'terminal_tab', 'workspace', ]), - label: z.string(), + label: z.string().max(MAX_CHAT_CONTEXT_LABEL_LENGTH), chatId: z.string().optional(), workflowId: z.string().optional(), knowledgeId: z.string().optional(), @@ -271,7 +276,7 @@ const ChatContextSchema = z const ChatMessageSchema = z .object({ - message: z.string(), + message: z.string().max(MAX_CHAT_MESSAGE_LENGTH), /* Bounded because it becomes part of a Postgres key in `chatSendIdempotency`; a client-supplied id longer than the btree entry limit would throw there. A generated id is 36 chars. */ @@ -294,7 +299,7 @@ const ChatMessageSchema = z .preprocess(dropUnaddressableAttachments, z.array(ResourceAttachmentSchema)) .optional(), provider: z.string().optional(), - contexts: z.array(ChatContextSchema).optional(), + contexts: z.array(ChatContextSchema).max(MAX_CHAT_CONTEXTS).optional(), commands: z.array(z.string()).optional(), userTimezone: z.string().optional(), effort: z.enum(['none', 'low', 'medium', 'high', 'xhigh', 'max']).optional(), diff --git a/apps/sim/lib/mothership/chat/process-contents.test.ts b/apps/sim/lib/mothership/chat/process-contents.test.ts index 4d2d69b5247..6694ec63841 100644 --- a/apps/sim/lib/mothership/chat/process-contents.test.ts +++ b/apps/sim/lib/mothership/chat/process-contents.test.ts @@ -452,6 +452,38 @@ describe('processContextsServer - skill contexts', () => { }) describe('processContextsServer - docs contexts', () => { + it.each([ + { + message: 'Explain loops', + contexts: Array.from({ length: 101 }, () => ({ kind: 'docs' as const, label: 'Docs' })), + }, + { message: 'x'.repeat(1024 * 1024 + 1), contexts: [{ kind: 'docs' as const, label: 'Docs' }] }, + { message: 'Explain loops', contexts: [{ kind: 'docs' as const, label: 'x'.repeat(1001) }] }, + ])('rejects context resolution beyond its input budget', async ({ message, contexts }) => { + await expect(processContextsServer(contexts, 'reader', message, 'workspace-1')).rejects.toThrow( + /exceeds|maximum/i + ) + }) + + it('preserves multiword block mentions while removing resource labels and unknown mentions', async () => { + searchDocsExecute.mockImplementationOnce(async ({ query }: { query: string }) => ({ + results: [], + note: query, + })) + const result = await processContextsServer( + [ + { kind: 'docs', label: 'Product docs' }, + { kind: 'blocks', label: 'Read [rows]' }, + { kind: 'integration', label: 'My service' }, + ], + 'reader', + '@Product docs compare @Read [rows] with @My service and @unknown', + 'workspace-1' + ) + const docs = result.find((context) => context.type === 'docs') + expect(JSON.parse(docs!.content).note).toBe('compare Read [rows] with and') + }) + it('routes @Docs to an unscoped search_docs query', async () => { const resolvedSecretTraceRegistry = new ResolvedSecretTraceRegistry() const results = [ diff --git a/apps/sim/lib/mothership/chat/process-contents.ts b/apps/sim/lib/mothership/chat/process-contents.ts index 56fd8393915..946793fc09b 100644 --- a/apps/sim/lib/mothership/chat/process-contents.ts +++ b/apps/sim/lib/mothership/chat/process-contents.ts @@ -31,6 +31,11 @@ import { createCopilotChatTablePrincipal } from '@/lib/mothership/auth/table-del import { getBlockVisibilityForCopilot } from '@/lib/mothership/block-visibility' import { readWorkspaceContext } from '@/lib/mothership/chat/application/workspace-context' import { WORKSPACE_TARGET_AUDIENCE } from '@/lib/mothership/chat/application/workspace-target' +import { + MAX_CHAT_CONTEXT_LABEL_LENGTH, + MAX_CHAT_CONTEXTS, + MAX_CHAT_MESSAGE_LENGTH, +} from '@/lib/mothership/chat/context-limits' import { isWorkspaceOwnedContext, type WorkspaceOwnedContext, @@ -140,6 +145,19 @@ export async function processContextsServer( organizationId?: string ): Promise { if (!Array.isArray(contexts) || contexts.length === 0) return [] + if (contexts.length > MAX_CHAT_CONTEXTS) { + throw new Error(`Context count exceeds the maximum of ${MAX_CHAT_CONTEXTS}`) + } + if ((userMessage?.length ?? 0) > MAX_CHAT_MESSAGE_LENGTH) { + throw new Error(`Message exceeds the maximum of ${MAX_CHAT_MESSAGE_LENGTH} characters`) + } + if (contexts.some((context) => context.label.length > MAX_CHAT_CONTEXT_LABEL_LENGTH)) { + throw new Error( + `Context label exceeds the maximum of ${MAX_CHAT_CONTEXT_LABEL_LENGTH} characters` + ) + } + const docsMessage = userMessage?.trim() ?? '' + const docsQueries = new Map() /** * An organization chat has no workspace of its own, so each workspace-owned @@ -342,8 +360,13 @@ export async function processContextsServer( const { searchDocsServerTool } = await import( '@/lib/mothership/tools/server/docs/search-docs' ) - const rawQuery = (userMessage || '').trim() || ctx.label || 'Sim documentation' - const query = sanitizeMessageForDocs(rawQuery, contexts) || ctx.label || 'Sim documentation' + const rawQuery = docsMessage || ctx.label || 'Sim documentation' + let sanitizedQuery = docsQueries.get(rawQuery) + if (sanitizedQuery === undefined) { + sanitizedQuery = sanitizeMessageForDocs(rawQuery, contexts) + docsQueries.set(rawQuery, sanitizedQuery) + } + const query = sanitizedQuery || ctx.label || 'Sim documentation' const res = await searchDocsServerTool.execute( { query }, { @@ -469,22 +492,20 @@ function sanitizeMessageForDocs(rawMessage: string, contexts: ChatContext[] | un let result = rawMessage - // 1) Remove all non-block mentions entirely - for (const label of nonBlockLabels) { - const pattern = new RegExp(`(^|\\s)@${escapeRegExp(label)}(?!\\S)`, 'g') + if (nonBlockLabels.size > 0) { + const labels = [...nonBlockLabels].map(escapeRegExp).join('|') + const pattern = new RegExp(`(^|\\s)@(?:${labels})(?!\\S)`, 'g') result = result.replace(pattern, ' ') } - // 2) For block mentions, strip the '@' but keep the block name - for (const label of blockLabels) { - const pattern = new RegExp(`@${escapeRegExp(label)}(?!\\S)`, 'g') - result = result.replace(pattern, label) + if (blockLabels.size > 0) { + const labels = [...blockLabels].map(escapeRegExp).join('|') + const pattern = new RegExp(`@(${labels})(?!\\S)`, 'g') + result = result.replace(pattern, (_match, label: string) => label) } - // 3) Remove any remaining @mentions (unknown or not in contexts) result = result.replace(/(^|\s)@([^\s]+)/g, ' ') - // Normalize whitespace result = result.replace(/\s{2,}/g, ' ').trim() return result }