fix(agent): seal document approval content

This commit is contained in:
RaresKeY
2026-08-15 07:01:36 +00:00
parent 2b72531eaa
commit 7a138e8a3f
7 changed files with 151 additions and 10 deletions
+16 -1
View File
@@ -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
),
+19 -5
View File
@@ -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": (
+12
View File
@@ -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"
],
+10
View File
@@ -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 {}):
+69
View File
@@ -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
+2
View File
@@ -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),
)
+23 -4
View File
@@ -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},