Refactor: Always include risk fields (#1052)

Co-authored-by: openhands <openhands@all-hands.dev>
Co-authored-by: Xingyao Wang <xingyao@all-hands.dev>

Author: 3 people

openhands-agent-server/openhands/agent_server/conversation_router.py

@@ -16,6 +16,7 @@
     GenerateTitleResponse,
     SendMessageRequest,
     SetConfirmationPolicyRequest,
+    SetSecurityAnalyzerRequest,
     StartConversationRequest,
     Success,
     UpdateConversationRequest,
@@ -237,6 +238,23 @@ async def set_conversation_confirmation_policy(
     return Success()
 
 
+@conversation_router.post(
+    "/{conversation_id}/security_analyzer",
+    responses={404: {"description": "Item not found"}},
+)
+async def set_conversation_security_analyzer(
+    conversation_id: UUID,
+    request: SetSecurityAnalyzerRequest,
+    conversation_service: ConversationService = Depends(get_conversation_service),
+) -> Success:
+    """Set the security analyzer for a conversation."""
+    event_service = await conversation_service.get_event_service(conversation_id)
+    if event_service is None:
+        raise HTTPException(status.HTTP_404_NOT_FOUND)
+    await event_service.set_security_analyzer(request.security_analyzer)
+    return Success()
+
+
 @conversation_router.patch(
     "/{conversation_id}", responses={404: {"description": "Item not found"}}
 )

openhands-agent-server/openhands/agent_server/event_service.py

@@ -20,6 +20,7 @@
     ConversationState,
 )
 from openhands.sdk.event.conversation_state import ConversationStateUpdateEvent
+from openhands.sdk.security.analyzer import SecurityAnalyzerBase
 from openhands.sdk.security.confirmation_policy import ConfirmationPolicyBase
 from openhands.sdk.utils.async_utils import AsyncCallbackWrapper
 from openhands.sdk.utils.cipher import Cipher
@@ -303,6 +304,17 @@ async def set_confirmation_policy(self, policy: ConfirmationPolicyBase):
             None, self._conversation.set_confirmation_policy, policy
         )
 
+    async def set_security_analyzer(
+        self, security_analyzer: SecurityAnalyzerBase | None
+    ):
+        """Set the security analyzer for the conversation."""
+        if not self._conversation:
+            raise ValueError("inactive_service")
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(
+            None, self._conversation.set_security_analyzer, security_analyzer
+        )
+
     async def close(self):
         await self._pub_sub.close()
         if self._conversation:

openhands-agent-server/openhands/agent_server/models.py

@@ -14,6 +14,7 @@
     ConversationState,
 )
 from openhands.sdk.llm.utils.metrics import MetricsSnapshot
+from openhands.sdk.security.analyzer import SecurityAnalyzerBase
 from openhands.sdk.security.confirmation_policy import (
     ConfirmationPolicyBase,
     NeverConfirm,
@@ -165,6 +166,14 @@ class SetConfirmationPolicyRequest(BaseModel):
     policy: ConfirmationPolicyBase = Field(description="The confirmation policy to set")
 
 
+class SetSecurityAnalyzerRequest(BaseModel):
+    "Payload to set security analyzer for a conversation"
+
+    security_analyzer: SecurityAnalyzerBase | None = Field(
+        description="The security analyzer to set"
+    )
+
+
 class UpdateConversationRequest(BaseModel):
     """Payload to update conversation metadata."""
 

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

@@ -1,7 +1,8 @@
 import json
 
-from pydantic import ValidationError
+from pydantic import ValidationError, model_validator
 
+import openhands.sdk.security.analyzer as analyzer
 import openhands.sdk.security.risk as risk
 from openhands.sdk.agent.base import AgentBase
 from openhands.sdk.agent.utils import fix_malformed_tool_arguments
@@ -41,7 +42,6 @@
     should_enable_observability,
 )
 from openhands.sdk.observability.utils import extract_action_name
