211 lines
8.4 KiB
Python
211 lines
8.4 KiB
Python
"""Application-layer bridge between the harness ``WriteGuardMiddleware`` and
|
|
the MAF ``ActionCoordinator``.
|
|
|
|
The harness middleware (``deerflow.agents.middlewares.write_guard_middleware``)
|
|
is constructed once at agent-build time and cannot hold per-run state. It
|
|
resolves two factories — ``make_policy`` / ``make_interceptor`` — by dotted
|
|
path (``resolve_variable``). These factories return lightweight singleton
|
|
objects that, at interception time, read the MAF service + per-run data from
|
|
the run context.
|
|
|
|
Wiring
|
|
------
|
|
``register_write_guard_service(service)`` is called from ``deps.py`` once the
|
|
``MafRoundtableService`` is built, stashing a module-level reference. The
|
|
factories use it to build an ``ActionCoordinator`` on demand.
|
|
|
|
Per-run data (``session_id`` / ``run_id`` / ``agent_id`` / ``turn_index`` /
|
|
``write_policy`` / ``capability_snapshot``) is injected into the DeerFlow run
|
|
context by ``runtime_client.start_agent_run`` (see Task 1.4) and read back
|
|
here from the ``run_context`` dict the middleware forwards.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Any
|
|
|
|
from app.maf_roundtable.action_coordinator import (
|
|
ActionCoordinator,
|
|
tool_needs_approval,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Module-level reference to the MAF service, set by ``register_write_guard_service``.
|
|
# Resolved lazily so importing this module never requires the service to exist.
|
|
_maf_service: Any = None
|
|
|
|
|
|
def register_write_guard_service(service: Any) -> None:
|
|
"""Stash the MAF service so the WriteGuard factories can reach it.
|
|
|
|
Called once from ``app.gateway.deps`` after ``MafRoundtableService`` is built.
|
|
"""
|
|
global _maf_service
|
|
_maf_service = service
|
|
|
|
|
|
def _get_service() -> Any:
|
|
if _maf_service is None: # pragma: no cover - defensive; wiring is enforced at startup
|
|
raise RuntimeError(
|
|
"WriteGuard bridge used before register_write_guard_service() — "
|
|
"MAF service is not wired."
|
|
)
|
|
return _maf_service
|
|
|
|
|
|
def _coordinator_for_context(run_context: dict) -> ActionCoordinator:
|
|
"""Build an ActionCoordinator bound to the run's user identity.
|
|
|
|
The user_id is sourced ONLY from the server-injected run context (set by
|
|
``runtime_client`` from the authenticated ``ExecutionPrincipal``), never
|
|
from tool arguments.
|
|
"""
|
|
service = _get_service()
|
|
user_id = str(run_context.get("maf_user_id") or run_context.get("user_id") or "")
|
|
return ActionCoordinator(service, user_id=user_id)
|
|
|
|
|
|
# ── Policy ─────────────────────────────────────────────────────────────────
|
|
|
|
|
|
class MafWriteGuardPolicy:
|
|
"""Decides whether a tool call must be intercepted for approval.
|
|
|
|
Reads the per-run ``write_policy`` and ``capability_snapshot`` from the
|
|
run context. Uses the same metadata-driven classification as the rest of
|
|
the ActionCoordinator (never name-prefix heuristics). Tools without
|
|
metadata use the safe default (treated as a write requiring approval).
|
|
"""
|
|
|
|
def should_intercept(self, tool_name: str, args: Any, run_context: dict | None = None) -> bool: # noqa: ANN401
|
|
"""Decide interception from the run context (read vs write aware).
|
|
|
|
Reads / skills (``side_effect`` of ``read`` / ``skill``) are NEVER
|
|
intercepted, so they pass straight through — including on the sync
|
|
path. Writes are gated by the active write policy via the same
|
|
metadata-driven ``tool_needs_approval`` classification used everywhere
|
|
else (never name-prefix heuristics). When no run context is available
|
|
we defer to the interceptor (legacy behaviour).
|
|
"""
|
|
if not tool_name:
|
|
return False
|
|
if not run_context:
|
|
# Cannot classify without context — let the interceptor decide.
|
|
return True
|
|
write_policy = str(run_context.get("write_policy") or "confirm_high_risk")
|
|
capability_snapshot = run_context.get("capability_snapshot")
|
|
agent_id = str(run_context.get("agent_id") or "")
|
|
return tool_needs_approval(tool_name, write_policy, capability_snapshot, agent_id)
|
|
|
|
|
|
# ── Interceptor ────────────────────────────────────────────────────────────
|
|
|
|
|
|
class MafWriteGuardInterceptor:
|
|
"""Records an intercepted write tool call as a pending action + approval.
|
|
|
|
Implements the harness ``WriteGuardInterceptor`` Protocol. Uses the
|
|
single-transaction ``store.create_pending_action_bundle`` (Task 2) so an
|
|
action + secret + approval + audit event are written atomically — a
|
|
partial action never reaches the database.
|
|
"""
|
|
|
|
async def intercept(
|
|
self,
|
|
*,
|
|
tool_name: str,
|
|
args: Any, # noqa: ANN401
|
|
tool_call_id: str,
|
|
run_context: dict,
|
|
) -> dict[str, Any]:
|
|
write_policy = str(run_context.get("write_policy") or "confirm_high_risk")
|
|
capability_snapshot = run_context.get("capability_snapshot")
|
|
agent_id = str(run_context.get("agent_id") or "")
|
|
run_id = str(run_context.get("maf_run_id") or "")
|
|
session_id = str(run_context.get("maf_session_id") or "")
|
|
turn_index = int(run_context.get("turn_index") or 0)
|
|
|
|
# Apply the real policy now that we have the run context.
|
|
if not tool_needs_approval(tool_name, write_policy, capability_snapshot, agent_id):
|
|
return {"status": "allowed", "tool_name": tool_name}
|
|
|
|
coordinator = _coordinator_for_context(run_context)
|
|
raw_args = args if isinstance(args, dict) else {"value": str(args)}
|
|
|
|
created = await coordinator.intercept_write_tool_calls(
|
|
run_id=run_id,
|
|
session_id=session_id,
|
|
agent_id=agent_id,
|
|
turn_index=turn_index,
|
|
tool_calls=[{"name": tool_name, "id": tool_call_id, "args": raw_args}],
|
|
capability_snapshot=capability_snapshot,
|
|
write_policy=write_policy,
|
|
)
|
|
|
|
if not created:
|
|
return {"status": "allowed", "tool_name": tool_name}
|
|
|
|
action = created[0]
|
|
return {
|
|
"status": "approval_required",
|
|
"action_id": action.get("id", ""),
|
|
"approval_id": action.get("approval_id", ""),
|
|
"tool_name": tool_name,
|
|
}
|
|
|
|
async def validate_approved_action(
|
|
self,
|
|
*,
|
|
tool_name: str,
|
|
run_context: dict,
|
|
) -> bool:
|
|
"""Validate a post-approval deterministic execution against the store.
|
|
|
|
Field presence is NOT trusted: this queries the bound action and confirms
|
|
it is ``running`` (the executor already occupied it), belongs to the
|
|
current ``run_id`` / ``session_id`` / ``agent_id``, and its ``tool_name``
|
|
matches. Returns True only when every check holds; otherwise the
|
|
middleware fails closed and the tool does not run.
|
|
"""
|
|
action_id = str(run_context.get("maf_approved_action_id") or "")
|
|
if not action_id:
|
|
return False
|
|
approved_tool = str(run_context.get("maf_approved_tool") or "")
|
|
if approved_tool and tool_name != approved_tool:
|
|
return False
|
|
|
|
run_id = str(run_context.get("maf_run_id") or "")
|
|
session_id = str(run_context.get("maf_session_id") or "")
|
|
agent_id = str(run_context.get("agent_id") or "")
|
|
|
|
service = _get_service()
|
|
action = await service.store.actions.get_action(action_id)
|
|
if action is None:
|
|
return False
|
|
if action.get("status") != "running":
|
|
return False
|
|
if run_id and str(action.get("run_id") or "") != run_id:
|
|
return False
|
|
if session_id and str(action.get("session_id") or "") != session_id:
|
|
return False
|
|
if agent_id and str(action.get("agent_id") or "") != agent_id:
|
|
return False
|
|
if str(action.get("tool_name") or "") != tool_name:
|
|
return False
|
|
return True
|
|
|
|
|
|
# ── Factories resolved by the harness via resolve_variable ──────────────────
|
|
|
|
|
|
def make_policy(**_: Any) -> MafWriteGuardPolicy:
|
|
"""Factory for the harness ``resolve_variable`` resolution."""
|
|
return MafWriteGuardPolicy()
|
|
|
|
|
|
def make_interceptor(**_: Any) -> MafWriteGuardInterceptor:
|
|
"""Factory for the harness ``resolve_variable`` resolution."""
|
|
return MafWriteGuardInterceptor()
|