From 1f216cfd0e519b2c78451de0dac76d9ff89fe8d4 Mon Sep 17 00:00:00 2001
From: RaresKeY <158580472+RaresKeY@users.noreply.github.com>
Date: Sat, 15 Aug 2026 04:13:56 +0000
Subject: [PATCH] fix(agent): taint stored document tool results

---
 src/tool_capabilities.py                 | 10 ++++++++--
 tests/test_external_context_tool_gate.py | 17 +++++++++++++++++
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/src/tool_capabilities.py b/src/tool_capabilities.py
index b80f515ea..121ff54db 100644
--- a/src/tool_capabilities.py
+++ b/src/tool_capabilities.py
@@ -132,7 +132,6 @@ _register(
         "create_session",
         "draft_email",
         "draft_email_reply",
-        "edit_document",
         "manage_calendar",
         "manage_contact",
         "manage_documents",
@@ -144,10 +143,17 @@ _register(
         "manage_tasks",
         "suggest_document",
         "todowrite",
-        "update_document",
     },
     ToolEffect.WRITE_PRIVATE,
 )
+_register(
+    {"edit_document", "update_document"},
+    ToolEffect.WRITE_PRIVATE,
+    # These tools can echo stored document content that was not present in
+    # their arguments.  edit_document returns the complete edited document;
+    # update_document also preserves stored email headers/thread history.
+    result_integrity=ResultIntegrity.EXTERNAL_UNTRUSTED,
+)
 _register(
     {"pipeline"},
     ToolEffect.NETWORK_EGRESS,
diff --git a/tests/test_external_context_tool_gate.py b/tests/test_external_context_tool_gate.py
index ed9355a06..055a1b206 100644
--- a/tests/test_external_context_tool_gate.py
+++ b/tests/test_external_context_tool_gate.py
@@ -341,6 +341,22 @@ def test_cross_model_results_taint_before_later_host_actions(tool_name):
     assert context.decision_for("bash").allowed is False
 
 
+@pytest.mark.parametrize("tool_name", ["edit_document", "update_document"])
+def test_stored_document_results_taint_before_later_host_actions(tool_name):
+    context = ToolRunSecurityContext()
+
+    capabilities = capabilities_for_tool(tool_name)
+    assert capabilities.result_integrity is ResultIntegrity.EXTERNAL_UNTRUSTED
+    context.observe_tool_result(
+        tool_name,
+        {"content": "stored attacker-controlled content", "exit_code": 0},
+        "model-proposed replacement",
+    )
+
+    assert context.external_untrusted_context_seen is True
+    assert context.decision_for("bash").allowed is False
+
+
 @pytest.mark.parametrize(
     "tool_name,content",
     [
@@ -403,6 +419,7 @@ def test_ambiguous_private_manager_action_fails_high():
         ("web_search", {"output": "external", "exit_code": 0}, True),
         ("web_search", {"error": "offline", "exit_code": 1}, False),
         ("list_served_models", {"output": "local status", "exit_code": 0}, False),
+        ("edit_document", {"content": "stored content", "exit_code": 0}, True),
     ],
 )
 def test_result_folding_is_transport_and_status_consistent(