From 7a138e8a3f93a7cb2da93c71b006c20ada5f75cf Mon Sep 17 00:00:00 2001 From: RaresKeY <158580472+RaresKeY@users.noreply.github.com> Date: Sat, 15 Aug 2026 07:01:36 +0000 Subject: [PATCH] fix(agent): seal document approval content --- src/agent_loop.py | 17 +++++- src/agent_tools/document_tools.py | 24 +++++++-- src/tool_approvals.py | 12 +++++ src/tool_execution.py | 10 ++++ tests/test_external_context_tool_gate.py | 69 ++++++++++++++++++++++++ tests/test_foreground_model_routing.py | 2 + tests/test_tool_approvals.py | 27 ++++++++-- 7 files changed, 151 insertions(+), 10 deletions(-) diff --git a/src/agent_loop.py b/src/agent_loop.py index d79e645a8..e24c9c0f0 100644 --- a/src/agent_loop.py +++ b/src/agent_loop.py @@ -43,7 +43,11 @@ from src.tool_capabilities import ( tool_result_is_successful, tool_result_should_arm_gate, ) -from src.tool_approvals import ExactToolApproval, tool_approval_store +from src.tool_approvals import ( + ExactToolApproval, + document_content_digest, + tool_approval_store, +) from src.tool_utils import _truncate, get_mcp_manager from src.agent_tools import ( parse_tool_blocks, @@ -5666,6 +5670,17 @@ async def stream_agent_loop( "version_count", None, ), + document_digest=( + document_content_digest( + getattr( + approval_document, + "current_content", + "", + ) + ) + if approval_document is not None + else None + ), external_untrusted_context_seen=( run_security.external_untrusted_context_seen ), diff --git a/src/agent_tools/document_tools.py b/src/agent_tools/document_tools.py index 8ffe24c93..58ec77b56 100644 --- a/src/agent_tools/document_tools.py +++ b/src/agent_tools/document_tools.py @@ -2,6 +2,7 @@ from typing import Any, Dict, List, Optional import logging import re from src.constants import MAX_READ_CHARS +from src.tool_approvals import document_content_digest from src.tool_utils import _parse_tool_args, get_upload_handler from src.upload_handler import reserve_upload_references @@ -82,14 +83,27 @@ def _most_recent_owned_document(db, Document, owner: Optional[str], active_only: def _approved_document_version_error(doc: Any, ctx: dict) -> Optional[Dict]: """Reject a sealed document action when its target changed meanwhile.""" - expected = ctx.get("expected_document_version") - if expected is None: + expected_version = ctx.get("expected_document_version") + expected_digest = ( + str(ctx.get("expected_document_digest") or "").strip().lower() + ) + if expected_version is None and not expected_digest: return None try: - unchanged = int(getattr(doc, "version_count", -1)) == int(expected) + version_unchanged = ( + expected_version is None + or int(getattr(doc, "version_count", -1)) == int(expected_version) + ) except (TypeError, ValueError): - unchanged = False - if unchanged: + version_unchanged = False + content_unchanged = True + if expected_digest: + content_unchanged = ( + doc is not None + and document_content_digest(getattr(doc, "current_content", "")) + == expected_digest + ) + if version_unchanged and content_unchanged: return None return { "error": ( diff --git a/src/tool_approvals.py b/src/tool_approvals.py index d82b651d1..8b7fd3d63 100644 --- a/src/tool_approvals.py +++ b/src/tool_approvals.py @@ -43,6 +43,11 @@ def _canonical_digest(payload: dict[str, Any]) -> str: return hashlib.sha256(encoded).hexdigest() +def document_content_digest(content: Any) -> str: + """Return the stable server-side fingerprint used to seal a document.""" + return hashlib.sha256(str(content or "").encode("utf-8")).hexdigest() + + def _binding_payload( *, owner: Any, @@ -53,6 +58,7 @@ def _binding_payload( workspace: Any, document_id: Any, document_version: Any, + document_digest: Any, external_untrusted_context_seen: bool, effects: tuple[str, ...], result_integrity: str, @@ -68,6 +74,7 @@ def _binding_payload( "document_version": ( int(document_version) if document_version is not None else None ), + "document_digest": str(document_digest or "").strip().lower(), "external_untrusted_context_seen": bool(external_untrusted_context_seen), "effects": list(effects), "result_integrity": str(result_integrity), @@ -85,6 +92,7 @@ class PendingToolApproval: workspace: str document_id: str document_version: int | None + document_digest: str external_untrusted_context_seen: bool effects: tuple[str, ...] result_integrity: str @@ -163,6 +171,7 @@ class ExactToolApproval: workspace=workspace, document_id=self.pending.document_id, document_version=self.pending.document_version, + document_digest=self.pending.document_digest, external_untrusted_context_seen=( self.pending.external_untrusted_context_seen ), @@ -245,6 +254,7 @@ class ToolApprovalStore: workspace: Any, document_id: Any = None, document_version: Any = None, + document_digest: Any = None, external_untrusted_context_seen: bool, capabilities: ToolCapabilities, ) -> PendingToolApproval: @@ -260,6 +270,7 @@ class ToolApprovalStore: workspace=workspace, document_id=document_id, document_version=document_version, + document_digest=document_digest, external_untrusted_context_seen=external_untrusted_context_seen, effects=effects, result_integrity=result_integrity, @@ -274,6 +285,7 @@ class ToolApprovalStore: workspace=payload["workspace"], document_id=payload["document_id"], document_version=payload["document_version"], + document_digest=payload["document_digest"], external_untrusted_context_seen=payload[ "external_untrusted_context_seen" ], diff --git a/src/tool_execution.py b/src/tool_execution.py index f33f8c3fa..8c0c83032 100644 --- a/src/tool_execution.py +++ b/src/tool_execution.py @@ -570,6 +570,7 @@ async def _document_tool_dispatch( owner: Optional[str] = None, document_id: Optional[str] = None, document_version: Optional[int] = None, + document_digest: Optional[str] = None, ) -> Optional[Dict]: """Route a document tool through TOOL_HANDLERS with the right ctx shape.""" from src.agent_tools import TOOL_HANDLERS @@ -578,6 +579,7 @@ async def _document_tool_dispatch( "owner": owner, "doc_id": document_id, "expected_document_version": document_version, + "expected_document_digest": document_digest, } if tool in TOOL_HANDLERS: return await TOOL_HANDLERS[tool](content, ctx) @@ -645,6 +647,7 @@ async def execute_tool_block( and ( not exact_approval.pending.document_id or exact_approval.pending.document_version is None + or not exact_approval.pending.document_digest ) ): return ( @@ -725,6 +728,11 @@ async def execute_tool_block( if approval_claimed else None ), + approved_document_digest=( + exact_approval.pending.document_digest + if approval_claimed + else None + ), ) if isinstance(security_context, ToolRunSecurityContext): security_context.observe_tool_result( @@ -746,6 +754,7 @@ async def _execute_tool_block_impl( tool_policy: Optional[Any] = None, approved_document_id: Optional[str] = None, approved_document_version: Optional[int] = None, + approved_document_digest: Optional[str] = None, ) -> Tuple[str, Dict]: """Execute a single tool block. Returns (description, result_dict). @@ -914,6 +923,7 @@ async def _execute_tool_block_impl( owner, document_id=approved_document_id, document_version=approved_document_version, + document_digest=approved_document_digest, ) \ or {"error": f"{tool}: execution failed", "exit_code": 1} if tool in ("edit_document", "suggest_document") and "title" in (result or {}): diff --git a/tests/test_external_context_tool_gate.py b/tests/test_external_context_tool_gate.py index c8906e4d2..a7ec73420 100644 --- a/tests/test_external_context_tool_gate.py +++ b/tests/test_external_context_tool_gate.py @@ -1000,6 +1000,75 @@ def test_tainted_document_edit_without_active_target_cannot_be_approved(monkeypa assert "ask_user" not in blocked[0] +def test_tainted_document_approval_seals_current_content(monkeypatch): + from types import SimpleNamespace + + from src.prompt_security import untrusted_context_message + from src.tool_approvals import document_content_digest + + import src.agent_loop as agent_loop + + monkeypatch.setattr( + agent_loop, + "get_setting", + lambda key, default=None: default, + raising=False, + ) + monkeypatch.setattr(agent_loop, "get_mcp_manager", lambda: None, raising=False) + monkeypatch.setattr(agent_loop, "estimate_tokens", lambda *args, **kwargs: 10) + + async def fake_stream(*args, **kwargs): + yield "data: " + json.dumps({ + "delta": "```update_document\nreplacement\n```", + }) + "\n\n" + yield "data: [DONE]\n\n" + + async def should_not_execute(*args, **kwargs): + raise AssertionError("unapproved document edit reached executor") + + monkeypatch.setattr(agent_loop, "stream_llm_with_fallback", fake_stream) + monkeypatch.setattr(agent_loop, "execute_tool_block", should_not_execute) + active_document = SimpleNamespace( + id="document-7", + title="Draft", + language="markdown", + current_content="original", + version_count=4, + ) + events = _collect_agent_events( + agent_loop.stream_agent_loop( + "http://local.test/v1", + "small-local-model", + [ + {"role": "user", "content": "update this document"}, + untrusted_context_message("stored context", "untrusted"), + ], + active_document=active_document, + session_id="document-approval-session", + owner="alice", + max_rounds=1, + relevant_tools={"update_document"}, + ) + ) + + approval = next( + event["ask_user"] + for event in events + if event.get("ask_user", {}).get("kind") == "tool_approval" + ) + pending = agent_loop.tool_approval_store.peek(approval["approval_id"]) + assert pending is not None + assert pending.document_id == "document-7" + assert pending.document_version == 4 + assert pending.document_digest == document_content_digest("original") + agent_loop.tool_approval_store.consume( + pending.approval_id, + decision="deny", + owner="alice", + session_id="document-approval-session", + ) + + def test_approval_pause_does_not_trigger_teacher_takeover(monkeypatch): from src.prompt_security import untrusted_context_message diff --git a/tests/test_foreground_model_routing.py b/tests/test_foreground_model_routing.py index 481d3d32e..e4928f86c 100644 --- a/tests/test_foreground_model_routing.py +++ b/tests/test_foreground_model_routing.py @@ -16,6 +16,7 @@ import routes.chat_routes as chat_routes import routes.chat_helpers as chat_helpers import routes.prefs_routes as prefs_routes from src.request_models import ChatRequest +from src.tool_approvals import document_content_digest from src.foreground_model_routing import ( FOREGROUND_AVAILABILITY_STATUSES, MAX_FOREGROUND_FALLBACKS, @@ -276,6 +277,7 @@ async def test_chat_stream_consumes_exact_tool_approval_for_own_session(monkeypa workspace=None, document_id="document-7", document_version=4, + document_digest=document_content_digest("original"), external_untrusted_context_seen=True, capabilities=capabilities_for_action("update_document", tool_content), ) diff --git a/tests/test_tool_approvals.py b/tests/test_tool_approvals.py index cce3f1ac4..2e13ffa32 100644 --- a/tests/test_tool_approvals.py +++ b/tests/test_tool_approvals.py @@ -5,7 +5,7 @@ from collections import namedtuple import pytest -from src.tool_approvals import ToolApprovalStore +from src.tool_approvals import ToolApprovalStore, document_content_digest from src.tool_capabilities import ToolRunSecurityContext, capabilities_for_action @@ -121,6 +121,7 @@ def test_public_payload_shows_complete_action_but_not_authority_fields(): content="printf safe\nSECOND_LINE", document_id="document-7", document_version=4, + document_digest=document_content_digest("original"), ) payload = pending.public_payload() @@ -184,6 +185,7 @@ async def test_dispatcher_uses_sealed_document_target(monkeypatch): content=content, document_id="document-7", document_version=4, + document_digest=document_content_digest("original"), capabilities=capabilities_for_action("update_document", content), ) grant = store.consume( @@ -199,6 +201,7 @@ async def test_dispatcher_uses_sealed_document_target(monkeypatch): ( kwargs.get("approved_document_id"), kwargs.get("approved_document_version"), + kwargs.get("approved_document_digest"), ) ) return "update_document", {"output": "ok", "exit_code": 0} @@ -220,7 +223,9 @@ async def test_dispatcher_uses_sealed_document_target(monkeypatch): ) assert result["exit_code"] == 0 - assert captured == [("document-7", 4)] + assert captured == [ + ("document-7", 4, document_content_digest("original")) + ] @pytest.mark.asyncio @@ -268,7 +273,11 @@ async def test_dispatcher_rejects_approved_document_action_without_target(monkey def test_approved_document_version_guard_rejects_changed_target(): from src.agent_tools.document_tools import _approved_document_version_error - doc = type("Document", (), {"version_count": 5})() + doc = type( + "Document", + (), + {"version_count": 5, "current_content": "original"}, + )() assert _approved_document_version_error( doc, @@ -276,8 +285,18 @@ def test_approved_document_version_guard_rejects_changed_target(): )["document_changed"] is True assert _approved_document_version_error( doc, - {"expected_document_version": 5}, + { + "expected_document_version": 5, + "expected_document_digest": document_content_digest("original"), + }, ) is None + assert _approved_document_version_error( + doc, + { + "expected_document_version": 5, + "expected_document_digest": document_content_digest("changed"), + }, + )["document_changed"] is True assert _approved_document_version_error( None, {"expected_document_version": 5},