fix(agent): harden approval lifecycle

This commit is contained in:
RaresKeY
2026-08-15 06:52:44 +00:00
parent 58b2a4bfa9
commit 2b72531eaa
25 changed files with 782 additions and 116 deletions
+184 -1
View File
@@ -499,6 +499,63 @@ def test_private_manager_write_aliases_keep_write_effect(tool_name, content):
assert capabilities.result_integrity is ResultIntegrity.EXTERNAL_UNTRUSTED
@pytest.mark.parametrize(
"tool_name,content",
[
("manage_calendar", '{"action":"delete_event"}'),
("manage_contact", '{"action":"delete"}'),
("manage_documents", '{"action":"tidy"}'),
("manage_endpoints", '{"action":"delete"}'),
("manage_bg_jobs", '{"action":"kill","job_id":"job-1"}'),
("manage_memory", "delete\nmemory-id"),
("manage_mcp", '{"action":"delete"}'),
("manage_notes", '{"action":"delete"}'),
("manage_research", '{"action":"delete"}'),
("manage_session", "truncate\nsession-id\n10"),
("manage_settings", '{"action":"reset","key":"theme"}'),
("manage_skills", '{"action":"delete"}'),
("manage_tasks", '{"action":"delete"}'),
("manage_tokens", '{"action":"delete"}'),
("manage_webhooks", '{"action":"delete"}'),
],
)
def test_multiplexed_destructive_actions_disclose_destructive_effect(
tool_name,
content,
):
capabilities = capabilities_for_action(tool_name, content)
assert any(
effect in capabilities.effects
for effect in (
ToolEffect.WRITE_PRIVATE,
ToolEffect.ADMIN_CHANGE,
ToolEffect.EXECUTE_CODE,
)
)
assert ToolEffect.DESTRUCTIVE in capabilities.effects
@pytest.mark.parametrize(
"tool_name,content",
[
("manage_bg_jobs", '{"action":"output","job_id":"job-1"}'),
("manage_endpoints", '{"action":"list"}'),
("manage_mcp", '{"action":"reconnect"}'),
("manage_settings", '{"action":"set","key":"theme","value":"dark"}'),
("manage_tokens", '{"action":"create","name":"automation"}'),
("manage_webhooks", '{"action":"disable"}'),
],
)
def test_multiplexed_non_destructive_actions_do_not_claim_destructive_effect(
tool_name,
content,
):
capabilities = capabilities_for_action(tool_name, content)
assert ToolEffect.DESTRUCTIVE not in capabilities.effects
def test_ambiguous_private_manager_action_fails_high():
capabilities = capabilities_for_action("manage_notes", "not json")
@@ -894,6 +951,100 @@ def test_tainted_native_route_keeps_action_schema_for_exact_approval(monkeypatch
assert "update_document" in seen_tools
def test_tainted_document_edit_without_active_target_cannot_be_approved(monkeypatch):
from src.prompt_security import untrusted_context_message
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("unsealed document edit reached executor")
monkeypatch.setattr(agent_loop, "stream_llm_with_fallback", fake_stream)
monkeypatch.setattr(agent_loop, "execute_tool_block", should_not_execute)
events = _collect_agent_events(
agent_loop.stream_agent_loop(
"http://local.test/v1",
"small-local-model",
[
{"role": "user", "content": "update a document"},
untrusted_context_message("stored context", "untrusted"),
],
max_rounds=1,
relevant_tools={"update_document"},
)
)
blocked = [
event
for event in events
if event.get("type") == "tool_output"
and event.get("tool") == "update_document"
]
assert blocked
assert "Open the exact document" in blocked[0]["output"]
assert "ask_user" not in blocked[0]
def test_approval_pause_does_not_trigger_teacher_takeover(monkeypatch):
from src.prompt_security import untrusted_context_message
import src.agent_loop as agent_loop
import src.teacher_escalation as teacher_escalation
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": "```bash\nprintf paused\n```"}) + "\n\n"
yield "data: [DONE]\n\n"
async def fail_teacher(*args, **kwargs):
raise AssertionError("approval pause reached teacher takeover")
yield # pragma: no cover
monkeypatch.setattr(agent_loop, "stream_llm_with_fallback", fake_stream)
monkeypatch.setattr(teacher_escalation, "run_teacher_inline", fail_teacher)
events = _collect_agent_events(
agent_loop.stream_agent_loop(
"http://local.test/v1",
"small-local-model",
[
{"role": "user", "content": "run it"},
untrusted_context_message("stored context", "untrusted"),
],
session_id="session-1",
max_rounds=1,
relevant_tools={"bash"},
)
)
assert any(
event.get("ask_user", {}).get("kind") == "tool_approval"
for event in events
)
def test_frontend_tool_approval_uses_opaque_id_and_fixed_decisions():
root = Path(__file__).parents[1]
chat = (root / "static/js/chat.js").read_text()
@@ -910,13 +1061,45 @@ def test_frontend_tool_approval_uses_opaque_id_and_fixed_decisions():
assert "if (isStreaming || _sendInFlight)" in chat
assert "_submitToolApprovalWhenIdle" in chat
assert "input.dispatchEvent(new Event('input'" in chat
assert "_pendingToolApproval.draft = input.value" in chat
assert "const approvalForSend = _pendingToolApproval" in chat
assert "!approvalForSend && fileHandlerModule.getPendingCount()" in chat
assert "if (!approvalForSend) _pendingRegenAttachments = null" in chat
assert "!approvalForSend && el('research-toggle').checked" in chat
assert "approvalForSend ? (approvalForSend.draft || '') : ''" in chat
assert "if (approvalForSend && documentSaved === false)" in chat
assert "if (!approvalForSend) {\n try {\n _sendPerf.mark('doc_silent_save_begin')" in chat
assert "document_id: aq.action && aq.action.document_id" in renderer
assert "const firstRound = (toolsByRound[0] || []).length ? 0 : 1" in renderer
assert "const r = ev.round ?? 1" in renderer
assert "/test-approval`" in skills
assert "approval_id: approval.approval_id" in skills
assert "['approve', 'Allow once'" in skills
assert index.count("app.js?v=20260815toolapproval3") == 2
assert index.count("app.js?v=20260815toolapproval4") == 2
assert "app.js?v=20260808startupshell1" not in index
approval_module_sources = [
(root / path).read_text()
for path in (
"static/app.js",
"static/index.html",
"static/js/chat.js",
"static/js/chatRenderer.js",
"static/js/chatStream.js",
"static/js/document.js",
"static/js/emailInbox.js",
"static/js/emailLibrary.js",
"static/js/settings.js",
"static/js/slashCommands.js",
)
]
assert all(
"20260722emailfastindex1" not in source
for source in approval_module_sources
)
assert all(
"20260815approvalsave1" in source
for source in approval_module_sources
)
def test_frontend_raw_fences_do_not_call_document_mutators():
+45
View File
@@ -93,6 +93,7 @@ def _chat_stream_endpoint(
agent_chunks=None,
chat_chunks=None,
capture_completion=False,
capture_context=False,
endpoint_url="https://selected.example/v1",
):
def add_message(message):
@@ -136,6 +137,8 @@ def _chat_stream_endpoint(
)
async def fake_build_context(*args, **kwargs):
if capture_context:
captured["build_context"] = kwargs
return context
async def fake_chat_stream(candidates, messages, **kwargs):
@@ -336,6 +339,48 @@ async def test_chat_stream_approval_restores_exact_shell_turn_toggle(monkeypatch
assert "bash" not in captured["approval_disabled_tools"]
@pytest.mark.asyncio
async def test_chat_stream_approval_ignores_research_and_new_attachments(monkeypatch):
from src.tool_capabilities import capabilities_for_action
captured = {}
endpoint = _chat_stream_endpoint(
monkeypatch,
"agent",
captured,
capture_context=True,
)
monkeypatch.setattr(chat_routes, "get_session_mode", lambda _session_id: "research_pending")
pending = chat_routes.tool_approval_store.create(
owner="alice",
session_id="session-1",
origin_run_id="run-1",
tool_name="bash",
content="printf exact",
workspace=None,
external_untrusted_context_seen=True,
capabilities=capabilities_for_action("bash", "printf exact"),
)
request = _RouteRequest("agent")
request._form.update(
{
"attachments": '["unrelated-upload"]',
"use_research": "true",
"tool_approval_id": pending.approval_id,
"tool_approval_decision": "approve",
}
)
response = await endpoint(request)
async for _ in response.body_iterator:
pass
assert captured["exact_approval"].pending == pending
assert captured["build_context"]["att_ids"] == []
assert "agent" in captured
assert "chat" not in captured
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ["chat", "agent"])
@pytest.mark.parametrize("endpoint_url", ["", None])
+46
View File
@@ -12,6 +12,7 @@ Three focused tests:
from __future__ import annotations
import json
from datetime import datetime, timezone
from types import SimpleNamespace
@@ -68,6 +69,51 @@ async def test_scheduler_agent_loop_path(monkeypatch):
assert msgs[2]["content"] == "run the digest"
async def test_scheduler_retires_unattended_exact_approval(monkeypatch):
from src.task_scheduler import TaskScheduler
from src.tool_approvals import tool_approval_store
from src.tool_capabilities import capabilities_for_action
pending = tool_approval_store.create(
owner="admin",
session_id="s",
origin_run_id="scheduled-run",
tool_name="bash",
content="printf exact",
workspace=None,
external_untrusted_context_seen=True,
capabilities=capabilities_for_action("bash", "printf exact"),
)
approval = pending.public_payload()
async def fake_stream_agent_loop(*args, **kwargs):
yield "data: " + json.dumps({
"type": "tool_output",
"tool": "bash",
"output": "Waiting for an exact user approval.",
"ask_user": approval,
}) + "\n\n"
monkeypatch.setattr(
"src.agent_loop.stream_agent_loop",
fake_stream_agent_loop,
)
monkeypatch.setattr(
"src.task_endpoint.resolve_task_candidates",
lambda **kwargs: [],
)
result = await TaskScheduler(session_manager=None)._run_agent_loop(
"http://ep/v1",
"model",
_make_task(),
"s",
)
assert "paused safely" in result
assert "That action was not executed" in result
assert tool_approval_store.peek(pending.approval_id) is None
# ---------------------------------------------------------------------------
# Test 2 — fallback path receives the same datetime context
# ---------------------------------------------------------------------------
+1 -1
View File
@@ -23,7 +23,7 @@ _IMPORT_REWRITES = {
"import uiModule, { autoResize, styledPrompt } from './ui.js';": (
"import uiModule, { autoResize, styledPrompt } from './ui.mjs';"
),
"import chatRenderer from './chatRenderer.js?v=20260815toolapproval3';": (
"import chatRenderer from './chatRenderer.js?v=20260815toolapproval4';": (
"import chatRenderer from './chatRenderer.mjs';"
),
"import { providerLogo } from './providers.js';": (
+82
View File
@@ -1,4 +1,5 @@
import asyncio
import json
from types import SimpleNamespace
import pytest
@@ -217,6 +218,87 @@ async def test_run_teacher_inline_triggers_tier2_escalation(monkeypatch):
assert any("skill_saved" in evt for evt in events)
@pytest.mark.asyncio
async def test_teacher_approval_keeps_parent_authority_and_skips_skill_save(
monkeypatch,
):
monkeypatch.setattr(
"src.settings.get_setting",
lambda key, default=None: {
"teacher_enabled": True,
"teacher_model": "teacher-model",
}.get(key, default),
)
monkeypatch.setattr(
"src.ai_interaction._resolve_model",
lambda spec, owner=None: (
"http://teacher.local/v1",
"teacher-model",
{},
),
)
monkeypatch.setattr(
"src.teacher_escalation.evaluate_turn_regex",
lambda *args: ("failure", "student failed"),
)
captured = {}
approval = {
"kind": "tool_approval",
"approval_id": "opaque-id",
"question": "Allow this exact action once?",
}
async def fake_stream_agent_loop(*args, **kwargs):
captured.update(kwargs)
yield "data: " + json.dumps({
"type": "tool_output",
"tool": "bash",
"output": "Waiting for an exact user approval.",
"ask_user": approval,
}) + "\n\n"
yield "data: [DONE]\n\n"
async def fail_skill_distillation(*args, **kwargs):
raise AssertionError("paused teacher trace was distilled into a skill")
monkeypatch.setattr(
"src.agent_loop.stream_agent_loop",
fake_stream_agent_loop,
)
monkeypatch.setattr(
"src.teacher_escalation._call_teacher",
fail_skill_distillation,
)
active_document = object()
active_email = {"uid": "email-1"}
policy = object()
events = []
async for evt in teacher_escalation.run_teacher_inline(
student_endpoint_url="http://student.local/v1",
student_messages=[{"role": "user", "content": "test request"}],
student_tool_events=[],
student_reply="student reply",
owner="alice",
session_id="session-1",
workspace="/workspace",
disabled_tools={"web_fetch"},
tool_policy=policy,
active_document=active_document,
active_email=active_email,
):
events.append(evt)
assert captured["session_id"] == "session-1"
assert captured["workspace"] == "/workspace"
assert captured["disabled_tools"] == {"web_fetch"}
assert captured["tool_policy"] is policy
assert captured["active_document"] is active_document
assert captured["active_email"] == active_email
assert any("opaque-id" in event for event in events)
assert not any("skill_saved" in event for event in events)
@pytest.mark.asyncio
async def test_run_teacher_inline_tier2_disabled_by_default(monkeypatch):
# Settings and gates (Tier 2 disabled)
+88 -3
View File
@@ -61,10 +61,9 @@ def test_approval_is_bound_to_exact_action_and_claimed_once():
)
def test_wrong_owner_and_deny_destructively_consume_pending_action():
def test_wrong_owner_cannot_consume_but_deny_retires_pending_action():
store = ToolApprovalStore()
wrong_owner = _pending(store)
denied = _pending(store)
assert store.consume(
wrong_owner.approval_id,
@@ -72,7 +71,9 @@ def test_wrong_owner_and_deny_destructively_consume_pending_action():
owner="mallory",
session_id="session-1",
) is None
assert store.peek(wrong_owner.approval_id) is None
assert store.peek(wrong_owner.approval_id) == wrong_owner
denied = _pending(store)
assert store.consume(
denied.approval_id,
decision="deny",
@@ -222,6 +223,48 @@ async def test_dispatcher_uses_sealed_document_target(monkeypatch):
assert captured == [("document-7", 4)]
@pytest.mark.asyncio
async def test_dispatcher_rejects_approved_document_action_without_target(monkeypatch):
import src.tool_execution as tool_execution
store = ToolApprovalStore()
content = "replacement"
pending = _pending(
store,
tool_name="update_document",
content=content,
capabilities=capabilities_for_action("update_document", content),
)
grant = store.consume(
pending.approval_id,
decision="approve",
owner="alice",
session_id="session-1",
)
async def should_not_run(*args, **kwargs):
raise AssertionError("unsealed document target reached implementation")
monkeypatch.setattr(
tool_execution,
"_execute_tool_block_impl",
should_not_run,
)
_, result = await tool_execution.execute_tool_block(
ToolBlock("update_document", content),
session_id="session-1",
owner="alice",
workspace=None,
security_context=ToolRunSecurityContext(
external_untrusted_context_seen=True
),
exact_approval=grant,
)
assert result["blocked"] is True
assert result["policy"] == "exact_tool_approval"
def test_approved_document_version_guard_rejects_changed_target():
from src.agent_tools.document_tools import _approved_document_version_error
@@ -235,6 +278,48 @@ def test_approved_document_version_guard_rejects_changed_target():
doc,
{"expected_document_version": 5},
) is None
assert _approved_document_version_error(
None,
{"expected_document_version": 5},
)["document_changed"] is True
@pytest.mark.asyncio
async def test_missing_sealed_document_does_not_fall_back_to_another(monkeypatch):
import src.agent_tools.document_tools as document_tools
class FakeDb:
def close(self):
pass
def rollback(self):
pass
monkeypatch.setattr("src.database.SessionLocal", lambda: FakeDb())
monkeypatch.setattr(
document_tools,
"_get_owned_document",
lambda *args, **kwargs: None,
)
def fail_fallback(*args, **kwargs):
raise AssertionError("sealed target fell back to a different document")
monkeypatch.setattr(
document_tools,
"_most_recent_owned_document",
fail_fallback,
)
result = await document_tools.UpdateDocumentTool().execute(
"replacement",
{
"doc_id": "deleted-document",
"expected_document_version": 4,
"owner": "alice",
},
)
assert result["document_changed"] is True
@pytest.mark.asyncio