Revert "always default to adding risk prediction"

This reverts commit 160a6a2.

Author: malhotra5

openhands-sdk/openhands/sdk/agent/agent.py

@@ -99,7 +99,10 @@ def init_state(
                 source="agent",
                 system_prompt=TextContent(text=self.system_message),
                 # Always include security_risk field in tools
-                tools=[t.to_openai_tool() for t in self.tools_map.values()],
+                tools=[
+                    t.to_openai_tool(add_security_risk_prediction=True)
+                    for t in self.tools_map.values()
+                ],
             )
             on_event(event)
 
@@ -175,13 +178,15 @@ def step(
                     tools=list(self.tools_map.values()),
                     include=None,
                     store=False,
+                    add_security_risk_prediction=True,
                     extra_body=self.llm.litellm_extra_body,
                 )
             else:
                 llm_response = self.llm.completion(
                     messages=_messages,
                     tools=list(self.tools_map.values()),
                     extra_body=self.llm.litellm_extra_body,
+                    add_security_risk_prediction=True,
                 )
         except FunctionCallValidationError as e:
             logger.warning(f"LLM generated malformed function call: {e}")

openhands-sdk/openhands/sdk/llm/llm.py

@@ -433,6 +433,7 @@ def completion(
         messages: list[Message],
         tools: Sequence[ToolDefinition] | None = None,
         _return_metrics: bool = False,
+        add_security_risk_prediction: bool = False,
         **kwargs,
     ) -> LLMResponse:
         """Generate a completion from the language model.
@@ -466,7 +467,12 @@ def completion(
         # Convert Tool objects to ChatCompletionToolParam once here
         cc_tools: list[ChatCompletionToolParam] = []
         if tools:
-            cc_tools = [t.to_openai_tool() for t in tools]
+            cc_tools = [
+                t.to_openai_tool(
+                    add_security_risk_prediction=add_security_risk_prediction
+                )
+                for t in tools
+            ]
 
         use_mock_tools = self.should_mock_tool_calls(cc_tools)
         if use_mock_tools:
@@ -566,6 +572,7 @@ def responses(
         include: list[str] | None = None,
         store: bool | None = None,
         _return_metrics: bool = False,
+        add_security_risk_prediction: bool = False,
         **kwargs,
     ) -> LLMResponse:
         """Alternative invocation path using OpenAI Responses API via LiteLLM.
