jpmschweitzer/odysseus
1
0
Fork 0
mirror of https://github.com/pewdiepie-archdaemon/odysseus.git synced 2026-09-25 17:42:20 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(agent): authorize exact actions after untrusted context

Browse Source
This commit is contained in:
RaresKeY
2026-08-15 05:37:47 +00:00
parent 2811c7e815
commit 1b09c568d8
16 changed files with 1510 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent_loop.py
+309 -18
View File
@@ -41,7 +41,9 @@ from src.tool_capabilities import (
capabilities_for_tool,
messages_contain_external_untrusted_context,
tool_result_is_successful,
tool_result_should_arm_gate,
)
from src.tool_approvals import ExactToolApproval, tool_approval_store
from src.tool_utils import _truncate, get_mcp_manager
from src.agent_tools import (
parse_tool_blocks,
@@ -3058,7 +3060,11 @@ def _append_tool_results(
result_message["metadata"] = {
"trusted": False,
"source": f"tool result: {tool_name}",
"tool_gate_untrusted": tool_result_is_successful(result),
"tool_gate_untrusted": tool_result_should_arm_gate(
tool_name,
result,
tool_content,
),
}
messages.append(result_message)
else:
@@ -3074,11 +3080,11 @@ def _append_tool_results(
# web/RAG context. THREAT_MODEL.md lists tool output as a surface that
# must go through untrusted_context_message.
arm_tool_gate = any(
tool_result_is_successful(record.get("result"))
and capabilities_for_action(
tool_result_should_arm_gate(
record.get("tool_name"),
record.get("result"),
record.get("content"),
).result_integrity is not ResultIntegrity.SYSTEM
)
for record in tool_result_records
)
messages.append(
@@ -3418,6 +3424,7 @@ async def stream_agent_loop(
uploaded_files: Optional[List[Dict]] = None,
workload: str = "foreground",
external_untrusted_context_seen: bool = False,
exact_approval: Optional[ExactToolApproval] = None,
_is_teacher_run: bool = False,
history_session=None,
defer_context_shaping: bool = False,
@@ -3436,6 +3443,10 @@ async def stream_agent_loop(
run_security = ToolRunSecurityContext(
external_untrusted_context_seen=(
bool(external_untrusted_context_seen)
or bool(
exact_approval
and exact_approval.pending.external_untrusted_context_seen
)
or messages_contain_external_untrusted_context(messages)
)
)
@@ -4435,15 +4446,11 @@ async def stream_agent_loop(
_exhausted_rounds = False
def _filter_route_tool_schemas(schemas):
if not run_security.external_untrusted_context_seen or not schemas:
return schemas
return [
schema
for schema in schemas
if run_security.decision_for(
(schema.get("function") or {}).get("name") or schema.get("name")
).allowed
]
# Keep candidate actions visible after taint so the model can propose
# the exact call that the server will seal for user approval. Schema
# visibility is not authority: both the loop and dispatcher still gate
# execution, and only a one-use server record can cross that boundary.
return schemas
def _tool_schemas_for_route(route_state):
route_mcp_schemas = route_state["mcp_schemas"]
@@ -4484,6 +4491,254 @@ async def stream_agent_loop(
schemas = route_mcp_schemas if wants_mcp and route_mcp_schemas else []
return _filter_route_tool_schemas(schemas)
_approved_result_injected = False
if exact_approval is not None:
approved = exact_approval.pending
approved_block = ToolBlock(approved.tool_name, approved.content)
approved_display = approved.content.strip()
approval_matches = exact_approval.matches(
owner=owner,
session_id=session_id,
tool_name=approved.tool_name,
content=approved.content,
workspace=workspace,
)
if approval_matches:
yield (
"data: "
+ json.dumps(
{
"type": "tool_start",
"tool": approved.tool_name,
"command": approved_display[:240],
"full_command": approved_display,
"round": 0,
"approved": True,
}
)
+ "\n\n"
)
approved_progress_q: asyncio.Queue = asyncio.Queue()
async def _push_approved_progress(payload):
await approved_progress_q.put(payload)
async def _run_approved_tool():
try:
return await execute_tool_block(
approved_block,
session_id=session_id,
disabled_tools=disabled_tools,
tool_policy=tool_policy,
owner=owner,
progress_cb=_push_approved_progress,
workspace=workspace,
security_context=run_security,
exact_approval=exact_approval,
)
finally:
await approved_progress_q.put(None)
approved_tool_task = asyncio.create_task(_run_approved_tool())
try:
while True:
progress_event = await approved_progress_q.get()
if progress_event is None:
break
yield (
"data: "
+ json.dumps(
{
"type": "tool_progress",
"tool": approved.tool_name,
"round": 0,
"approved": True,
**progress_event,
}
)
+ "\n\n"
)
desc, approved_result = await approved_tool_task
finally:
if not approved_tool_task.done():
approved_tool_task.cancel()
try:
await approved_tool_task
except (asyncio.CancelledError, Exception):
pass
total_tool_calls += 1
if tool_result_is_successful(approved_result):
for doc_event in _document_stream_events(approved_block):
yield f"data: {json.dumps(doc_event)}\n\n"
if approved_result.get("action") == "suggest":
yield (
"data: "
+ json.dumps(
{
"type": "doc_suggestions",
"doc_id": approved_result.get("doc_id"),
"suggestions": approved_result.get("suggestions", []),
}
)
+ "\n\n"
)
elif approved_result.get("doc_id") and approved_result.get("content") is not None:
yield (
"data: "
+ json.dumps(
{
"type": "doc_update",
"doc_id": approved_result["doc_id"],
"title": approved_result.get("title", ""),
"language": approved_result.get("language", ""),
"content": approved_result.get("content", ""),
"version": approved_result.get("version", 1),
}
)
+ "\n\n"
)
if approved_result.get("ui_event"):
yield (
"data: "
+ json.dumps({"type": "ui_control", "data": approved_result})
+ "\n\n"
)
approved_output = str(
approved_result.get("output")
or approved_result.get("stdout")
or approved_result.get("response")
or approved_result.get("results")
or approved_result.get("content")
or approved_result.get("error")
or "(no output)"
)
approved_event = {
"type": "tool_output",
"tool": approved.tool_name,
"command": approved_display[:240] if approval_matches else "",
"output": _truncate(approved_output),
"exit_code": approved_result.get("exit_code"),
"approved": True,
}
for key in (
"image_url",
"image_id",
"image_prompt",
"image_model",
"image_size",
"image_quality",
"doc_id",
"title",
"language",
"content",
"version",
"action",
"ui_event",
"diff",
):
if key in approved_result:
approved_event[key] = approved_result[key]
if approved_result.get("images"):
approved_image = approved_result["images"][0]
approved_event["screenshot"] = (
f"data:{approved_image['mimeType']};base64,{approved_image['data']}"
)
yield "data: " + json.dumps(approved_event) + "\n\n"
if approved_result.get("image_url"):
yield (
"data: "
+ json.dumps(
{
"type": "generated_image",
"url": approved_result["image_url"],
**{
key: approved_result[key]
for key in (
"image_url",
"image_id",
"image_prompt",
"image_model",
"image_size",
"image_quality",
)
if key in approved_result
},
}
)
+ "\n\n"
)
approved_research_id = approved_result.get("research_session_id")
if approved_research_id:
approved_anchor = (
f"\n\n[Open in Deep Research](#research-{approved_research_id})\n"
)
full_response += approved_anchor
yield "data: " + json.dumps({"delta": approved_anchor}) + "\n\n"
approved_note_id = approved_result.get("note_id")
if approved_note_id and approved.tool_name == "manage_notes":
approved_note_title = str(
approved_result.get("note_title") or ""
).strip()
approved_note_label = (
f"View note: {approved_note_title}"
if approved_note_title
else "View note"
)
approved_anchor = (
f"\n\n[{approved_note_label}](#note-{approved_note_id})\n"
)
full_response += approved_anchor
yield "data: " + json.dumps({"delta": approved_anchor}) + "\n\n"
approved_tool_event = {
"round": 0,
"tool": approved.tool_name,
"desc": desc,
"command": approved_display[:240] if approval_matches else "",
"output": _truncate(approved_output),
"exit_code": approved_result.get("exit_code"),
"approved": True,
"approval_digest": approved.digest[:16],
}
for key in (
"image_url",
"image_prompt",
"image_model",
"image_size",
"image_quality",
"diff",
):
if approved_result.get(key):
approved_tool_event[key] = approved_result[key]
if approved_result.get("doc_id"):
approved_tool_event["doc_id"] = approved_result["doc_id"]
approved_tool_event["doc_title"] = approved_result.get("title", "")
tool_events.append(approved_tool_event)
if approved.tool_name in _VERIFIER_EFFECTFUL_TOOLS:
_effectful_used = True
formatted_approved_result = format_tool_result(desc, approved_result)
_append_tool_results(
messages,
"",
[],
[formatted_approved_result],
[formatted_approved_result],
False,
0,
tool_result_records=[
{
"tool_name": approved.tool_name,
"content": approved.content,
"result": approved_result,
"text": formatted_approved_result,
}
],
)
_approved_result_injected = True
for round_num in range(1, max_rounds + 1):
round_response = ""
round_reasoning = "" # reasoning_content deltas (DeepSeek-thinking, vLLM --reasoning-parser)
@@ -4504,7 +4759,7 @@ async def stream_agent_loop(
_route_state.get("compaction_state", {}) if round_num == 1 else {}
),
}
if round_num == 1:
if round_num == 1 and not _approved_result_injected:
_active_route_state["request_messages"] = _initial_route_request_messages
all_tool_schemas = _tool_schemas_for_route(_active_route_state)
agent_stream_timeout = int(get_setting("agent_stream_timeout_seconds", 300) or 300)
@@ -5368,12 +5623,44 @@ async def stream_agent_loop(
and block.tool_type in {"manage_notes", "manage_calendar", "manage_tasks"}
)
if not security_decision.allowed:
desc, result = blocked_tool_result(
block.tool_type,
security_decision.reason or "Tool blocked by external-context policy.",
approval_document = (
active_document
if block.tool_type
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(
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(
"Tool blocked before start by external-context policy: %s",
"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:
@@ -5857,6 +6144,10 @@ async def stream_agent_loop(
and not result.get("error")
):
_ody_doc_tool_completed = True
if _pending_ask_user_event:
# An approval card is a turn boundary. Never execute a later
# model-supplied call from the same batch after this request.
break
# If budget was hit, stop the loop
if budget_hit:
src/agent_tools/document_tools.py
+33
View File
@@ -80,6 +80,27 @@ def _most_recent_owned_document(db, Document, owner: Optional[str], active_only:
return q.order_by(Document.updated_at.desc()).first()
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:
return None
try:
unchanged = int(getattr(doc, "version_count", -1)) == int(expected)
except (TypeError, ValueError):
unchanged = False
if unchanged:
return None
return {
"error": (
"The target document changed after this action was proposed. "
"Review the latest version and request the edit again."
),
"exit_code": 1,
"document_changed": True,
}
# ---------------------------------------------------------------------------
# Document tools — create/update/edit/suggest living documents
# ---------------------------------------------------------------------------
@@ -463,6 +484,10 @@ class UpdateDocumentTool:
if not doc:
return {"error": "No documents exist to update"}
version_error = _approved_document_version_error(doc, ctx)
if version_error:
return version_error
is_email_doc = doc.language == "email" or _looks_like_email_document(doc.current_content or "", doc.title or "")
new_content = _coerce_email_document_content(doc.current_content or "", content) if is_email_doc else content.strip()
if is_email_doc:
@@ -541,6 +566,10 @@ class EditDocumentTool:
if not doc:
return {"error": "No documents exist to edit"}
version_error = _approved_document_version_error(doc, ctx)
if version_error:
return version_error
is_email_doc = doc.language == "email" or _looks_like_email_document(doc.current_content or "", doc.title or "")
blank_find_edits = [e for e in edits if not (e.get("find") or "").strip()]
if blank_find_edits:
@@ -677,6 +706,10 @@ class SuggestDocumentTool:
if not doc:
return {"error": f"Document {target_id} not found"}
version_error = _approved_document_version_error(doc, ctx)
if version_error:
return version_error
# Validate that FIND text exists in document
valid = []
for s in suggestions:
src/integrations.py
+8 -1
View File
@@ -719,7 +719,14 @@ async def execute_api_call(
output = f"HTTP {status}\n{formatted}"
if status >= 400:
return {"error": output, "exit_code": 1}
return {
"error": output,
"exit_code": 1,
# The error string includes the remote response body. Preserve
# it for diagnostics, but make its provenance explicit so the
# agent gate does not treat HTTP failure as content-free.
"untrusted_content": True,
}
return {"output": output, "exit_code": 0}
src/tool_approvals.py
+340
View File
@@ -0,0 +1,340 @@
"""Opaque, exact, one-use approvals for tainted model-requested actions.
The model may propose an action after untrusted context, but only the server
stores and later executes the exact approved tool input. Browser-visible
fields are display copies, never authority.
"""
from __future__ import annotations
import hashlib
import json
import os
import secrets
import threading
import time
from dataclasses import dataclass, field
from typing import Any
from src.tool_capabilities import ToolCapabilities, capabilities_for_action
DEFAULT_APPROVAL_TTL_SECONDS = 10 * 60
DEFAULT_MAX_PENDING_APPROVALS = 2048
def _normalized_owner(owner: Any) -> str:
return str(owner or "").strip().casefold()
def _normalized_workspace(workspace: Any) -> str:
if not isinstance(workspace, str) or not workspace.strip():
return ""
return os.path.realpath(os.path.expanduser(workspace))
def _canonical_digest(payload: dict[str, Any]) -> str:
encoded = json.dumps(
payload,
sort_keys=True,
separators=(",", ":"),
ensure_ascii=False,
).encode("utf-8")
return hashlib.sha256(encoded).hexdigest()
def _binding_payload(
*,
owner: Any,
session_id: Any,
origin_run_id: Any,
tool_name: Any,
content: Any,
workspace: Any,
document_id: Any,
document_version: Any,
external_untrusted_context_seen: bool,
effects: tuple[str, ...],
result_integrity: str,
) -> dict[str, Any]:
return {
"owner": _normalized_owner(owner),
"session_id": str(session_id or ""),
"origin_run_id": str(origin_run_id or ""),
"tool_name": str(tool_name or ""),
"content": str(content or ""),
"workspace": _normalized_workspace(workspace),
"document_id": str(document_id or ""),
"document_version": (
int(document_version) if document_version is not None else None
),
"external_untrusted_context_seen": bool(external_untrusted_context_seen),
"effects": list(effects),
"result_integrity": str(result_integrity),
}
@dataclass(frozen=True)
class PendingToolApproval:
approval_id: str
owner: str
session_id: str
origin_run_id: str
tool_name: str
content: str
workspace: str
document_id: str
document_version: int | None
external_untrusted_context_seen: bool
effects: tuple[str, ...]
result_integrity: str
digest: str
created_at: float
expires_at: float
def public_payload(self, *, reason: str | None = None) -> dict[str, Any]:
return {
"kind": "tool_approval",
"approval_id": self.approval_id,
"question": "Allow this exact action once?",
"description": reason or (
"Untrusted context influenced this run, so this action needs "
"your explicit approval."
),
"options": [
{
"label": "Allow once",
"value": "approve",
"description": "Execute only the sealed action shown here.",
},
{
"label": "Deny",
"value": "deny",
"description": "Do not execute it.",
},
],
"action": {
"tool": self.tool_name,
# Show the complete sealed input so approval never hides
# trailing lines. This is not read back as authority.
"content": self.content,
"digest": self.digest[:16],
"effects": list(self.effects),
"workspace": self.workspace or None,
"document_id": self.document_id or None,
"document_version": self.document_version,
},
}
@dataclass
class ExactToolApproval:
"""A consumed grant that the dispatcher can claim exactly once."""
pending: PendingToolApproval
_claimed: bool = field(default=False, init=False, repr=False)
_lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
def _matches_unlocked(
self,
*,
owner: Any,
session_id: Any,
tool_name: Any,
content: Any,
workspace: Any,
) -> bool:
if self._claimed:
return False
capabilities = capabilities_for_action(tool_name, content)
effects = tuple(sorted(effect.value for effect in capabilities.effects))
result_integrity = capabilities.result_integrity.value
if (
effects != self.pending.effects
or result_integrity != self.pending.result_integrity
):
return False
expected = _binding_payload(
owner=owner,
session_id=session_id,
origin_run_id=self.pending.origin_run_id,
tool_name=tool_name,
content=content,
workspace=workspace,
document_id=self.pending.document_id,
document_version=self.pending.document_version,
external_untrusted_context_seen=(
self.pending.external_untrusted_context_seen
),
effects=effects,
result_integrity=result_integrity,
)
return _canonical_digest(expected) == self.pending.digest
def matches(
self,
*,
owner: Any,
session_id: Any,
tool_name: Any,
content: Any,
workspace: Any,
) -> bool:
with self._lock:
return self._matches_unlocked(
owner=owner,
session_id=session_id,
tool_name=tool_name,
content=content,
workspace=workspace,
)
def claim(
self,
*,
owner: Any,
session_id: Any,
tool_name: Any,
content: Any,
workspace: Any,
) -> bool:
with self._lock:
if not self._matches_unlocked(
owner=owner,
session_id=session_id,
tool_name=tool_name,
content=content,
workspace=workspace,
):
return False
self._claimed = True
return True
class ToolApprovalStore:
"""Thread-safe pending approval registry with destructive consumption."""
def __init__(
self,
*,
ttl_seconds: int = DEFAULT_APPROVAL_TTL_SECONDS,
max_pending: int = DEFAULT_MAX_PENDING_APPROVALS,
):
self._ttl_seconds = max(1, int(ttl_seconds))
self._max_pending = max(1, int(max_pending))
self._pending: dict[str, PendingToolApproval] = {}
self._lock = threading.Lock()
def _purge_expired_locked(self, now: float) -> None:
expired = [
approval_id
for approval_id, pending in self._pending.items()
if pending.expires_at <= now
]
for approval_id in expired:
self._pending.pop(approval_id, None)
def create(
self,
*,
owner: Any,
session_id: Any,
origin_run_id: Any,
tool_name: Any,
content: Any,
workspace: Any,
document_id: Any = None,
document_version: Any = None,
external_untrusted_context_seen: bool,
capabilities: ToolCapabilities,
) -> PendingToolApproval:
now = time.time()
effects = tuple(sorted(effect.value for effect in capabilities.effects))
result_integrity = capabilities.result_integrity.value
payload = _binding_payload(
owner=owner,
session_id=session_id,
origin_run_id=origin_run_id,
tool_name=tool_name,
content=content,
workspace=workspace,
document_id=document_id,
document_version=document_version,
external_untrusted_context_seen=external_untrusted_context_seen,
effects=effects,
result_integrity=result_integrity,
)
pending = PendingToolApproval(
approval_id=secrets.token_urlsafe(32),
owner=payload["owner"],
session_id=payload["session_id"],
origin_run_id=payload["origin_run_id"],
tool_name=payload["tool_name"],
content=payload["content"],
workspace=payload["workspace"],
document_id=payload["document_id"],
document_version=payload["document_version"],
external_untrusted_context_seen=payload[
"external_untrusted_context_seen"
],
effects=effects,
result_integrity=result_integrity,
digest=_canonical_digest(payload),
created_at=now,
expires_at=now + self._ttl_seconds,
)
with self._lock:
self._purge_expired_locked(now)
# The UI exposes one pending card per chat. Supersede any older
# action for the same owner/session so stale history cannot retain
# parallel grants and the in-memory registry stays bounded.
superseded = [
approval_id
for approval_id, existing in self._pending.items()
if (
existing.owner == pending.owner
and existing.session_id == pending.session_id
)
]
for approval_id in superseded:
self._pending.pop(approval_id, None)
while len(self._pending) >= self._max_pending:
oldest_id = min(
self._pending,
key=lambda approval_id: self._pending[approval_id].created_at,
)
self._pending.pop(oldest_id, None)
self._pending[pending.approval_id] = pending
return pending
def consume(
self,
approval_id: Any,
*,
decision: Any,
owner: Any,
session_id: Any,
) -> ExactToolApproval | None:
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
if str(decision or "").strip().lower() != "approve":
return None
return ExactToolApproval(pending)
def peek(self, approval_id: Any) -> PendingToolApproval | None:
now = time.time()
with self._lock:
self._purge_expired_locked(now)
return self._pending.get(str(approval_id or ""))
tool_approval_store = ToolApprovalStore()
src/tool_capabilities.py
+50 -10
View File
@@ -8,6 +8,7 @@ run-local integrity gates before dispatch.
from __future__ import annotations
import json
import uuid
from dataclasses import dataclass, field
from enum import Enum
from types import MappingProxyType
@@ -79,9 +80,17 @@ _register(
"list_models",
"list_serve_presets",
"list_served_models",
"search_hf_models",
},
ToolEffect.READ_PUBLIC,
ToolEffect.READ_PRIVATE,
# These readers return provider-controlled model identifiers or durable
# user/admin-authored Cookbook and process state. Local brokering does not
# make the returned text server-authored.
result_integrity=ResultIntegrity.EXTERNAL_UNTRUSTED,
)
_register(
{"search_hf_models"},
ToolEffect.BROKERED_NETWORK_READ,
result_integrity=ResultIntegrity.EXTERNAL_UNTRUSTED,
)
_register(
{"get_workspace", "glob", "grep", "ls", "read_file"},
@@ -205,15 +214,8 @@ _register(
_register(
{
"adopt_served_model",
"api_call",
"app_api",
"cancel_download",
"download_model",
"manage_endpoints",
"manage_mcp",
"manage_settings",
"manage_tokens",
"manage_webhooks",
"serve_model",
"serve_preset",
"stop_served_model",
@@ -221,6 +223,22 @@ _register(
},
ToolEffect.ADMIN_CHANGE,
)
_register(
{
"api_call",
"app_api",
"manage_endpoints",
"manage_mcp",
"manage_settings",
"manage_tokens",
"manage_webhooks",
},
ToolEffect.ADMIN_CHANGE,
# api_call/app_api return remote or stored application data, and the
# admin managers can echo user-controlled configuration. Conservatively
# retain the action effect while treating every successful result as data.
result_integrity=ResultIntegrity.EXTERNAL_UNTRUSTED,
)
TOOL_CAPABILITIES: Mapping[str, ToolCapabilities] = MappingProxyType(dict(_REGISTRY))
@@ -432,6 +450,27 @@ def tool_result_is_successful(result: Any) -> bool:
)
def tool_result_should_arm_gate(
tool_name: Any,
result: Any,
content: Any = None,
) -> bool:
"""Return whether a result introduced non-system content to the model.
A content-free transport or validation failure does not change authority.
Producers set ``untrusted_content`` when a failed response still carries a
remote/private body, so HTTP status alone cannot launder that body.
"""
if not isinstance(result, dict):
return False
if result.get("blocked") or result.get("approval_required"):
return False
capabilities = capabilities_for_action(tool_name, content)
if capabilities.result_integrity is ResultIntegrity.SYSTEM:
return False
return tool_result_is_successful(result) or result.get("untrusted_content") is True
POST_EXTERNAL_BLOCKED_EFFECTS = frozenset(
{
ToolEffect.READ_PRIVATE,
@@ -500,6 +539,7 @@ class ToolRunSecurityContext:
external_untrusted_context_seen: bool = False
external_sources: list[str] = field(default_factory=list)
run_id: str = field(default_factory=lambda: uuid.uuid4().hex)
def observe_messages(self, messages: Iterable[dict]) -> None:
"""Promote any server-labelled untrusted prompt context into the gate."""
@@ -531,7 +571,7 @@ class ToolRunSecurityContext:
result: Any,
content: Any = None,
) -> None:
if not tool_result_is_successful(result):
if not tool_result_should_arm_gate(tool_name, result, content):
return
capabilities = capabilities_for_action(tool_name, content)
if capabilities.result_integrity is not ResultIntegrity.SYSTEM:
src/tool_execution.py
+79 -3
View File
@@ -28,6 +28,7 @@ from src.tool_security import (
owner_is_admin_or_single_user,
)
from src.tool_capabilities import ToolRunSecurityContext, blocked_tool_result
from src.tool_approvals import ExactToolApproval
from src.tool_policy import ToolPolicy
from src.constants import MAX_OUTPUT_CHARS, MAX_READ_CHARS, MAX_DIFF_LINES, DATA_DIR
from src.tool_utils import _truncate, get_mcp_manager
@@ -567,10 +568,17 @@ async def _document_tool_dispatch(
content: str,
session_id: Optional[str] = None,
owner: Optional[str] = None,
document_id: Optional[str] = None,
document_version: Optional[int] = None,
) -> Optional[Dict]:
"""Route a document tool through TOOL_HANDLERS with the right ctx shape."""
from src.agent_tools import TOOL_HANDLERS
ctx = {"session_id": session_id, "owner": owner}
ctx = {
"session_id": session_id,
"owner": owner,
"doc_id": document_id,
"expected_document_version": document_version,
}
if tool in TOOL_HANDLERS:
return await TOOL_HANDLERS[tool](content, ctx)
return None
@@ -593,6 +601,7 @@ async def execute_tool_block(
| _NoToolSecurityContext
| _MissingToolSecurityContext
) = _MISSING_TOOL_SECURITY_CONTEXT,
exact_approval: Optional[ExactToolApproval] = None,
) -> Tuple[str, Dict]:
"""Execute a single tool block. Returns (description, result_dict).
@@ -614,7 +623,55 @@ async def execute_tool_block(
"NO_TOOL_SECURITY_CONTEXT"
)
if isinstance(security_context, ToolRunSecurityContext):
approval_claimed = False
if exact_approval is not None:
if (
not isinstance(security_context, ToolRunSecurityContext)
or not security_context.external_untrusted_context_seen
or not exact_approval.pending.external_untrusted_context_seen
):
return (
f"{getattr(block, 'tool_type', None)}: BLOCKED",
{
"error": "Exact-action approval requires an armed run security context.",
"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 (
f"{getattr(block, 'tool_type', None)}: BLOCKED",
{
"error": (
"The approved workspace is no longer a valid safe "
"directory. Review the action again."
),
"exit_code": 1,
"blocked": True,
"policy": "exact_tool_approval",
},
)
approval_claimed = exact_approval.claim(
owner=owner,
session_id=session_id,
tool_name=getattr(block, "tool_type", None),
content=getattr(block, "content", None),
workspace=workspace,
)
if not approval_claimed:
return (
f"{getattr(block, 'tool_type', None)}: BLOCKED",
{
"error": "The exact-action approval did not match this tool request.",
"exit_code": 1,
"blocked": True,
"policy": "exact_tool_approval",
},
)
if isinstance(security_context, ToolRunSecurityContext) and not approval_claimed:
decision = security_context.decision_for(
getattr(block, "tool_type", None),
getattr(block, "content", None),
@@ -638,6 +695,16 @@ async def execute_tool_block(
owner=owner,
progress_cb=progress_cb,
tool_policy=tool_policy,
approved_document_id=(
exact_approval.pending.document_id
if approval_claimed
else None
),
approved_document_version=(
exact_approval.pending.document_version
if approval_claimed
else None
),
)
if isinstance(security_context, ToolRunSecurityContext):
security_context.observe_tool_result(
@@ -657,6 +724,8 @@ async def _execute_tool_block_impl(
owner: Optional[str] = None,
progress_cb: Optional[Callable[[Dict], Awaitable[None]]] = None,
tool_policy: Optional[Any] = None,
approved_document_id: Optional[str] = None,
approved_document_version: Optional[int] = None,
) -> Tuple[str, Dict]:
"""Execute a single tool block. Returns (description, result_dict).
@@ -818,7 +887,14 @@ async def _execute_tool_block_impl(
elif tool in ("create_document", "update_document", "edit_document",
"suggest_document", "manage_documents"):
desc = f"{tool}: {content.split(chr(10))[0][:80]}"
result = await _document_tool_dispatch(tool, content, session_id, owner) \
result = await _document_tool_dispatch(
tool,
content,
session_id,
owner,
document_id=approved_document_id,
document_version=approved_document_version,
) \
or {"error": f"{tool}: execution failed", "exit_code": 1}
if tool in ("edit_document", "suggest_document") and "title" in (result or {}):
desc = f"{tool}: {result.get('title', '')}"
src/tools/system.py
+1
View File
@@ -725,6 +725,7 @@ async def do_app_api(content: str, owner: Optional[str] = None) -> Dict:
"status_code": resp.status_code,
"body": preview,
"exit_code": 1,
"untrusted_content": True,
}
return {
"output": f"{method} {path} -> {resp.status_code}\n{preview}",