fix(agent): authorize exact actions after untrusted context

src/tool_execution.py

@@ -28,6 +28,7 @@ from src.tool_security import (
     owner_is_admin_or_single_user,
 )
 from src.tool_capabilities import ToolRunSecurityContext, blocked_tool_result
+from src.tool_approvals import ExactToolApproval
 from src.tool_policy import ToolPolicy
 from src.constants import MAX_OUTPUT_CHARS, MAX_READ_CHARS, MAX_DIFF_LINES, DATA_DIR
 from src.tool_utils import _truncate, get_mcp_manager
@@ -567,10 +568,17 @@ async def _document_tool_dispatch(
     content: str,
     session_id: Optional[str] = None,
     owner: Optional[str] = None,
+    document_id: Optional[str] = None,
+    document_version: Optional[int] = None,
 ) -> Optional[Dict]:
     """Route a document tool through TOOL_HANDLERS with the right ctx shape."""
     from src.agent_tools import TOOL_HANDLERS
-    ctx = {"session_id": session_id, "owner": owner}
+    ctx = {
+        "session_id": session_id,
+        "owner": owner,
+        "doc_id": document_id,
+        "expected_document_version": document_version,
+    }
     if tool in TOOL_HANDLERS:
         return await TOOL_HANDLERS[tool](content, ctx)
     return None
@@ -593,6 +601,7 @@ async def execute_tool_block(
         | _NoToolSecurityContext
         | _MissingToolSecurityContext
     ) = _MISSING_TOOL_SECURITY_CONTEXT,
+    exact_approval: Optional[ExactToolApproval] = None,
 ) -> Tuple[str, Dict]:
     """Execute a single tool block. Returns (description, result_dict).
 
@@ -614,7 +623,55 @@ async def execute_tool_block(
             "NO_TOOL_SECURITY_CONTEXT"
         )
 
-    if isinstance(security_context, ToolRunSecurityContext):
+    approval_claimed = False
+    if exact_approval is not None:
+        if (
+            not isinstance(security_context, ToolRunSecurityContext)
+            or not security_context.external_untrusted_context_seen
+            or not exact_approval.pending.external_untrusted_context_seen
+        ):
+            return (
+                f"{getattr(block, 'tool_type', None)}: BLOCKED",
+                {
+                    "error": "Exact-action approval requires an armed run security context.",
+                    "exit_code": 1,
+                    "blocked": True,
+                    "policy": "exact_tool_approval",
+                },
+            )
+        sealed_workspace = exact_approval.pending.workspace
+        if sealed_workspace and vet_workspace(sealed_workspace) != sealed_workspace:
+            return (
+                f"{getattr(block, 'tool_type', None)}: BLOCKED",
+                {
+                    "error": (
+                        "The approved workspace is no longer a valid safe "
+                        "directory. Review the action again."
+                    ),
+                    "exit_code": 1,
+                    "blocked": True,
+                    "policy": "exact_tool_approval",
+                },
+            )
+        approval_claimed = exact_approval.claim(
+            owner=owner,
+            session_id=session_id,
+            tool_name=getattr(block, "tool_type", None),
+            content=getattr(block, "content", None),
+            workspace=workspace,
+        )
+        if not approval_claimed:
+            return (
+                f"{getattr(block, 'tool_type', None)}: BLOCKED",
+                {
+                    "error": "The exact-action approval did not match this tool request.",
+                    "exit_code": 1,
+                    "blocked": True,
+                    "policy": "exact_tool_approval",
+                },
+            )
+
+    if isinstance(security_context, ToolRunSecurityContext) and not approval_claimed:
         decision = security_context.decision_for(
             getattr(block, "tool_type", None),
             getattr(block, "content", None),
@@ -638,6 +695,16 @@ async def execute_tool_block(
             owner=owner,
             progress_cb=progress_cb,
             tool_policy=tool_policy,
+            approved_document_id=(
+                exact_approval.pending.document_id
+                if approval_claimed
+                else None
+            ),
+            approved_document_version=(
+                exact_approval.pending.document_version
+                if approval_claimed
+                else None
+            ),
         )
         if isinstance(security_context, ToolRunSecurityContext):
             security_context.observe_tool_result(
@@ -657,6 +724,8 @@ async def _execute_tool_block_impl(
     owner: Optional[str] = None,
     progress_cb: Optional[Callable[[Dict], Awaitable[None]]] = None,
     tool_policy: Optional[Any] = None,
+    approved_document_id: Optional[str] = None,
+    approved_document_version: Optional[int] = None,
 ) -> Tuple[str, Dict]:
     """Execute a single tool block. Returns (description, result_dict).
 
@@ -818,7 +887,14 @@ async def _execute_tool_block_impl(
     elif tool in ("create_document", "update_document", "edit_document",
                   "suggest_document", "manage_documents"):
         desc = f"{tool}: {content.split(chr(10))[0][:80]}"
-        result = await _document_tool_dispatch(tool, content, session_id, owner) \
+        result = await _document_tool_dispatch(
+            tool,
+            content,
+            session_id,
+            owner,
+            document_id=approved_document_id,
+            document_version=approved_document_version,
+        ) \
             or {"error": f"{tool}: execution failed", "exit_code": 1}
         if tool in ("edit_document", "suggest_document") and "title" in (result or {}):
             desc = f"{tool}: {result.get('title', '')}"