fix: bind MCP approvals to tool calls

Signed-off-by: SabhyaC26 <sabhyachhabria@gmail.com>

omnigent/server/routes/_sessions/common.py

@@ -568,11 +568,19 @@ class _PendingPolicyAskWrites:
         itself, so the events handler skips write application for
         these entries to avoid double-applying non-idempotent ops
         (e.g. ``INCREMENT`` state updates for cost-budget counters).
+    :param session_id: Session that owns an MCP approval.
+    :param tool_name: Exact MCP tool authorized by the approval.
+    :param arguments_hash: SHA-256 of the canonical MCP arguments.
+    :param transformed_arguments: Policy-transformed arguments to execute.
     """
 
     state_updates: list[StateUpdate] | None
     set_labels: dict[str, str] | None
     from_mcp: bool = False
+    session_id: str | None = None
+    tool_name: str | None = None
+    arguments_hash: str | None = None
+    transformed_arguments: dict[str, Any] | None = None
 
 
 _pending_policy_ask_writes: cachetools.LRUCache[str, _PendingPolicyAskWrites] = (

omnigent/server/routes/_sessions/orchestration.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 
 import asyncio
 import contextlib
+import hashlib
 import json
 import secrets
 import time
@@ -8282,6 +8283,32 @@ async def _child_session_summaries_from_conversations(
     ]
 
 
+def _mcp_arguments_hash(arguments: dict[str, Any]) -> str:
+    """Hash MCP arguments using a stable JSON representation."""
+    canonical = json.dumps(
+        arguments,
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=False,
+    ).encode()
+    return hashlib.sha256(canonical).hexdigest()
+
+
+def _mcp_approval_matches(
+    pending: _PendingPolicyAskWrites,
+    session_id: str,
+    tool_name: str,
+    arguments: dict[str, Any],
+) -> bool:
+    """Return whether a pending approval authorizes this invocation."""
+    return (
+        pending.from_mcp
+        and pending.session_id == session_id
+        and pending.tool_name == tool_name
+        and pending.arguments_hash == _mcp_arguments_hash(arguments)
+    )
+
+
 async def _handle_mcp_tools_call(
     rpc_id: int | str | None,
     session_id: str,
@@ -8422,7 +8449,8 @@ async def _handle_mcp_tools_call(
             # was genuinely issued by the server (present in the
             # server-side pending map) and that the user approved it.
             elicitation_id_from_state: str = state.get("elicitation_id", "")
-            if elicitation_id_from_state not in _pending_policy_ask_writes:
+            pending = _pending_policy_ask_writes.get(elicitation_id_from_state)
+            if pending is None or not pending.from_mcp:
                 # The elicitation_id is not in the server-side map.
                 # Either it was forged, already consumed, or expired.
                 # Check inputResponses: if the caller claims approval
@@ -8437,29 +8465,40 @@ async def _handle_mcp_tools_call(
                         "Elicitation not found or already resolved",
                     )
                 return _mcp_error_response(rpc_id, -32000, "Tool call denied by user")
+            if not _mcp_approval_matches(pending, session_id, namespaced_name, arguments):
+                return _mcp_error_response(
+                    rpc_id,
+                    -32000,
+                    "Approval does not match this tool call",
+                )
+            # Compare and consume without yielding so concurrent retries cannot
+            # both spend the same approval. Mismatches leave it intact.
+            _pending_policy_ask_writes.pop(elicitation_id_from_state, None)
             approval = input_responses.get(elicitation_id_from_state) or {}
             if approval.get("action") != "accept":
                 return _mcp_error_response(rpc_id, -32000, "Tool call denied by user")
-            # Recover any policy-transformed args that were serialised into
-            # requestState on the initial ASK — the client re-sends the
-            # original arguments which we must not use when a transform was set.
-            if state.get("transformed_arguments") is not None:
-                arguments = state["transformed_arguments"]
+            if pending.transformed_arguments is not None:
+                arguments = pending.transformed_arguments
             # Apply the deciding policy's deferred writes now that the
             # user approved (POLICIES.md §7.2: only on accept).
-            _pending = _pending_policy_ask_writes.pop(elicitation_id_from_state, None)
-            if _pending is not None:
-                if _pending.set_labels:
-                    await asyncio.to_thread(engine.apply_label_writes, _pending.set_labels)
-                if _pending.state_updates:
-                    with contextlib.suppress(ConversationNotFoundError):
-                        await asyncio.to_thread(engine.apply_state_updates, _pending.state_updates)
+            if pending.set_labels:
+                await asyncio.to_thread(engine.apply_label_writes, pending.set_labels)
+            if pending.state_updates:
+                with contextlib.suppress(ConversationNotFoundError):
+                    await asyncio.to_thread(engine.apply_state_updates, pending.state_updates)
         else:
             # ALLOW — policy no longer requires approval (e.g. label
             # state changed between the original ASK and this retry).
-            # Recover transformed args if present, then fall through.
-            if state.get("transformed_arguments") is not None:
-                arguments = state["transformed_arguments"]
+            # Consume a matching stale approval even though it is no longer
+            # required, so it cannot authorize a later ASK retry.
+            elicitation_id_from_state = state.get("elicitation_id", "")
+            pending = _pending_policy_ask_writes.get(elicitation_id_from_state)
+            if pending is not None and _mcp_approval_matches(
+                pending, session_id, namespaced_name, arguments
+            ):
+                _pending_policy_ask_writes.pop(elicitation_id_from_state, None)
+                if pending.transformed_arguments is not None:
+                    arguments = pending.transformed_arguments
         # Fall through to execution.
     else:
         # ── First call: evaluate TOOL_CALL policy ────────────────────
@@ -8513,17 +8552,15 @@ async def _handle_mcp_tools_call(
                 state_updates=call_result.state_updates,
                 set_labels=call_result.set_labels,
                 from_mcp=True,
+                session_id=session_id,
+                tool_name=namespaced_name,
+                arguments_hash=_mcp_arguments_hash(arguments),
+                transformed_arguments=cast("dict[str, Any] | None", call_result.data),
             )
             request_state_payload: dict[str, Any] = {
                 "elicitation_id": elicitation_id,
                 "session_id": session_id,
             }
-            # If the policy returned transformed args alongside ASK (e.g.
-            # PII-redacted arguments), persist them so the retry path can
-            # apply them after the user approves — the client re-sends the
-            # original arguments, which would silently bypass the transform.
-            if call_result.data is not None:
-                request_state_payload["transformed_arguments"] = call_result.data
             request_state = json.dumps(request_state_payload)
             return _mcp_input_required_response(
                 rpc_id,

tests/server/routes/test_sessions_mcp_proxy_policy_retry.py

@@ -12,6 +12,7 @@ dependencies, following the pattern in ``test_sessions_mcp_proxy.py``.
 
 from __future__ import annotations
 
+import hashlib
 import json
 from dataclasses import dataclass
 from typing import Any
@@ -36,17 +37,17 @@ from omnigent.spec.types import PolicyAction
 _SESSION_ID = "conv_test_policy_retry"
 
 
-def _make_conversation() -> Conversation:
+def _make_conversation(session_id: str = _SESSION_ID) -> Conversation:
     """
     Build a minimal :class:`Conversation` with an agent binding.
 
     :returns: A :class:`Conversation` pointing at agent ``"ag_test"``.
     """
     return Conversation(
-        id=_SESSION_ID,
+        id=session_id,
         created_at=0,
         updated_at=0,
-        root_conversation_id=_SESSION_ID,
+        root_conversation_id=session_id,
         agent_id="ag_test",
     )
 
@@ -195,6 +196,51 @@ def _parse_rpc_error(response: Any) -> dict[str, Any]:
     return payload["error"]
 
 
+def _arguments_hash(arguments: dict[str, Any]) -> str:
+    canonical = json.dumps(
+        arguments, sort_keys=True, separators=(",", ":"), ensure_ascii=False
+    ).encode()
+    return hashlib.sha256(canonical).hexdigest()
+
+
+async def _issue_ask(
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    session_id: str = _SESSION_ID,
+    tool_name: str = "sys_os_shell",
+    arguments: dict[str, Any] | None = None,
+) -> tuple[str, str]:
+    """Issue a real first-call ASK and return its id and request state."""
+    ask_engine = _FixedPolicyEngine(
+        result=PolicyResult(
+            action=PolicyAction.ASK,
+            reason="approval required",
+            deciding_policies=["test-gate"],
+        )
+    )
+    monkeypatch.setattr(
+        sessions_mod,
+        "_load_agent_spec_for_session",
+        lambda conv, agent_store: "fake_spec",
+    )
+    monkeypatch.setattr(
+        sessions_mod,
+        "_build_policy_engine_from_spec",
+        _engine_factory_expecting_preload(ask_engine),
+    )
+    response = await _handle_mcp_tools_call(
+        rpc_id=10,
+        session_id=session_id,
+        params={"name": tool_name, "arguments": arguments or {"command": "id -un"}},
+        conversation_store=_StubConversationStore(_make_conversation(session_id)),  # type: ignore[arg-type]
+        agent_store=_StubAgentStore(),  # type: ignore[arg-type]
+        runner_router=None,
+    )
+    result = json.loads(bytes(response.body))["result"]
+    elicitation_id = next(iter(result["inputRequests"]))
+    return elicitation_id, result["requestState"]
+
+
 # ---------------------------------------------------------------------------
 # Tests
 # ---------------------------------------------------------------------------
@@ -315,6 +361,107 @@ async def test_forged_retry_with_ask_policy_rejects_unknown_elicitation(
     )
 
 
+@pytest.mark.parametrize(
+    ("replay_session", "replay_tool", "replay_arguments"),
+    [
+        (_SESSION_ID, "sys_os_read", {"path": "/etc/passwd"}),
+        (_SESSION_ID, "sys_os_shell", {"command": "whoami"}),
+        ("conv_other_session", "sys_os_shell", {"command": "id -un"}),
+    ],
+    ids=["tool", "arguments", "session"],
+)
+@pytest.mark.asyncio
+async def test_approval_is_bound_to_exact_invocation(
+    monkeypatch: pytest.MonkeyPatch,
+    replay_session: str,
+    replay_tool: str,
+    replay_arguments: dict[str, Any],
+) -> None:
+    """An ASK approval cannot authorize a different tool call."""
+    elicitation_id, request_state = await _issue_ask(monkeypatch)
+    state = json.loads(request_state)
+    state["session_id"] = replay_session
+
+    try:
+        response = await _handle_mcp_tools_call(
+            rpc_id=11,
+            session_id=replay_session,
+            params={
+                "name": replay_tool,
+                "arguments": replay_arguments,
+                "requestState": json.dumps(state),
+                "inputResponses": {elicitation_id: {"action": "accept"}},
+            },
+            conversation_store=_StubConversationStore(_make_conversation(replay_session)),  # type: ignore[arg-type]
+            agent_store=_StubAgentStore(),  # type: ignore[arg-type]
+            runner_router=None,
+        )
+        assert "does not match" in _parse_rpc_error(response)["message"].lower()
+        assert elicitation_id in _pending_policy_ask_writes
+    finally:
+        _pending_policy_ask_writes.pop(elicitation_id, None)
+
+
+@pytest.mark.asyncio
+async def test_approval_is_consumed_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A matching approval cannot be replayed after its first retry."""
+    elicitation_id, request_state = await _issue_ask(
+        monkeypatch,
+        arguments={"a": 1, "b": 2},
+    )
+    params = {
+        "name": "sys_os_shell",
+        "arguments": {"b": 2, "a": 1},
+        "requestState": request_state,
+        "inputResponses": {elicitation_id: {"action": "accept"}},
+    }
+    store = _StubConversationStore(_make_conversation())
+
+    try:
+        first = await _handle_mcp_tools_call(
+            rpc_id=12,
+            session_id=_SESSION_ID,
+            params=params,
+            conversation_store=store,  # type: ignore[arg-type]
+            agent_store=_StubAgentStore(),  # type: ignore[arg-type]
+            runner_router=None,
+        )
+        assert "runner" in _parse_rpc_error(first)["message"].lower()
+
+        second = await _handle_mcp_tools_call(
+            rpc_id=13,
+            session_id=_SESSION_ID,
+            params=params,
+            conversation_store=store,  # type: ignore[arg-type]
+            agent_store=_StubAgentStore(),  # type: ignore[arg-type]
+            runner_router=None,
+        )
+        assert "already resolved" in _parse_rpc_error(second)["message"].lower()
+    finally:
+        _pending_policy_ask_writes.pop(elicitation_id, None)
+
+
+@pytest.mark.asyncio
+async def test_matching_decline_consumes_approval(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A matching decline is terminal and cannot be replayed later."""
+    elicitation_id, request_state = await _issue_ask(monkeypatch)
+    response = await _handle_mcp_tools_call(
+        rpc_id=14,
+        session_id=_SESSION_ID,
+        params={
+            "name": "sys_os_shell",
+            "arguments": {"command": "id -un"},
+            "requestState": request_state,
+            "inputResponses": {elicitation_id: {"action": "decline"}},
+        },
+        conversation_store=_StubConversationStore(_make_conversation()),  # type: ignore[arg-type]
+        agent_store=_StubAgentStore(),  # type: ignore[arg-type]
+        runner_router=None,
+    )
+    assert "denied by user" in _parse_rpc_error(response)["message"].lower()
+    assert elicitation_id not in _pending_policy_ask_writes
+
+
 @pytest.mark.asyncio
 async def test_retry_with_allow_policy_falls_through(
     monkeypatch: pytest.MonkeyPatch,
@@ -370,6 +517,35 @@ async def test_retry_with_allow_policy_falls_through(
         )
 
 
+@pytest.mark.asyncio
+async def test_matching_retry_consumed_when_policy_changes_to_allow(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A now-unneeded matching approval does not linger for later replay."""
+    elicitation_id, request_state = await _issue_ask(monkeypatch)
+    allow_engine = _FixedPolicyEngine(result=PolicyResult(action=PolicyAction.ALLOW, reason=None))
+    monkeypatch.setattr(
+        sessions_mod,
+        "_build_policy_engine_from_spec",
+        _engine_factory_expecting_preload(allow_engine),
+    )
+
+    await _handle_mcp_tools_call(
+        rpc_id=15,
+        session_id=_SESSION_ID,
+        params={
+            "name": "sys_os_shell",
+            "arguments": {"command": "id -un"},
+            "requestState": request_state,
+            "inputResponses": {elicitation_id: {"action": "accept"}},
+        },
+        conversation_store=_StubConversationStore(_make_conversation()),  # type: ignore[arg-type]
+        agent_store=_StubAgentStore(),  # type: ignore[arg-type]
+        runner_router=None,
+    )
+    assert elicitation_id not in _pending_policy_ask_writes
+
+
 @pytest.mark.asyncio
 async def test_retry_session_mismatch_still_rejected(
     monkeypatch: pytest.MonkeyPatch,
@@ -464,6 +640,9 @@ async def test_legitimate_retry_with_pending_entry_proceeds(
         state_updates=None,
         set_labels=None,
         from_mcp=True,
+        session_id=_SESSION_ID,
+        tool_name="sys_os_shell",
+        arguments_hash=_arguments_hash({"command": "id -un"}),
     )
 
     try: