diff --git a/src/openrouter_agent/conversation_state.py b/src/openrouter_agent/conversation_state.py index 7d1bf0e..f187da0 100644 --- a/src/openrouter_agent/conversation_state.py +++ b/src/openrouter_agent/conversation_state.py @@ -7,7 +7,7 @@ from dataclasses import replace from typing import Any, Dict, List, Mapping, Optional, Sequence -from ._utils import dump, json_dumps, maybe_await +from ._utils import dump, json_dumps, maybe_await, validate_schema from .tool_types import ( ConversationState, ParsedToolCall, @@ -195,8 +195,6 @@ async def tool_requires_approval( context: Mapping[str, Any], call_level_check: Any = None, ) -> bool: - if call_level_check is not None: - return bool(await maybe_await(call_level_check(tool_call, context))) matching = next( ( candidate @@ -205,11 +203,37 @@ async def tool_requires_approval( ), None, ) + + # Normalize call-level policy arguments when possible while preserving its raw-call fallback. + if call_level_check is not None: + normalized_call = tool_call + if matching is not None: + try: + normalized_arguments = validate_schema( + get_tool_function(matching).get("input_schema"), + tool_call.arguments, + ) + except Exception: # noqa: BLE001 - policy callback gets the raw call on invalid input + pass + else: + normalized_call = replace(tool_call, arguments=normalized_arguments) + return bool(await maybe_await(call_level_check(normalized_call, context))) + if not matching: return False requirement = get_tool_function(matching).get("require_approval") if callable(requirement): - return bool(await maybe_await(requirement(tool_call.arguments, context))) + # Match execute_tool schema semantics and fail closed on invalid input. + if isinstance(tool_call.arguments, str): + return True + try: + normalized_arguments = validate_schema( + get_tool_function(matching).get("input_schema"), + tool_call.arguments, + ) + except Exception: # noqa: BLE001 - approval checks fail closed on invalid input + return True + return bool(await maybe_await(requirement(normalized_arguments, context))) return bool(requirement) diff --git a/src/openrouter_agent/model_result.py b/src/openrouter_agent/model_result.py index 11016c0..9a16a75 100644 --- a/src/openrouter_agent/model_result.py +++ b/src/openrouter_agent/model_result.py @@ -493,7 +493,7 @@ async def _run_tool_with_hooks( return {"type": "parse_error", "call": call, "error_message": error_message} effective_call = call - if self._hooks: + if self._hooks and call.pre_tool_use_applied is not True: original_input = call.arguments if isinstance(call.arguments, dict) else {} pre = await self._hooks.emit( "PreToolUse", @@ -543,6 +543,118 @@ async def _run_tool_with_hooks( return {"type": "execution", "call": effective_call, "result": result} + async def _prepare_tool_calls_for_approval( + self, + calls: Sequence[ParsedToolCall], + tools: Sequence[Tool], + ) -> Tuple[List[ParsedToolCall], List[UnsentToolResult]]: + """Apply PreToolUse once before approval classification.""" + prepared: List[ParsedToolCall] = [] + blocked: List[UnsentToolResult] = [] + for call in calls: + tool = self._find_tool(call.name, tools) + if ( + not self._hooks + or call.pre_tool_use_applied is True + or isinstance(call.arguments, str) + or not tool + or not is_auto_resolvable_tool(tool) + ): + prepared.append(call) + continue + + original_input = call.arguments if isinstance(call.arguments, dict) else {} + pre = await self._hooks.emit( + "PreToolUse", + {"tool_name": call.name, "tool_input": original_input}, + tool_name=call.name, + session_id=self._session_id, + ) + if pre.blocked: + block = next((result.get("block") for result in pre.results if result.get("block")), None) + reason = block if isinstance(block, str) else "Blocked by PreToolUse hook" + blocked.append(create_rejected_result(call.id, call.name, reason)) + continue + + arguments = pre.final_payload.get("tool_input") if pre.mutated else call.arguments + prepared.append(dataclasses.replace(call, arguments=arguments, pre_tool_use_applied=True)) + + return prepared, blocked + + async def _resolve_approval_gate( + self, + calls: Sequence[ParsedToolCall], + tools: Sequence[Tool], + turn_context: Mapping[str, Any], + ) -> Tuple[List[ParsedToolCall], List[ParsedToolCall], List[UnsentToolResult]]: + """Return pending calls, executable calls, and resolved rejection results.""" + prepared, resolved_results = await self._prepare_tool_calls_for_approval(calls, tools) + partition = await partition_tool_calls( + prepared, + tools, + turn_context, + self.options.get("require_approval"), + ) + needs_approval = list(partition["requires_approval"]) + pending: List[ParsedToolCall] = [] + allowed_occurrences: set[int] = set() + denied_occurrences: set[int] = set() + + if needs_approval and self._hooks: + for call in needs_approval: + decision, reason = await self._emit_permission_request(call, tools) + if decision == "allow": + allowed_occurrences.add(id(call)) + elif decision == "deny": + denied_occurrences.add(id(call)) + resolved_results.append( + create_rejected_result(call.id, call.name, reason or "Denied by PermissionRequest hook") + ) + else: + pending.append(call) + else: + pending.extend(needs_approval) + + pending_occurrences = {id(call) for call in pending} + auto_execute_occurrences = {id(call) for call in partition["auto_execute"]} + executable = [ + call + for call in prepared + if id(call) not in pending_occurrences + and id(call) not in denied_occurrences + and (id(call) in allowed_occurrences or id(call) in auto_execute_occurrences) + ] + return pending, executable, resolved_results + + async def _execute_calls_to_unsent_results( + self, + calls: Sequence[ParsedToolCall], + tools: Sequence[Tool], + turn_context: Mapping[str, Any], + ) -> List[UnsentToolResult]: + """Execute and persist auto-resolvable calls before an approval pause.""" + results: List[UnsentToolResult] = [] + for call in calls: + tool = self._find_tool(call.name, tools) + if not tool or not is_auto_resolvable_tool(tool): + continue + outcome = await self._run_tool_with_hooks(tool, call, {**turn_context, "tool_call": call}) + if outcome["type"] == "parse_error": + results.append(create_rejected_result(call.id, call.name, outcome["error_message"])) + elif outcome["type"] == "hook_blocked": + results.append(create_rejected_result(call.id, call.name, outcome["reason"])) + elif outcome["result"] is None: + # HITL tools cannot produce an unsent result until resumed. + continue + elif outcome["result"].get("error") is not None: + results.append(create_rejected_result(call.id, call.name, str(outcome["result"]["error"]))) + else: + effective_call = outcome["call"] + results.append( + create_unsent_result(effective_call.id, effective_call.name, outcome["result"].get("result")) + ) + return results + async def _record_preliminary(self, call_id: str, value: Any) -> None: await self._append_event( { @@ -714,9 +826,32 @@ async def _run(self) -> Any: final_response_enabled = allow_final_response is not False resolvable_pending = [c for c in calls if self._call_is_auto_resolvable(c, tools)] if final_response_enabled and resolvable_pending: - final_outputs: List[Any] = [] turn_context = {"number_of_turns": turn_number + 1, "turn_request": current_request} - for call in calls: + requires_approval, executable_calls, gate_results = await self._resolve_approval_gate( + calls, tools, turn_context + ) + if requires_approval: + if self.options.get("state") is None: + names = ", ".join(call.name for call in requires_approval) + raise ValueError( + f"Tool(s) require approval but no state accessor is configured: {names}" + ) + gate_results.extend( + await self._execute_calls_to_unsent_results(executable_calls, tools, turn_context) + ) + await self._save_state( + pending_tool_calls=requires_approval, + unsent_tool_results=gate_results or None, + status="awaiting_approval", + ) + self._final_response = final_response + return final_response + + final_outputs: List[Any] = unsent_results_to_api_format(gate_results) + resolved_gate_ids = {result.call_id for result in gate_results} + for call in executable_calls: + if call.id in resolved_gate_ids: + continue matching_tool = self._find_tool(call.name, tools) if matching_tool and is_auto_resolvable_tool(matching_tool): outcome = await self._run_tool_with_hooks(matching_tool, call, turn_context) @@ -756,59 +891,20 @@ async def _run(self) -> Any: continue break - partition = await partition_tool_calls( - calls, tools, {"number_of_turns": turn_number + 1}, self.options.get("require_approval") + turn_context = {"number_of_turns": turn_number + 1, "turn_request": current_request} + requires_approval, executable_calls, hook_resolved_unsent = await self._resolve_approval_gate( + calls, + tools, + turn_context, ) - requires_approval = list(partition["requires_approval"]) - hook_resolved_unsent: List[UnsentToolResult] = [] - if requires_approval and hooks: - still_pending: List[ParsedToolCall] = [] - for call in requires_approval: - decision, reason = await self._emit_permission_request(call, tools) - if decision == "allow": - promo_tool = self._find_tool(call.name, tools) - if promo_tool and is_auto_resolvable_tool(promo_tool): - outcome = await self._run_tool_with_hooks( - promo_tool, - call, - { - "number_of_turns": turn_number + 1, - "tool_call": call, - "turn_request": current_request, - }, - ) - if outcome["type"] == "parse_error": - hook_resolved_unsent.append( - create_rejected_result(call.id, call.name, outcome["error_message"]) - ) - elif outcome["type"] == "hook_blocked": - hook_resolved_unsent.append( - create_rejected_result(call.id, call.name, outcome["reason"]) - ) - elif outcome["result"] is None: - still_pending.append(call) - elif outcome["result"].get("error") is not None: - hook_resolved_unsent.append( - create_rejected_result(call.id, call.name, str(outcome["result"]["error"])) - ) - else: - hook_resolved_unsent.append( - create_unsent_result(call.id, call.name, outcome["result"].get("result")) - ) - else: - still_pending.append(call) - elif decision == "deny": - hook_resolved_unsent.append( - create_rejected_result(call.id, call.name, reason or "Denied by PermissionRequest hook") - ) - else: - still_pending.append(call) - requires_approval = still_pending if requires_approval: if self.options.get("state") is None: names = ", ".join(call.name for call in requires_approval) raise ValueError(f"Tool(s) require approval but no state accessor is configured: {names}") + hook_resolved_unsent.extend( + await self._execute_calls_to_unsent_results(executable_calls, tools, turn_context) + ) save_kwargs: Dict[str, Any] = { "pending_tool_calls": requires_approval, "status": "awaiting_approval", @@ -822,7 +918,7 @@ async def _run(self) -> Any: outputs: List[Any] = unsent_results_to_api_format(hook_resolved_unsent) if hook_resolved_unsent else [] paused: List[ParsedToolCall] = [] executed_calls: List[ParsedToolCall] = [] - for call in partition["auto_execute"]: + for call in executable_calls: tool = self._find_tool(call.name, tools) if not tool or not is_auto_resolvable_tool(tool): continue diff --git a/src/openrouter_agent/tool_types.py b/src/openrouter_agent/tool_types.py index f6db4b9..a28aa7a 100644 --- a/src/openrouter_agent/tool_types.py +++ b/src/openrouter_agent/tool_types.py @@ -24,6 +24,8 @@ class ParsedToolCall: id: str name: str arguments: Any + # Prevent PreToolUse from running twice after an approval resume. + pre_tool_use_applied: Optional[bool] = None @dataclass(frozen=True) diff --git a/tests/unit/test_approval_gate_regressions.py b/tests/unit/test_approval_gate_regressions.py new file mode 100644 index 0000000..3250492 --- /dev/null +++ b/tests/unit/test_approval_gate_regressions.py @@ -0,0 +1,293 @@ +"""Approval-gate regressions ported from upstream PR #94.""" + +from __future__ import annotations + +from typing import Any, List + +from pydantic import BaseModel + +from openrouter_agent import HookEntry, HookName, HooksManager, call_model, step_count_is, tool +from tests._fixtures import ( + MemoryStateAccessor, + QueuedClient, + function_call_item, + make_response, + text_response, + tool_call_response, +) + + +async def test_allow_final_response_pauses_instead_of_executing_an_approval_gated_tool() -> None: + state = MemoryStateAccessor() + executed: List[Any] = [] + + def execute(params: Any, _context: Any) -> Any: + executed.append(params) + return {"ok": True} + + gated = tool( + name="deploy", + input_schema=dict, + output_schema=dict, + execute=execute, + require_approval=True, + ) + client = QueuedClient([tool_call_response("r1", "deploy")]) + + response = await call_model( + client, + { + "model": "test-model", + "input": "deploy", + "tools": [gated], + "state": state, + "stop_when": step_count_is(0), + "allow_final_response": True, + }, + ).get_response() + + assert response["id"] == "r1" + assert executed == [] + assert len(client.requests) == 1 + assert state.stored is not None + assert state.stored.status == "awaiting_approval" + assert [call.name for call in state.stored.pending_tool_calls or []] == ["deploy"] + + +async def test_allow_final_response_honors_permission_request_deny() -> None: + executed: List[Any] = [] + permission_requests: List[Any] = [] + hooks = HooksManager() + + def deny(payload: Any, _context: Any) -> Any: + permission_requests.append(payload) + return {"decision": "deny", "reason": "policy denied"} + + hooks.on(HookName.PermissionRequest.value, HookEntry(handler=deny)) + + def execute(params: Any, _context: Any) -> Any: + executed.append(params) + return {"ok": True} + + gated = tool( + name="deploy", + input_schema=dict, + output_schema=dict, + execute=execute, + require_approval=True, + ) + client = QueuedClient([tool_call_response("r1", "deploy"), text_response("r2", "not deployed")]) + + text = await call_model( + client, + { + "model": "test-model", + "input": "deploy", + "tools": [gated], + "hooks": hooks, + "stop_when": step_count_is(0), + "allow_final_response": True, + }, + ).get_text() + + assert text == "not deployed" + assert executed == [] + assert len(permission_requests) == 1 + rejection = next(item for item in client.requests[1]["input"] if item.get("type") == "function_call_output") + assert "policy denied" in rejection["output"] + + +async def test_allow_final_response_persists_ungated_results_before_approval_pause() -> None: + state = MemoryStateAccessor() + executions = {"lookup": 0, "deploy": 0} + + def execute_lookup(_params: Any, _context: Any) -> Any: + executions["lookup"] += 1 + return {"value": 42} + + def execute_deploy(_params: Any, _context: Any) -> Any: + executions["deploy"] += 1 + return {"ok": True} + + lookup = tool(name="lookup", input_schema=dict, output_schema=dict, execute=execute_lookup) + deploy = tool( + name="deploy", + input_schema=dict, + output_schema=dict, + execute=execute_deploy, + require_approval=True, + ) + response = make_response( + "r1", + [function_call_item("call_lookup", "lookup"), function_call_item("call_deploy", "deploy")], + ) + + await call_model( + QueuedClient([response]), + { + "model": "test-model", + "input": "lookup then deploy", + "tools": [lookup, deploy], + "state": state, + "stop_when": step_count_is(0), + }, + ).get_response() + + assert executions == {"lookup": 1, "deploy": 0} + assert state.stored is not None + assert [call.name for call in state.stored.pending_tool_calls or []] == ["deploy"] + assert [result.name for result in state.stored.unsent_tool_results or []] == ["lookup"] + + +async def test_tool_approval_predicate_receives_schema_defaults() -> None: + class DeployInput(BaseModel): + dangerous: bool = True + + state = MemoryStateAccessor() + observed: List[bool] = [] + executed: List[Any] = [] + + def needs_approval(params: DeployInput, _context: Any) -> bool: + observed.append(params.dangerous) + return params.dangerous + + def execute(params: Any, _context: Any) -> Any: + executed.append(params) + return {"ok": True} + + gated = tool( + name="deploy", + input_schema=DeployInput, + output_schema=dict, + execute=execute, + require_approval=needs_approval, + ) + + await call_model( + QueuedClient([tool_call_response("r1", "deploy", arguments="{}")]), + {"model": "test-model", "input": "deploy", "tools": [gated], "state": state}, + ).get_response() + + assert observed == [True] + assert executed == [] + assert state.stored is not None and state.stored.status == "awaiting_approval" + + +async def test_call_level_approval_predicate_receives_schema_normalized_call() -> None: + class DeployInput(BaseModel): + dangerous: bool = True + + state = MemoryStateAccessor() + observed: List[bool] = [] + + def needs_approval(call: Any, _context: Any) -> bool: + observed.append(call.arguments.dangerous) + return call.arguments.dangerous + + deploy = tool( + name="deploy", + input_schema=DeployInput, + output_schema=dict, + execute=lambda params, ctx: {"ok": True}, + ) + + await call_model( + QueuedClient([tool_call_response("r1", "deploy", arguments="{}")]), + { + "model": "test-model", + "input": "deploy", + "tools": [deploy], + "state": state, + "require_approval": needs_approval, + }, + ).get_response() + + assert observed == [True] + assert state.stored is not None and state.stored.status == "awaiting_approval" + + +async def test_pre_tool_use_mutation_is_approved_once_and_persists_across_resume() -> None: + class DeployInput(BaseModel): + dangerous: bool = False + + state = MemoryStateAccessor() + hooks = HooksManager() + pre_tool_use_calls: List[Any] = [] + predicate_values: List[bool] = [] + executed: List[DeployInput] = [] + + def mutate(payload: Any, _context: Any) -> Any: + pre_tool_use_calls.append(payload) + return {"mutated_input": {"dangerous": True}} + + def needs_approval(params: DeployInput, _context: Any) -> bool: + predicate_values.append(params.dangerous) + return params.dangerous + + def execute(params: DeployInput, _context: Any) -> Any: + executed.append(params) + return {"ok": True} + + hooks.on(HookName.PreToolUse.value, HookEntry(handler=mutate)) + gated = tool( + name="deploy", + input_schema=DeployInput, + output_schema=dict, + execute=execute, + require_approval=needs_approval, + ) + + await call_model( + QueuedClient([tool_call_response("r1", "deploy", arguments='{"dangerous": false}')]), + {"model": "test-model", "input": "deploy", "tools": [gated], "state": state, "hooks": hooks}, + ).get_response() + + assert predicate_values == [True] + assert len(pre_tool_use_calls) == 1 + assert state.stored is not None + pending = state.stored.pending_tool_calls or [] + assert pending[0].arguments == {"dangerous": True} + assert pending[0].pre_tool_use_applied is True + + resumed = call_model( + QueuedClient([text_response("r2", "deployed")]), + { + "model": "test-model", + "input": "continue", + "tools": [gated], + "state": state, + "hooks": hooks, + "approve_tool_calls": [pending[0].id], + }, + ) + assert await resumed.get_text() == "deployed" + assert len(pre_tool_use_calls) == 1 + assert len(executed) == 1 and executed[0].dangerous is True + + +async def test_invalid_tool_input_fails_closed_before_predicate() -> None: + class DeployInput(BaseModel): + dangerous: bool + + state = MemoryStateAccessor() + predicate_calls: List[Any] = [] + + def needs_approval(params: Any, _context: Any) -> bool: + predicate_calls.append(params) + return False + + gated = tool( + name="deploy", + input_schema=DeployInput, + output_schema=dict, + execute=lambda params, ctx: {"ok": True}, + require_approval=needs_approval, + ) + + await call_model( + QueuedClient([tool_call_response("r1", "deploy", arguments="{}")]), + {"model": "test-model", "input": "deploy", "tools": [gated], "state": state}, + ).get_response() + + assert predicate_calls == [] + assert state.stored is not None and state.stored.status == "awaiting_approval" diff --git a/tests/unit/test_conversation_state_serialization.py b/tests/unit/test_conversation_state_serialization.py index 74fe1b6..381ba0f 100644 --- a/tests/unit/test_conversation_state_serialization.py +++ b/tests/unit/test_conversation_state_serialization.py @@ -43,7 +43,14 @@ def test_round_trips_a_rich_awaiting_client_tools_state() -> None: "status": "completed", }, ], - pending_tool_calls=[ParsedToolCall(id="call_manual_1", name="exec_command", arguments={"command": "ls"})], + pending_tool_calls=[ + ParsedToolCall( + id="call_manual_1", + name="exec_command", + arguments={"command": "ls"}, + pre_tool_use_applied=True, + ) + ], unsent_tool_results=[ UnsentToolResult(call_id="call_auto_1", name="auto_search", output={"result": "found it"}) ], @@ -54,7 +61,12 @@ def test_round_trips_a_rich_awaiting_client_tools_state() -> None: assert restored == rich assert restored.status == "awaiting_client_tools" assert restored.pending_tool_calls == [ - ParsedToolCall(id="call_manual_1", name="exec_command", arguments={"command": "ls"}) + ParsedToolCall( + id="call_manual_1", + name="exec_command", + arguments={"command": "ls"}, + pre_tool_use_applied=True, + ) ] assert restored.unsent_tool_results is not None assert restored.unsent_tool_results[0].call_id == "call_auto_1"