-from openhands.sdk.security.confirmation_policy import NeverConfirm
 from openhands.sdk.security.llm_analyzer import LLMSecurityAnalyzer
 from openhands.sdk.tool import (
     Action,
@@ -72,9 +72,20 @@ class Agent(AgentBase):
         >>> agent = Agent(llm=llm, tools=tools)
     """
 
-    @property
-    def _add_security_risk_prediction(self) -> bool:
-        return isinstance(self.security_analyzer, LLMSecurityAnalyzer)
+    @model_validator(mode="before")
+    @classmethod
+    def _add_security_prompt_as_default(cls, data):
+        """Ensure llm_security_analyzer=True is always set before initialization."""
+        if not isinstance(data, dict):
+            return data
+
+        kwargs = data.get("system_prompt_kwargs") or {}
+        if not isinstance(kwargs, dict):
+            kwargs = {}
+
+        kwargs.setdefault("llm_security_analyzer", True)
+        data["system_prompt_kwargs"] = kwargs
+        return data
 
     def init_state(
         self,
@@ -85,18 +96,6 @@ def init_state(
         # TODO(openhands): we should add test to test this init_state will actually
         # modify state in-place
 
-        # Validate security analyzer configuration once during initialization
-        if self._add_security_risk_prediction and isinstance(
-            state.confirmation_policy, NeverConfirm
-        ):
-            # If security analyzer is enabled, we always need a policy that is not
-            # NeverConfirm, otherwise we are just predicting risks without using them,
-            # and waste tokens!
-            logger.warning(
-                "LLM security analyzer is enabled but confirmation "
-                "policy is set to NeverConfirm"
-            )
-
         llm_convertible_messages = [
             event for event in state.events if isinstance(event, LLMConvertibleEvent)
         ]
@@ -105,10 +104,15 @@ def init_state(
             event = SystemPromptEvent(
                 source="agent",
                 system_prompt=TextContent(text=self.system_message),
+                # Always expose a 'security_risk' parameter in tool schemas.
+                # This ensures the schema remains consistent, even if the
+                # security analyzer is disabled. Validation of this field
+                # happens dynamically at runtime depending on the analyzer
+                # configured. This allows weaker models to omit risk field
+                # and bypass validation requirements when analyzer is disabled.
+                # For detailed logic, see `_extract_security_risk` method.
                 tools=[
-                    t.to_openai_tool(
-                        add_security_risk_prediction=self._add_security_risk_prediction
-                    )
+                    t.to_openai_tool(add_security_risk_prediction=True)
                     for t in self.tools_map.values()
                 ],
             )
@@ -176,15 +180,15 @@ def step(
                     tools=list(self.tools_map.values()),
                     include=None,
                     store=False,
-                    add_security_risk_prediction=self._add_security_risk_prediction,
+                    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=self._add_security_risk_prediction,
+                    add_security_risk_prediction=True,
                 )
         except FunctionCallValidationError as e:
             logger.warning(f"LLM generated malformed function call: {e}")
@@ -230,6 +234,7 @@ def step(
                     tool_call,
                     llm_response_id=llm_response.id,
                     on_event=on_event,
+                    security_analyzer=state.security_analyzer,
                     thought=thought_content
                     if i == 0
                     else [],  # Only first gets thought
@@ -300,10 +305,10 @@ def _requires_user_confirmation(
 
         # If a security analyzer is registered, use it to grab the risks of the actions
         # involved. If not, we'll set the risks to UNKNOWN.
-        if self.security_analyzer is not None:
+        if state.security_analyzer is not None:
             risks = [
                 risk
-                for _, risk in self.security_analyzer.analyze_pending_actions(
+                for _, risk in state.security_analyzer.analyze_pending_actions(
                     action_events
                 )
             ]
@@ -319,11 +324,44 @@ def _requires_user_confirmation(
 
         return False
 
+    def _extract_security_risk(
+        self,
+        arguments: dict,
+        tool_name: str,
+        read_only_tool: bool,
+        security_analyzer: analyzer.SecurityAnalyzerBase | None = None,
+    ) -> risk.SecurityRisk:
+        requires_sr = isinstance(security_analyzer, LLMSecurityAnalyzer)
+        raw = arguments.pop("security_risk", None)
+
+        # Default risk value for action event
+        # Tool is marked as read-only so security risk can be ignored
+        if read_only_tool:
+            return risk.SecurityRisk.UNKNOWN
+
+        # Raises exception if failed to pass risk field when expected
+        # Exception will be sent back to agent as error event
+        # Strong models like GPT-5 can correct itself by retrying
+        if requires_sr and raw is None:
+            raise ValueError(
+                f"Failed to provide security_risk field in tool '{tool_name}'"
+            )
+
+        # When using weaker models without security analyzer
+        # safely ignore missing security risk fields
+        if not requires_sr and raw is None:
+            return risk.SecurityRisk.UNKNOWN
+
+        # Raises exception if invalid risk enum passed by LLM
+        security_risk = risk.SecurityRisk(raw)
+        return security_risk
+
     def _get_action_event(
         self,
         tool_call: MessageToolCall,
         llm_response_id: str,
         on_event: ConversationCallbackType,
+        security_analyzer: analyzer.SecurityAnalyzerBase | None = None,
         thought: list[TextContent] | None = None,
         reasoning_content: str | None = None,
         thinking_blocks: list[ThinkingBlock | RedactedThinkingBlock] | None = None,
@@ -369,25 +407,18 @@ def _get_action_event(
 
             # Fix malformed arguments (e.g., JSON strings for list/dict fields)
             arguments = fix_malformed_tool_arguments(arguments, tool.action_type)
-
-            # if the tool has a security_risk field (when security analyzer is set),
-            # pop it out as it's not part of the tool's action schema
-            if (
-                _predicted_risk := arguments.pop("security_risk", None)
-            ) is not None and self.security_analyzer is not None:
-                try:
-                    security_risk = risk.SecurityRisk(_predicted_risk)
-                except ValueError:
-                    logger.warning(
-                        f"Invalid security_risk value from LLM: {_predicted_risk}"
-                    )
-
+            security_risk = self._extract_security_risk(
+                arguments,
+                tool.name,
+                tool.annotations.readOnlyHint if tool.annotations else False,
+                security_analyzer,
+            )
             assert "security_risk" not in arguments, (
                 "Unexpected 'security_risk' key found in tool arguments"
             )
 
             action: Action = tool.action_from_arguments(arguments)
-        except (json.JSONDecodeError, ValidationError) as e:
+        except (json.JSONDecodeError, ValidationError, ValueError) as e:
             err = (
                 f"Error validating args {tool_call.arguments} for tool "
                 f"'{tool.name}': {e}"

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

@@ -1,20 +1,20 @@
 import os
 import re
 import sys
+import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Generator, Iterable
 from typing import TYPE_CHECKING, Any
 
-from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator
 
-import openhands.sdk.security.analyzer as analyzer
 from openhands.sdk.context.agent_context import AgentContext
 from openhands.sdk.context.condenser import CondenserBase, LLMSummarizingCondenser
 from openhands.sdk.context.prompts.prompt import render_template
 from openhands.sdk.llm import LLM
 from openhands.sdk.logger import get_logger
 from openhands.sdk.mcp import create_mcp_tools
-from openhands.sdk.security.llm_analyzer import LLMSecurityAnalyzer
+from openhands.sdk.security import analyzer
 from openhands.sdk.tool import BUILT_IN_TOOLS, Tool, ToolDefinition, resolve_tool
 from openhands.sdk.utils.models import DiscriminatedUnionMixin
 from openhands.sdk.utils.pydantic_diff import pretty_pydantic_diff
@@ -27,6 +27,13 @@
 logger = get_logger(__name__)
 
 
+AGENT_SECURITY_ANALYZER_DEPRECATION_WARNING = (
+    "Agent.security_analyzer is deprecated and will be removed "
+    "in a future release.\n\n use `conversation = Conversation();"
+    "conversation.set_security_analyzer(...)` instead."
+)
+
+
 class AgentBase(DiscriminatedUnionMixin, ABC):
     """Abstract base class for OpenHands agents.
 
@@ -122,11 +129,13 @@ class AgentBase(DiscriminatedUnionMixin, ABC):
         description="Optional kwargs to pass to the system prompt Jinja2 template.",
         examples=[{"cli_mode": True}],
     )
