From 9ac7e9b5d2268fdf8fab426a5c867a59515ca5a1 Mon Sep 17 00:00:00 2001 From: Kevin Lin Date: Tue, 3 Dec 2024 19:37:17 -0800 Subject: [PATCH] fix: escape strings in sandbox (#2156) Co-authored-by: Sarah Wooders --- letta/services/tool_execution_sandbox.py | 4 +--- tests/integration_test_tool_execution_sandbox.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/letta/services/tool_execution_sandbox.py b/letta/services/tool_execution_sandbox.py index 14068715..3cce7f25 100644 --- a/letta/services/tool_execution_sandbox.py +++ b/letta/services/tool_execution_sandbox.py @@ -287,13 +287,12 @@ class ToolExecutionSandbox: f"{self.LOCAL_SANDBOX_RESULT_VAR_NAME} = base64.b64encode(pickle.dumps({self.LOCAL_SANDBOX_RESULT_VAR_NAME})).decode('utf-8')\n" ) code += f"{self.LOCAL_SANDBOX_RESULT_VAR_NAME}\n" - return code def _convert_param_to_value(self, param_type: str, raw_value: str) -> str: if param_type == "string": - value = '"' + raw_value + '"' + value = "pickle.loads(" + str(pickle.dumps(raw_value)) + ")" elif param_type == "integer" or param_type == "boolean" or param_type == "number": value = raw_value @@ -306,7 +305,6 @@ class ToolExecutionSandbox: else: raise TypeError(f"Unsupported type: {param_type}, raw_value={raw_value}") - return str(value) def initialize_param(self, name: str, raw_value: str) -> str: diff --git a/tests/integration_test_tool_execution_sandbox.py b/tests/integration_test_tool_execution_sandbox.py index 1c5dd05f..da51f8ec 100644 --- a/tests/integration_test_tool_execution_sandbox.py +++ b/tests/integration_test_tool_execution_sandbox.py @@ -388,6 +388,18 @@ def test_e2b_sandbox_core_memory_replace(check_e2b_key_is_set, core_memory_repla assert result.func_return is None +@pytest.mark.e2b_sandbox +def test_e2b_sandbox_escape_strings_in_args(check_e2b_key_is_set, core_memory_replace_tool, test_user, agent_state): + new_name = "Matt" + args = {"label": "human", "old_content": "Chad", "new_content": new_name + "\n"} + sandbox = ToolExecutionSandbox(core_memory_replace_tool.name, args, user_id=test_user.id) + + # run the sandbox + result = sandbox.run(agent_state=agent_state) + assert new_name in result.agent_state.memory.get_block("human").value + assert result.func_return is None + + @pytest.mark.e2b_sandbox def test_e2b_sandbox_core_memory_replace_errors(check_e2b_key_is_set, core_memory_replace_tool, test_user, agent_state): nonexistent_name = "Alexander Wang"