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
+64 -34
View File
@@ -5629,40 +5629,64 @@ async def stream_agent_loop(
in {"edit_document", "suggest_document", "update_document"}
else None
)
pending_approval = tool_approval_store.create(
owner=owner,
session_id=session_id,
origin_run_id=run_security.run_id,
tool_name=block.tool_type,
content=block.content,
workspace=workspace,
document_id=getattr(approval_document, "id", None),
document_version=getattr(
approval_document,
"version_count",
None,
),
external_untrusted_context_seen=(
run_security.external_untrusted_context_seen
),
capabilities=capabilities_for_action(
if (
block.tool_type
in {"edit_document", "suggest_document", "update_document"}
and (
approval_document is None
or getattr(approval_document, "id", None) is None
or getattr(approval_document, "version_count", None) is None
)
):
# These legacy tools otherwise fall back to a process-global
# or most-recent document at dispatch time. That target can
# change while an approval card is pending, so there is no
# exact action to seal until the user opens a real document.
desc = f"{block.tool_type}: BLOCKED"
result = {
"error": (
"Open the exact document to edit, then request this "
"action again so its id and version can be sealed."
),
"exit_code": 1,
"blocked": True,
"policy": "exact_tool_approval_target",
}
else:
pending_approval = tool_approval_store.create(
owner=owner,
session_id=session_id,
origin_run_id=run_security.run_id,
tool_name=block.tool_type,
content=block.content,
workspace=workspace,
document_id=getattr(approval_document, "id", None),
document_version=getattr(
approval_document,
"version_count",
None,
),
external_untrusted_context_seen=(
run_security.external_untrusted_context_seen
),
capabilities=capabilities_for_action(
block.tool_type,
block.content,
),
)
desc = f"{block.tool_type}: APPROVAL REQUIRED"
result = {
"output": "Waiting for an exact user approval.",
"exit_code": None,
"approval_required": True,
"ask_user": pending_approval.public_payload(
reason=security_decision.reason,
),
}
logger.info(
"Exact approval required before tool start: %s",
block.tool_type,
block.content,
),
)
desc = f"{block.tool_type}: APPROVAL REQUIRED"
result = {
"output": "Waiting for an exact user approval.",
"exit_code": None,
"approval_required": True,
"ask_user": pending_approval.public_payload(
reason=security_decision.reason,
),
}
logger.info(
"Exact approval required before tool start: %s",
block.tool_type,
)
)
elif tool_policy and tool_policy.blocks(block.tool_type) and not _ody_clamped_tool_allowed:
desc = f"{block.tool_type}: BLOCKED"
result = {
@@ -6325,7 +6349,7 @@ async def stream_agent_loop(
# gets a turn (with its own tool calls forwarded to the user) and
# a skill is saved ONLY if the teacher actually succeeds. Skipped
# when we ARE the teacher to avoid recursion.
if not _is_teacher_run and not guide_only:
if not _is_teacher_run and not guide_only and not _awaiting_user:
try:
from src.teacher_escalation import run_teacher_inline
async for evt in run_teacher_inline(
@@ -6334,6 +6358,12 @@ async def stream_agent_loop(
student_tool_events=tool_events,
student_reply=full_response,
owner=owner,
session_id=session_id,
workspace=workspace,
disabled_tools=disabled_tools,
tool_policy=tool_policy,
active_document=active_document,
active_email=active_email,
):
yield evt
except Exception as _esc_err:
+12
View File
@@ -475,6 +475,12 @@ class UpdateDocumentTool:
doc = None
if target_id:
doc = _get_owned_document(db, Document, target_id, owner)
if (
not doc
and target_id
and ctx.get("expected_document_version") is not None
):
return _approved_document_version_error(None, ctx)
if not doc:
doc = _most_recent_owned_document(db, Document, owner)
if doc:
@@ -555,6 +561,12 @@ class EditDocumentTool:
doc = None
if target_id:
doc = _get_owned_document(db, Document, target_id, owner)
if (
not doc
and target_id
and ctx.get("expected_document_version") is not None
):
return _approved_document_version_error(None, ctx)
if not doc:
# Fallback: most recently updated document. Avoids "no active doc" errors
# after server restart or when the agent loses track of which doc to edit.
+36
View File
@@ -1883,6 +1883,7 @@ class TaskScheduler:
pass
full_text = ""
tool_results = []
approval_pause = None
# Honor per-task max_steps (defense against runaway agent loops).
# Falls back to 20 if not set — the historical default.
@@ -1929,9 +1930,44 @@ class TaskScheduler:
tool_summary = data.get("stdout") or data.get("output") or data.get("result") or ""
if isinstance(tool_summary, str) and tool_summary.strip():
tool_results.append(f"[{data.get('tool', '?')}] {tool_summary[:500]}")
approval = data.get("ask_user")
if (
isinstance(approval, dict)
and approval.get("kind") == "tool_approval"
):
approval_pause = {
"tool": data.get("tool") or "tool",
"approval_id": approval.get("approval_id"),
}
# Scheduled tasks have no interactive surface that
# can safely resume a one-use grant. Retire the
# record immediately instead of leaving it pending
# and report an explicit manual-action boundary.
try:
from src.tool_approvals import tool_approval_store
tool_approval_store.consume(
approval_pause["approval_id"],
decision="deny",
owner=task.owner,
session_id=session_id,
)
except Exception:
logger.debug(
"Could not retire scheduled-task approval",
exc_info=True,
)
break
except (json.JSONDecodeError, KeyError):
pass
if approval_pause is not None:
return (
"Scheduled task paused safely: "
f"{approval_pause['tool']} requested an exact action after "
"untrusted context. That action was not executed. Run this task "
"interactively to inspect and approve the action."
)
# Grace summarization — if the model exhausted rounds on tool calls
# without producing a final text response, do one last LLM call
# asking it to summarize what it did. Guarantees output.
+23 -2
View File
@@ -563,6 +563,12 @@ async def run_teacher_inline(
student_tool_events: List[Dict[str, Any]],
student_reply: str,
owner: Optional[str] = None,
session_id: Optional[str] = None,
workspace: Optional[str] = None,
disabled_tools: Optional[set[str]] = None,
tool_policy: Any = None,
active_document: Any = None,
active_email: Optional[Dict[str, str]] = None,
):
"""Async generator. Yields SSE event strings.
@@ -668,6 +674,12 @@ async def run_teacher_inline(
messages=teacher_messages,
headers=teacher_headers,
owner=owner,
session_id=session_id,
workspace=workspace,
disabled_tools=disabled_tools,
tool_policy=tool_policy,
active_document=active_document,
active_email=active_email,
_is_teacher_run=True,
):
# Swallow teacher's own [DONE] — outer loop emits the real one
@@ -683,12 +695,15 @@ async def run_teacher_inline(
payload["teacher"] = True
typ = payload.get("type")
if typ == "tool_output":
captured_tool_events.append({
captured_tool_event = {
"tool": payload.get("tool"),
"command": payload.get("command"),
"output": payload.get("output"),
"exit_code": payload.get("exit_code"),
})
}
if isinstance(payload.get("ask_user"), dict):
captured_tool_event["ask_user"] = payload["ask_user"]
captured_tool_events.append(captured_tool_event)
if "delta" in payload and isinstance(payload["delta"], str):
if payload.get("thinking"):
continue
@@ -697,6 +712,12 @@ async def run_teacher_inline(
continue
yield evt_str
# A takeover that paused for a question or exact action has not completed
# yet. Its server-owned approval card is already in the live/persisted tool
# events; do not evaluate the partial trace or distill it into a skill.
if any(event.get("ask_user") for event in captured_tool_events):
return
teacher_text = "".join(captured_text_parts).strip()
t_status, t_reason = evaluate_turn_regex(captured_tool_events, teacher_text)
if t_status == "failure":
+13 -8
View File
@@ -323,14 +323,19 @@ class ToolApprovalStore:
now = time.time()
with self._lock:
self._purge_expired_locked(now)
pending = self._pending.pop(str(approval_id or ""), None)
if pending is None:
return None
if (
pending.owner != _normalized_owner(owner)
or pending.session_id != str(session_id or "")
):
return None
approval_key = str(approval_id or "")
pending = self._pending.get(approval_key)
if pending is None:
return None
if (
pending.owner != _normalized_owner(owner)
or pending.session_id != str(session_id or "")
):
# Authentication is checked before destructive consumption so
# a leaked/guessed opaque id cannot be used to invalidate
# another owner's pending action.
return None
self._pending.pop(approval_key, None)
if str(decision or "").strip().lower() != "approve":
return None
return ExactToolApproval(pending)
+34 -2
View File
@@ -332,6 +332,26 @@ _PRIVATE_ACTION_WRITES: Mapping[str, frozenset[str]] = MappingProxyType(
}
)
_ACTION_DESTRUCTIVE: Mapping[str, frozenset[str]] = MappingProxyType(
{
"manage_calendar": frozenset({"delete_event"}),
"manage_contact": frozenset({"delete"}),
"manage_documents": frozenset({"delete", "tidy"}),
"manage_endpoints": frozenset({"delete"}),
"manage_bg_jobs": frozenset({"kill", "stop", "cancel", "terminate"}),
"manage_memory": frozenset({"delete"}),
"manage_mcp": frozenset({"delete"}),
"manage_notes": frozenset({"delete"}),
"manage_research": frozenset({"delete"}),
"manage_session": frozenset({"delete", "truncate"}),
"manage_settings": frozenset({"delete", "reset"}),
"manage_skills": frozenset({"delete"}),
"manage_tasks": frozenset({"delete"}),
"manage_tokens": frozenset({"delete"}),
"manage_webhooks": frozenset({"delete"}),
}
)
_ACTION_DEFAULTS: Mapping[str, str] = MappingProxyType(
{
"manage_calendar": "list_events",
@@ -415,18 +435,30 @@ def _action_from_content(tool_name: str, content: Any) -> str | None:
def capabilities_for_action(tool_name: Any, content: Any) -> ToolCapabilities:
"""Classify a sealed multiplexed action; ambiguous actions fail high."""
base = capabilities_for_tool(tool_name)
if not isinstance(tool_name, str) or tool_name not in _PRIVATE_ACTION_READS:
if not isinstance(tool_name, str):
return base
action = _action_from_content(tool_name, content)
destructive = action in _ACTION_DESTRUCTIVE.get(tool_name, ())
if tool_name not in _PRIVATE_ACTION_READS:
if not destructive:
return base
return ToolCapabilities(
frozenset(set(base.effects) | {ToolEffect.DESTRUCTIVE}),
base.result_integrity,
known=base.known,
)
if action in _PRIVATE_ACTION_READS[tool_name]:
return _capabilities(
ToolEffect.READ_PRIVATE,
result_integrity=ResultIntegrity.EXTERNAL_UNTRUSTED,
)
if action in _PRIVATE_ACTION_WRITES[tool_name]:
effects = set(base.effects)
if destructive:
effects.add(ToolEffect.DESTRUCTIVE)
return ToolCapabilities(
base.effects,
frozenset(effects),
ResultIntegrity.EXTERNAL_UNTRUSTED,
known=base.known,
)
+20
View File
@@ -639,6 +639,26 @@ async def execute_tool_block(
"policy": "exact_tool_approval",
},
)
if (
exact_approval.pending.tool_name
in {"edit_document", "suggest_document", "update_document"}
and (
not exact_approval.pending.document_id
or exact_approval.pending.document_version is None
)
):
return (
f"{getattr(block, 'tool_type', None)}: BLOCKED",
{
"error": (
"The approved document action has no sealed target and "
"cannot be executed."
),
"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 (