+
     security_analyzer: analyzer.SecurityAnalyzerBase | None = Field(
         default=None,
         description="Optional security analyzer to evaluate action risks.",
         examples=[{"kind": "LLMSecurityAnalyzer"}],
     )
+
     condenser: CondenserBase | None = Field(
         default=None,
         description="Optional condenser to use for condensing conversation history.",
@@ -147,6 +156,22 @@ class AgentBase(DiscriminatedUnionMixin, ABC):
     # Runtime materialized tools; private and non-serializable
     _tools: dict[str, ToolDefinition] = PrivateAttr(default_factory=dict)
 
+    @model_validator(mode="before")
+    @classmethod
+    def _coerce_inputs(cls, data):
+        if not isinstance(data, dict):
+            return data
+        d = dict(data)
+
+        if "security_analyzer" in d and d["security_analyzer"]:
+            warnings.warn(
+                AGENT_SECURITY_ANALYZER_DEPRECATION_WARNING,
+                DeprecationWarning,
+                stacklevel=3,
+            )
+
+        return d
+
     @property
     def prompt_dir(self) -> str:
         """Returns the directory where this class's module file is located."""
@@ -164,13 +189,7 @@ def name(self) -> str:
     @property
     def system_message(self) -> str:
         """Compute system message on-demand to maintain statelessness."""
-        # Prepare template kwargs, including cli_mode if available
         template_kwargs = dict(self.system_prompt_kwargs)
-        if self.security_analyzer:
-            template_kwargs["llm_security_analyzer"] = bool(
-                isinstance(self.security_analyzer, LLMSecurityAnalyzer)
-            )
-
         system_message = render_template(
             prompt_dir=self.prompt_dir,
             template_name=self.system_prompt_filename,
@@ -198,6 +217,16 @@ def init_state(
 
     def _initialize(self, state: "ConversationState"):
         """Create an AgentBase instance from an AgentSpec."""
+
+        # 1) Migrate deprecated analyzer → state (if present)
+        if self.security_analyzer and not state.security_analyzer:
+            state.security_analyzer = self.security_analyzer
+            # 2) Clear on the immutable model (allowed via object.__setattr__)
+            try:
+                object.__setattr__(self, "security_analyzer", None)
+            except Exception:
+                logger.warning("Could not clear deprecated Agent.security_analyzer")
+
         if self._tools:
             logger.warning("Agent already initialized; skipping re-initialization.")
             return
@@ -297,8 +326,6 @@ def resolve_diff_from_deserialized(self, persisted: "AgentBase") -> "AgentBase":
                 updates["condenser"] = new_condenser
 
         # Allow security_analyzer to differ - use the runtime (self) version
-        # This allows users to add/remove security analyzers mid-conversation
-        # (e.g., when switching to weaker LLMs that can't handle security_risk field)
         updates["security_analyzer"] = self.security_analyzer
 
         # Create maps by tool name for easy lookup