@@ -582,7 +589,16 @@ def responses(
 
         # Convert Tool objects to Responses ToolParam
         # (Responses path always supports function tools)
-        resp_tools = [t.to_responses_tool() for t in tools] if tools else None
+        resp_tools = (
+            [
+                t.to_responses_tool(
+                    add_security_risk_prediction=add_security_risk_prediction
+                )
+                for t in tools
+            ]
+            if tools
+            else None
+        )
 
         # Normalize/override Responses kwargs consistently
         call_kwargs = select_responses_options(

openhands-sdk/openhands/sdk/llm/router/base.py

@@ -51,6 +51,7 @@ def completion(
         messages: list[Message],
         tools: Sequence[ToolDefinition] | None = None,
         return_metrics: bool = False,
+        add_security_risk_prediction: bool = False,
         **kwargs,
     ) -> LLMResponse:
         """
@@ -68,6 +69,7 @@ def completion(
             messages=messages,
             tools=tools,
             _return_metrics=return_metrics,
+            add_security_risk_prediction=add_security_risk_prediction,
             **kwargs,
         )
 

openhands-sdk/openhands/sdk/mcp/tool.py

@@ -242,6 +242,7 @@ def to_mcp_tool(
 
     def to_openai_tool(
         self,
+        add_security_risk_prediction: bool = False,
         action_type: type[Schema] | None = None,
     ) -> ChatCompletionToolParam:
         """Convert a Tool to an OpenAI tool.
@@ -250,6 +251,12 @@ def to_openai_tool(
         from the MCP tool input schema, and pass it to the parent method.
         It will use the .model_fields from this pydantic model to
         generate the OpenAI-compatible tool schema.
+
+        Args:
+            add_security_risk_prediction: Whether to add a `security_risk` field
+                to the action schema for LLM to predict. This is useful for
+                tools that may have safety risks, so the LLM can reason about
+                the risk level before calling the tool.
         """
         if action_type is not None:
             raise ValueError(
@@ -259,5 +266,6 @@ def to_openai_tool(
         assert self.name == self.mcp_tool.name
         mcp_action_type = _create_mcp_action_type(self.mcp_tool)
         return super().to_openai_tool(
+            add_security_risk_prediction=add_security_risk_prediction,
             action_type=mcp_action_type,
         )

openhands-sdk/openhands/sdk/tool/tool.py

@@ -360,42 +360,72 @@ def to_mcp_tool(
 
     def _get_tool_schema(
         self,
+        add_security_risk_prediction: bool = False,
         action_type: type[Schema] | None = None,
     ) -> dict[str, Any]:
         action_type = action_type or self.action_type
-        action_type_with_risk = _create_action_type_with_risk(action_type)
-        schema = action_type_with_risk.to_mcp_schema()
+
+        if add_security_risk_prediction:
+            # Always include security_risk field when prediction is enabled
+            # This ensures consistent tool schemas regardless of tool type
+            # (including read-only tools)
+            action_type_with_risk = _create_action_type_with_risk(action_type)
+            schema = action_type_with_risk.to_mcp_schema()
+        else:
+            schema = action_type.to_mcp_schema()
+
         return schema
 
     def to_openai_tool(
         self,
+        add_security_risk_prediction: bool = False,
         action_type: type[Schema] | None = None,
     ) -> ChatCompletionToolParam:
-        """Convert a Tool to an OpenAI tool."""
+        """Convert a Tool to an OpenAI tool.
+
+        Args:
+            add_security_risk_prediction: Whether to include the `security_risk`
+                field in the tool schema. When enabled, the field is included
+                for all tool types (including read-only tools).
+            action_type: Optionally override the action_type to use for the schema.
+                This is useful for MCPTool to use a dynamically created action type
+                based on the tool's input schema.
+        """
         return ChatCompletionToolParam(
             type="function",
             function=ChatCompletionToolParamFunctionChunk(
                 name=self.name,
                 description=self.description,
-                parameters=self._get_tool_schema(action_type),
+                parameters=self._get_tool_schema(
+                    add_security_risk_prediction, action_type
+                ),
             ),
         )
 
     def to_responses_tool(
         self,
+        add_security_risk_prediction: bool = False,
         action_type: type[Schema] | None = None,
     ) -> FunctionToolParam:
         """Convert a Tool to a Responses API function tool (LiteLLM typed).
 
         For Responses API, function tools expect top-level keys:
         { "type": "function", "name": ..., "description": ..., "parameters": ... }
+
+        Args:
+            add_security_risk_prediction: Whether to include the `security_risk`
+                field in the tool schema. When enabled, the field is included
+                for all tool types (including read-only tools).
+            action_type: Optionally override the action_type to use for the schema.
         """
 
         return {
             "type": "function",
             "name": self.name,
             "description": self.description,
-            "parameters": self._get_tool_schema(action_type),
+            "parameters": self._get_tool_schema(
+                add_security_risk_prediction, action_type
+            ),
             "strict": False,
         }
 

tests/cross/test_remote_conversation_live_server.py

@@ -147,6 +147,7 @@ def fake_completion(
         messages,
         tools,
         return_metrics=False,
+        add_security_risk_prediction=False,
         **kwargs,
     ):  # type: ignore[no-untyped-def]
         from openhands.sdk.llm.llm_response import LLMResponse
@@ -447,6 +448,7 @@ def fake_completion_with_cost(
         messages,
         tools,
         return_metrics=False,
+        add_security_risk_prediction=False,
         **kwargs,
     ):  # type: ignore[no-untyped-def]
         from openhands.sdk.llm.llm_response import LLMResponse

tests/sdk/llm/test_llm_completion.py

@@ -349,7 +349,7 @@ def test_llm_completion_non_function_call_mode(mock_completion):
     tools = list(_MockTool.create())
 
     # Verify that tools should be mocked (non-function call path)
-    cc_tools = [t.to_openai_tool() for t in tools]
+    cc_tools = [t.to_openai_tool(add_security_risk_prediction=False) for t in tools]
     assert llm.should_mock_tool_calls(cc_tools)
 
     # Call completion - this should go through the prompt-based tool calling path

tests/sdk/mcp/test_mcp_security_risk.py

@@ -66,7 +66,7 @@ def test_mcp_tool_to_openai_with_security_risk():
     tool = tools[0]
 
     # Generate OpenAI tool schema WITH security risk prediction
-    openai_tool = tool.to_openai_tool()
+    openai_tool = tool.to_openai_tool(add_security_risk_prediction=True)
 
     function_params = openai_tool["function"]["parameters"]  # type: ignore[typeddict-item]
     properties = function_params["properties"]

tests/sdk/tool/test_to_responses_tool_security.py

@@ -48,7 +48,7 @@ def test_to_responses_tool_security_gating():
         observation_type=None,
         annotations=ToolAnnotations(readOnlyHint=True),
     )
-    t = readonly.to_responses_tool()
+    t = readonly.to_responses_tool(add_security_risk_prediction=True)
     params = t["parameters"]
     assert isinstance(params, dict)
     props = params.get("properties") or {}
@@ -62,7 +62,7 @@ def test_to_responses_tool_security_gating():
         observation_type=None,
         annotations=ToolAnnotations(readOnlyHint=False),
     )
-    t2 = writable.to_responses_tool()
+    t2 = writable.to_responses_tool(add_security_risk_prediction=True)
     params2 = t2["parameters"]
     assert isinstance(params2, dict)
     props2 = params2.get("properties") or {}
@@ -76,7 +76,7 @@ def test_to_responses_tool_security_gating():
         observation_type=None,
         annotations=None,
     )
-    t3 = noflag.to_responses_tool()
+    t3 = noflag.to_responses_tool(add_security_risk_prediction=False)
     params3 = t3["parameters"]
     assert isinstance(params3, dict)
     props3 = params3.get("properties") or {}

tests/sdk/tool/test_tool_definition.py

@@ -579,7 +579,9 @@ def test_security_risk_added_for_all_tools_when_enabled(self):
         )
 
         # Test read-only tool - security_risk should be added when enabled
-        readonly_openai_tool = readonly_tool.to_openai_tool()
+        readonly_openai_tool = readonly_tool.to_openai_tool(
+            add_security_risk_prediction=True
+        )
         readonly_function = readonly_openai_tool["function"]
         assert "parameters" in readonly_function
         readonly_params = readonly_function["parameters"]
@@ -588,27 +590,35 @@ def test_security_risk_added_for_all_tools_when_enabled(self):
         )  # Included for read-only tools too
 
         # Test writable tool - security_risk SHOULD be added
-        writable_openai_tool = writable_tool.to_openai_tool()
+        writable_openai_tool = writable_tool.to_openai_tool(
+            add_security_risk_prediction=True
+        )
         writable_function = writable_openai_tool["function"]
         assert "parameters" in writable_function
         writable_params = writable_function["parameters"]
         assert "security_risk" in writable_params["properties"]
 
         # Test tool with no annotations - security_risk SHOULD be added
-        no_annotations_openai_tool = no_annotations_tool.to_openai_tool()
+        no_annotations_openai_tool = no_annotations_tool.to_openai_tool(
+            add_security_risk_prediction=True
+        )
         no_annotations_function = no_annotations_openai_tool["function"]
         assert "parameters" in no_annotations_function
         no_annotations_params = no_annotations_function["parameters"]
         assert "security_risk" in no_annotations_params["properties"]
 
         # Test that when add_security_risk_prediction=False, no security_risk is added
-        readonly_no_risk = readonly_tool.to_openai_tool()
+        readonly_no_risk = readonly_tool.to_openai_tool(
+            add_security_risk_prediction=False
+        )
         readonly_no_risk_function = readonly_no_risk["function"]
         assert "parameters" in readonly_no_risk_function
         readonly_no_risk_params = readonly_no_risk_function["parameters"]
         assert "security_risk" not in readonly_no_risk_params["properties"]
 
-        writable_no_risk = writable_tool.to_openai_tool()
+        writable_no_risk = writable_tool.to_openai_tool(
+            add_security_risk_prediction=False
+        )
         writable_no_risk_function = writable_no_risk["function"]
         assert "parameters" in writable_no_risk_function
         writable_no_risk_params = writable_no_risk_function["parameters"]
@@ -633,7 +643,7 @@ def test_security_risk_is_required_field_in_schema(self):
             observation_type=ToolMockObservation,
         )
 
-        openai_tool = tool.to_openai_tool()
+        openai_tool = tool.to_openai_tool(add_security_risk_prediction=True)
         function_chunk = openai_tool["function"]
         assert "parameters" in function_chunk
         function_params = function_chunk["parameters"]
@@ -657,7 +667,9 @@ def test_security_risk_is_required_field_in_schema(self):
             annotations=writable_annotations,
         )
 
-        writable_openai_tool = writable_tool.to_openai_tool()
+        writable_openai_tool = writable_tool.to_openai_tool(
+            add_security_risk_prediction=True
+        )
         writable_function_chunk = writable_openai_tool["function"]
         assert "parameters" in writable_function_chunk
         writable_function_params = writable_function_chunk["parameters"]