Merge pull request #2383 from onmete/dont-leak-token-to-llm

Don't capture the literal error in case of wrong tool args

Author: openshift-merge-bot[bot]

ols/src/tools/tools.py

@@ -49,11 +49,8 @@ def execute_oc_tool_calls(
             logger.error(tool_output)
         else:
             try:
-                # inject token into tool args and immediately remove it
-                # to avoid leaking
-                tool_args["token"] = token
-                tool_output = tool.invoke(tool_args)
-                del tool_args["token"]
+                # create a new dict with the tool args and the token
+                tool_output = tool.invoke({**tool_args, "token": token})
             except ValidationError:
                 tool_output = (
                     f"Error executing {tool_name}: tool arguments are in wrong format"
@@ -64,11 +61,6 @@ def execute_oc_tool_calls(
             except Exception as e:
                 tool_output = f"Error executing {tool_name}: {e}"
                 logger.exception(tool_output)
-            finally:
-                # remove token from tool args if it was not removed
-                # in the try block
-                if "token" in tool_args:
-                    del tool_args["token"]
 
         logger.debug(
             "Tool: %s | Args: %s | Output: %s", tool_name, tool_args, tool_output

tests/unit/tools/test_tools.py

@@ -69,3 +69,30 @@ def tool1(some_arg: str):
         in caplog.text
     )
     assert "fake-token" not in caplog.text
+
+
+def test_execute_oc_tool_calls_not_leaks_token_into_output(caplog):
+    """Test execute_oc_tool_calls does not leak token into output."""
+    caplog.set_level(10)  # set debug level
+
+    @tool
+    def tool1(some_args: list):
+        """Tool 1."""
+        return "bla"
+
+    tools_map = {"tool1": tool1}
+
+    # missing args
+    tool_calls = [{"id": 1, "name": "tool1"}]
+    tool_messages = execute_oc_tool_calls(tools_map, tool_calls, "fake-token")
+    assert len(tool_messages) == 1
+    assert "fake-token" not in tool_messages[0].content
+
+    # unknown args
+    tool_calls = [{"id": 1, "name": "tool1", "args": {"unknown_args": "blo"}}]
+    tool_messages = execute_oc_tool_calls(tools_map, tool_calls, "fake-token")
+    assert len(tool_messages) == 1
+    assert "fake-token" not in tool_messages[0].content
+
+    # ensure the token is also not in the logs
+    assert "fake-token" not in caplog.text