From 26abdc4fe415981edcea4f84e40180b4821d2808 Mon Sep 17 00:00:00 2001 From: Damian Loch Date: Mon, 27 Apr 2026 17:39:58 -0400 Subject: [PATCH 01/13] feat(chat): interactive command gate with Yes/No/Yes-Always Add a unified command gate (utils/auth/command_gate.py) that runs the existing signature match, org allow/deny policy, and LLM safety judge layers behind a single entry point, and prompts the user in foreground chats when any layer blocks. Background chats and RCAs keep the pre-existing hard-block behavior. - Yes: approve for this invocation (silent to the agent). - No: abort with code USER_DECLINED, distinct from static policy codes so the agent sees an explicit user rejection. - Yes-Always: only offered on org-policy blocks; mutates the policy (disables the hit deny rule and/or adds an editable allow pattern) so future runs including background RCAs inherit the decision. Extend CommandVerdict with deny_rule_id and allowlist_exhausted so the gate can plan the right mutation, add plan_yes_always / apply_yes_always and a pattern-derivation helper in command_policy. Extend the existing WebSocket confirmation protocol with block_layer, yes_always_effect, an execute_always option, and edited_patterns handling. Thread the new fields through useMessageHandler and render them in ToolExecutionWidget with an editable regex input for proposed allow rules. Idempotency: two contextvars short-circuit duplicate prompts and duplicate guardrail LLM calls when one shell tool routes into another (terminal_exec -> cloud_exec) and when terminal_run._check_guardrails would otherwise re-evaluate a command the gate just approved. No new feature flag: the gate is the interactive surface of the existing guardrails, so GUARDRAILS_ENABLED and per-org list toggles continue to govern whether anything blocks (and therefore prompts) at all. Made-with: Cursor --- client/src/app/chat/types.ts | 17 + .../tool-calls/ToolExecutionWidget.tsx | 170 +++++++--- client/src/hooks/useMessageHandler.ts | 4 +- .../backend/agent/tools/cloud_exec_tool.py | 21 +- .../agent/tools/kubectl_onprem_tool.py | 35 +- .../backend/agent/tools/tailscale_ssh_tool.py | 30 +- .../backend/agent/tools/terminal_exec_tool.py | 19 +- server/utils/auth/command_gate.py | 315 ++++++++++++++++++ server/utils/auth/command_policy.py | 128 ++++++- .../cloud/infrastructure_confirmation.py | 103 +++++- server/utils/terminal/terminal_run.py | 14 + 11 files changed, 731 insertions(+), 125 deletions(-) create mode 100644 server/utils/auth/command_gate.py diff --git a/client/src/app/chat/types.ts b/client/src/app/chat/types.ts index 9922066a6..95f6c1e8d 100644 --- a/client/src/app/chat/types.ts +++ b/client/src/app/chat/types.ts @@ -1,4 +1,17 @@ // Message type for deploy page chat +export interface PolicyChange { + action: 'disable_deny_rule' | 'add_allow_rule'; + rule_id?: number | null; + pattern?: string | null; + description?: string | null; + editable: boolean; +} + +export interface YesAlwaysEffect { + summary: string; + changes: PolicyChange[]; +} + export interface ToolCall { id: string; tool_name: string; @@ -9,6 +22,10 @@ export interface ToolCall { timestamp: string; confirmation_id?: string; confirmation_message?: string; + // Set when the confirmation originates from the command gate. Drives the + // Yes-Always button visibility and the policy-change disclosure. + block_layer?: string; + yes_always_effect?: YesAlwaysEffect; command?: string; // Add command field to store final_command isExpanded?: boolean; // Track whether the tool output is expanded } diff --git a/client/src/components/tool-calls/ToolExecutionWidget.tsx b/client/src/components/tool-calls/ToolExecutionWidget.tsx index 2d26b0219..e70af71d5 100644 --- a/client/src/components/tool-calls/ToolExecutionWidget.tsx +++ b/client/src/components/tool-calls/ToolExecutionWidget.tsx @@ -466,56 +466,13 @@ const ToolExecutionWidget = ({ tool, className, sendMessage, sendRaw, onToolUpda {/* Show message when awaiting confirmation */} {tool.status === "awaiting_confirmation" && !tool.output && !tool.error && ( -
- - {(tool as any).confirmation_message || "This action requires confirmation"} - -
- - -
-
+ )} {/* Show shimmer effect while tool is running and no output yet */} @@ -586,3 +543,116 @@ const ToolExecutionWidget = ({ tool, className, sendMessage, sendRaw, onToolUpda } export default ToolExecutionWidget + +interface ConfirmationPanelProps { + tool: ToolCall + userId?: string + sessionId?: string + sendRaw?: (data: string) => boolean + onToolUpdate?: (updatedTool: Partial) => void +} + +const ConfirmationPanel = ({ tool, userId, sessionId, sendRaw, onToolUpdate }: ConfirmationPanelProps) => { + const effect = tool.yes_always_effect + const allowYesAlways = !!(effect && effect.changes.length > 0) + + // Pre-fill editable change patterns. Map is keyed by change index (as string + // to match the backend's JSON-parsed dict key). + const [edited, setEdited] = React.useState>(() => { + const m: Record = {} + if (effect) { + effect.changes.forEach((c, i) => { + if (c.editable && c.pattern) m[String(i)] = c.pattern + }) + } + return m + }) + + const respond = (decision: 'execute' | 'cancel' | 'execute_always') => { + const confirmationId = tool.confirmation_id + if (!confirmationId || !sendRaw || !userId) return + const payload: Record = { + type: 'confirmation_response', + confirmation_id: confirmationId, + decision, + user_id: userId, + session_id: sessionId, + } + if (decision === 'execute_always') payload.edited_patterns = edited + sendRaw(JSON.stringify(payload)) + if (decision === 'cancel') { + onToolUpdate?.({ status: 'completed', output: 'Operation cancelled by user' }) + } else { + onToolUpdate?.({ status: 'running' }) + } + } + + return ( +
+
+ + {tool.confirmation_message || "This action requires confirmation"} + +
+ + +
+
+ + {allowYesAlways && effect && ( +
+
{effect.summary}
+
    + {effect.changes.map((c, i) => ( +
  • + {c.action === 'disable_deny_rule' ? ( + + Disable deny rule: {c.description || 'rule'} + {c.pattern ? <> (pattern: {c.pattern}) : null} + + ) : ( + <> + Add allow rule (you can edit the pattern): + setEdited(prev => ({ ...prev, [String(i)]: e.target.value }))} + spellCheck={false} + aria-label="Allow-rule regex pattern" + /> + + )} +
    • + ))} +
+
+ +
+
+ )} +
+ ) +} diff --git a/client/src/hooks/useMessageHandler.ts b/client/src/hooks/useMessageHandler.ts index 62534a333..00fa5787d 100644 --- a/client/src/hooks/useMessageHandler.ts +++ b/client/src/hooks/useMessageHandler.ts @@ -115,7 +115,9 @@ export const useMessageHandler = ({ ...tc, status: 'awaiting_confirmation' as const, confirmation_id: confirmationId, - confirmation_message: message.data.message + confirmation_message: message.data.message, + block_layer: message.data.block_layer, + yes_always_effect: message.data.yes_always_effect, }; } return tc; diff --git a/server/chat/backend/agent/tools/cloud_exec_tool.py b/server/chat/backend/agent/tools/cloud_exec_tool.py index 61ef35ff4..c6e11b5a8 100644 --- a/server/chat/backend/agent/tools/cloud_exec_tool.py +++ b/server/chat/backend/agent/tools/cloud_exec_tool.py @@ -1409,25 +1409,22 @@ def cloud_exec(provider: str, command: str, user_id: Optional[str] = None, sessi normalized_provider = _normalize_cloud_exec_provider(provider) provider = normalized_provider - # Org command policy check -- must run before any execution path branches + # Unified gate: signature + org policy + LLM judge + HITL (foreground). # Prepend CLI prefix so patterns like ^aws\s+ match (cloud_exec receives # the subcommand without the provider prefix, e.g. "ecs list-clusters"). _CLI_PREFIX = {"aws": "aws", "gcp": "gcloud", "azure": "az", "scaleway": "scw", "ovh": "ovhcloud"} - from utils.auth.command_policy import evaluate_compound_command - from utils.auth.stateless_auth import get_org_id_for_user - org_id = get_org_id_for_user(user_id) if user_id else None prefix = _CLI_PREFIX.get(provider.lower(), "") - policy_cmd = f"{prefix} {command}" if prefix and not command.strip().startswith(prefix) else command - verdict = evaluate_compound_command(org_id, policy_cmd) - if not verdict.allowed: - reason = (verdict.rule_description or "Matched organization policy")[:200] - logger.warning("Policy denied cloud command for user %s (%s)", - user_id, reason) + gated_cmd = f"{prefix} {command}" if prefix and not command.strip().startswith(prefix) else command + from utils.auth.command_gate import gate_command + gate = gate_command(user_id=user_id, tool_name="cloud_exec", command=gated_cmd) + if not gate.allowed: + logger.warning("cloud_exec blocked for user %s (%s): %s", + user_id, gate.code, gate.block_reason[:200]) return json.dumps({ "success": False, - "error": f"Command blocked by organization policy: {reason}", - "code": "POLICY_DENIED", + "error": gate.block_reason, + "code": gate.code, "final_command": command, "provider": provider.lower(), }) diff --git a/server/chat/backend/agent/tools/kubectl_onprem_tool.py b/server/chat/backend/agent/tools/kubectl_onprem_tool.py index b4b592358..b63e7eb5e 100644 --- a/server/chat/backend/agent/tools/kubectl_onprem_tool.py +++ b/server/chat/backend/agent/tools/kubectl_onprem_tool.py @@ -42,35 +42,18 @@ def on_prem_kubectl( if command.lower().startswith('kubectl '): command = command[8:].strip() - # Org command policy check against full kubectl command - from utils.auth.command_policy import evaluate_compound_command - from utils.auth.stateless_auth import get_org_id_for_user + # Unified gate: signature + org policy + LLM judge + HITL (foreground). full_command = f"kubectl {command}" - org_id = get_org_id_for_user(user_id) if user_id else None - verdict = evaluate_compound_command(org_id, full_command) - if not verdict.allowed: - reason = (verdict.rule_description or "Matched organization policy")[:200] - logger.warning("Policy denied kubectl_onprem command for user %s (%s)", - user_id, reason) + from utils.auth.command_gate import gate_command + gate = gate_command(user_id=user_id, tool_name="kubectl_onprem", command=full_command) + if not gate.allowed: + logger.warning("kubectl_onprem blocked for user %s (%s): %s", + user_id, gate.code, gate.block_reason[:200]) return json.dumps({ 'success': False, - 'error': f"Command blocked by organization policy: {reason}", - 'code': 'POLICY_DENIED', - 'chat_output': f"$ {full_command}\nBlocked by organization policy: {reason}", - 'command': full_command, - 'return_code': 1, - 'provider': 'onprem_kubectl', - }) - - from utils.security.command_safety import evaluate_command - decision = evaluate_command(full_command, tool="kubectl_onprem", user_id=user_id, session_id=session_id) - if decision.blocked: - code = 'SIGNATURE_MATCHED' if decision.layer == 'signature_match' else 'SAFETY_BLOCKED' - return json.dumps({ - 'success': False, - 'error': f"Command blocked by safety guardrail: {decision.reason}", - 'code': code, - 'chat_output': f"$ {full_command}\nBlocked by safety guardrail: {decision.reason}", + 'error': gate.block_reason, + 'code': gate.code, + 'chat_output': f"$ {full_command}\n{gate.block_reason}", 'command': full_command, 'return_code': 1, 'provider': 'onprem_kubectl', diff --git a/server/chat/backend/agent/tools/tailscale_ssh_tool.py b/server/chat/backend/agent/tools/tailscale_ssh_tool.py index cd57b110b..1229ddc43 100644 --- a/server/chat/backend/agent/tools/tailscale_ssh_tool.py +++ b/server/chat/backend/agent/tools/tailscale_ssh_tool.py @@ -223,30 +223,16 @@ def tailscale_ssh( "error": "Command cannot be empty" }) - # Org command policy check (shared allow/deny firewall across all tools) - from utils.auth.command_policy import evaluate_compound_command - from utils.auth.stateless_auth import get_org_id_for_user - org_id = get_org_id_for_user(user_id) if user_id else None - verdict = evaluate_compound_command(org_id, command) - if not verdict.allowed: - reason = (verdict.rule_description or "Matched organization policy")[:200] - logger.warning("Policy denied tailscale_ssh command for user %s (%s)", - user_id, reason) + # Unified gate: signature + org policy + LLM judge + HITL (foreground). + from utils.auth.command_gate import gate_command + gate = gate_command(user_id=user_id, tool_name="tailscale_ssh", command=command) + if not gate.allowed: + logger.warning("tailscale_ssh blocked for user %s (%s): %s", + user_id, gate.code, gate.block_reason[:200]) return json.dumps({ "success": False, - "error": f"Command blocked by organization policy: {reason}", - "code": "POLICY_DENIED", - "provider": "tailscale_ssh", - }) - - from utils.security.command_safety import evaluate_command - decision = evaluate_command(command, tool="tailscale_ssh", user_id=user_id, session_id=session_id) - if decision.blocked: - code = "SIGNATURE_MATCHED" if decision.layer == "signature_match" else "SAFETY_BLOCKED" - return json.dumps({ - "success": False, - "error": f"Command blocked by safety guardrail: {decision.reason}", - "code": code, + "error": gate.block_reason, + "code": gate.code, "provider": "tailscale_ssh", }) diff --git a/server/chat/backend/agent/tools/terminal_exec_tool.py b/server/chat/backend/agent/tools/terminal_exec_tool.py index cb3d55fbe..fdfd0d27f 100644 --- a/server/chat/backend/agent/tools/terminal_exec_tool.py +++ b/server/chat/backend/agent/tools/terminal_exec_tool.py @@ -256,19 +256,16 @@ def terminal_exec( has_shell_syntax = _has_shell_metacharacters(command) allow_routing = not force_shell and not has_shell_syntax - # Org command policy check (shared allow/deny firewall across all tools) - from utils.auth.command_policy import evaluate_compound_command - from utils.auth.stateless_auth import get_org_id_for_user - org_id = get_org_id_for_user(user_id) if user_id else None - verdict = evaluate_compound_command(org_id, command) - if not verdict.allowed: - reason = (verdict.rule_description or "Matched organization policy")[:200] - logger.warning("Policy denied terminal command for user %s (%s)", - user_id, reason) + # Unified gate: signature + org policy + LLM judge + HITL (foreground). + from utils.auth.command_gate import gate_command + gate = gate_command(user_id=user_id, tool_name="terminal_exec", command=command) + if not gate.allowed: + logger.warning("terminal_exec blocked for user %s (%s): %s", + user_id, gate.code, gate.block_reason[:200]) return json.dumps({ "success": False, - "error": f"Command blocked by organization policy: {reason}", - "code": "POLICY_DENIED", + "error": gate.block_reason, + "code": gate.code, }) # Define routing table for cloud commands diff --git a/server/utils/auth/command_gate.py b/server/utils/auth/command_gate.py new file mode 100644 index 000000000..a7a831677 --- /dev/null +++ b/server/utils/auth/command_gate.py @@ -0,0 +1,315 @@ +"""Command-execution gate — unified policy + safety enforcement with HITL. + +Centralizes the four-layer defense-in-depth check that every shell tool must +run before execution: + + 1. Signature match (utils/security/signature_match.py via command_safety) + 2. Org allow/deny (utils/auth/command_policy.py) + 3. LLM safety judge (utils/security/command_safety.py) + +Behavior: + +* **Background chats** (``State.is_background == True``): on any block, returns + a deny decision with the layer's reason. Matches the pre-existing invariant + that destructive/denied actions cannot execute without an interactive user. +* **Foreground chats**: on any block, prompts the user via the WebSocket HITL + channel with Yes / No / (optionally) Yes-Always. + - **Yes**: approve this single invocation. Tool result looks like a normal + success to the agent (intentionally: we don't teach the LLM to reason + about the gate). + - **No**: abort with ``code="USER_DECLINED"``, distinct from policy/safety + codes so the agent sees explicit user rejection rather than a static + rule failure. + - **Yes-Always**: only offered when the block originated from + ``org_command_policies`` (deny rule hit or allowlist exhausted). Applies + the (possibly user-edited) policy mutation and then allows this + invocation. Future runs — including background RCAs — inherit the change. + +The gate has no independent on/off switch: when ``GUARDRAILS_ENABLED=false`` +and both org lists are disabled, no layer blocks anything, so the prompt +never fires. The gate is strictly the interactive surface of the existing +security layers. + +Two contextvars prevent duplicate prompts and duplicate guardrail LLM calls +for a single logical command as it passes through multiple tool layers +(e.g. ``terminal_exec`` routing into ``cloud_exec``): + +* ``_gate_inflight_command`` — set to the command hash during gating; re-entry + with a matching hash is a no-op "already approved" result. +* ``_guardrails_approved_command`` — read by ``terminal_run._check_guardrails`` + to skip the redundant signature+judge call on the agent path. Direct callers + (no contextvar set) still run the full check. +""" + +from __future__ import annotations + +import contextvars +import hashlib +import logging +from dataclasses import dataclass +from typing import Optional + +logger = logging.getLogger(__name__) + +_gate_inflight_command: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar( + "_gate_inflight_command", default=None, +) +_guardrails_approved_command: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar( + "_guardrails_approved_command", default=None, +) + + +def _hash(command: str) -> str: + return hashlib.sha256(command.encode("utf-8", errors="replace")).hexdigest() + + +def guardrails_approved_hash() -> Optional[str]: + """Accessor for ``terminal_run._check_guardrails`` to skip duplicate checks.""" + return _guardrails_approved_command.get() + + +@dataclass(frozen=True) +class GateDecision: + allowed: bool + code: str = "" # "" on allow, otherwise POLICY_DENIED / SAFETY_BLOCKED / + # SIGNATURE_MATCHED / USER_DECLINED + block_reason: str = "" + + +_ALLOWED = GateDecision(allowed=True) + + +def _block(code: str, reason: str) -> GateDecision: + return GateDecision(allowed=False, code=code, block_reason=reason) + + +def _is_foreground() -> bool: + """Return True iff the current execution is a foreground (interactive) chat.""" + try: + from utils.cloud.cloud_utils import get_state_context + state = get_state_context() + if state is None: + return False + return not bool(getattr(state, "is_background", False)) + except Exception: + return False + + +def _session_id() -> Optional[str]: + try: + from utils.cloud.cloud_utils import get_state_context + state = get_state_context() + return getattr(state, "session_id", None) if state else None + except Exception: + return None + + +def gate_command( + *, + user_id: Optional[str], + tool_name: str, + command: str, +) -> GateDecision: + """Run the full pre-execution gauntlet for *command*. + + Returns a :class:`GateDecision`. The caller is responsible for converting a + blocked decision into the tool's error response (``{"success": False, + "error": decision.block_reason, "code": decision.code}``). + """ + if not user_id: + # Without a user there is no org context and no HITL channel; defer to + # existing per-tool behavior by allowing through. Individual tools + # still enforce their own auth. + return _ALLOWED + + cmd_hash = _hash(command) + if _gate_inflight_command.get() == cmd_hash: + # Re-entry for the same command (e.g. terminal_exec -> cloud_exec). + return _ALLOWED + + token = _gate_inflight_command.set(cmd_hash) + try: + return _gate_impl(user_id=user_id, tool_name=tool_name, command=command, + cmd_hash=cmd_hash) + finally: + _gate_inflight_command.reset(token) + + +def _gate_impl(*, user_id: str, tool_name: str, command: str, cmd_hash: str) -> GateDecision: + from utils.auth.command_policy import ( + evaluate_compound_command, CommandVerdict, plan_yes_always, + apply_yes_always, validate_pattern, PolicyChange, + ) + from utils.auth.stateless_auth import get_org_id_for_user + from utils.security.command_safety import evaluate_command as safety_evaluate + + org_id = get_org_id_for_user(user_id) + foreground = _is_foreground() + session_id = _session_id() + + # --- Layer 1 & 3: signature_match + LLM judge ---------------------------- + # command_safety.evaluate_command bundles both: signature first (fast, + # cached), LLM judge second (network call). We read decision.layer to know + # which fired. + safety_decision = safety_evaluate( + command, tool=tool_name, user_id=user_id, session_id=session_id, + ) + if safety_decision.blocked: + layer = safety_decision.layer or "llm_judge" + code = "SIGNATURE_MATCHED" if layer == "signature_match" else "SAFETY_BLOCKED" + reason = f"Command blocked by safety guardrail: {safety_decision.reason}" + if not foreground: + return _block(code, reason) + # Foreground: ask user. No Yes-Always for this layer because there is + # no policy slot to flip. + return _prompt_user( + user_id=user_id, + session_id=session_id, + tool_name=tool_name, + command=command, + block_code=code, + block_reason=reason, + block_layer=layer, + allow_yes_always=False, + yes_always_changes=[], + org_id=None, + cmd_hash=cmd_hash, + ) + + # --- Layer 2: org allow/deny -------------------------------------------- + policy_verdict: CommandVerdict = evaluate_compound_command(org_id, command) + if policy_verdict.allowed: + # Tell terminal_run._check_guardrails it may skip re-running + # signature+judge for this command on the same invocation. + _guardrails_approved_command.set(cmd_hash) + return _ALLOWED + + reason = (policy_verdict.rule_description or "Matched organization policy")[:200] + block_reason = f"Command blocked by organization policy: {reason}" + layer = ( + "policy_both" if policy_verdict.deny_rule_id and policy_verdict.allowlist_exhausted + else "policy_deny" if policy_verdict.deny_rule_id + else "policy_allow_exhausted" + ) + if not foreground or not org_id: + return _block("POLICY_DENIED", block_reason) + + changes = plan_yes_always(policy_verdict, command) + decision = _prompt_user( + user_id=user_id, + session_id=session_id, + tool_name=tool_name, + command=command, + block_code="POLICY_DENIED", + block_reason=block_reason, + block_layer=layer, + allow_yes_always=bool(changes), + yes_always_changes=changes, + org_id=org_id, + cmd_hash=cmd_hash, + ) + # On Yes / Yes-Always, also clear the guardrails-approved marker so the + # downstream signature+judge re-check (if any) still runs normally for + # commands that bypassed the policy via user consent — the safety layers + # already passed above. + if decision.allowed: + _guardrails_approved_command.set(cmd_hash) + return decision + + +def _prompt_user( + *, + user_id: str, + session_id: Optional[str], + tool_name: str, + command: str, + block_code: str, + block_reason: str, + block_layer: str, + allow_yes_always: bool, + yes_always_changes: list, + org_id: Optional[str], + cmd_hash: str, +) -> GateDecision: + """Ask the user Yes / No / (Yes-Always) and apply the chosen effect.""" + from utils.cloud.infrastructure_confirmation import wait_for_user_confirmation_ex + from utils.auth.command_policy import apply_yes_always, validate_pattern + + options = [ + {"text": "Yes", "value": "execute"}, + {"text": "No", "value": "cancel"}, + ] + extra = { + "block_layer": block_layer, + "block_reason": block_reason, + "command": command, + } + if allow_yes_always and yes_always_changes: + options.append({"text": "Yes, Always", "value": "execute_always"}) + extra["yes_always_effect"] = { + "summary": "This will modify your organization's command policy:", + "changes": [ + { + "action": ch.action, + "rule_id": ch.rule_id, + "pattern": ch.pattern, + "description": ch.description, + "editable": ch.editable, + } + for ch in yes_always_changes + ], + } + + # A single-line human prompt shown in the UI alongside the options. + message = f"{tool_name}: {block_reason}. Command: {command[:500]}" + result = wait_for_user_confirmation_ex( + user_id=user_id, + message=message, + tool_name=tool_name, + session_id=session_id, + options=options, + extra=extra, + ) + decision = result.get("decision") + + if decision == "execute": + return _ALLOWED + if decision == "execute_always" and allow_yes_always and org_id: + edited = result.get("edited_patterns") or {} + applied = [] + for idx, ch in enumerate(yes_always_changes): + if ch.action == "add_allow_rule" and ch.editable: + # Users may tighten/loosen the pattern via the editable input. + # Indices are sent as strings by JSON. + override = edited.get(str(idx)) or edited.get(idx) + pattern = (override or ch.pattern or "").strip() + err = validate_pattern(pattern) if pattern else "empty pattern" + if err: + logger.warning( + "[CommandGate] Yes-Always rejected: invalid regex '%s' (%s). " + "Treating as cancel.", pattern, err, + ) + return _block( + "USER_DECLINED", + f"Tool call not allowed by user (invalid regex: {err})", + ) + applied.append(type(ch)( + action=ch.action, rule_id=ch.rule_id, pattern=pattern, + description=ch.description, editable=ch.editable, + )) + else: + applied.append(ch) + try: + apply_yes_always(org_id, applied, user_id) + logger.info( + "[CommandGate] Yes-Always applied %d change(s) for org %s by %s", + len(applied), org_id, user_id, + ) + except Exception: + logger.exception("[CommandGate] Failed to persist Yes-Always changes") + return _block("POLICY_DENIED", + "Failed to update organization policy; command not executed") + return _ALLOWED + + # Timeout or explicit cancel: USER_DECLINED. + return _block("USER_DECLINED", "Tool call not allowed by user") diff --git a/server/utils/auth/command_policy.py b/server/utils/auth/command_policy.py index adf44b236..ad61827ce 100644 --- a/server/utils/auth/command_policy.py +++ b/server/utils/auth/command_policy.py @@ -17,6 +17,7 @@ import logging import re +import shlex import time from dataclasses import dataclass, field from typing import Dict, List, Optional, Tuple @@ -40,6 +41,10 @@ class CommandVerdict: allowed: bool rule_description: Optional[str] = None + # Populated only when denylist blocked the command (id of the matched deny rule). + deny_rule_id: Optional[int] = None + # True when allowlist was enabled and no allow rule matched. + allowlist_exhausted: bool = False @dataclass(frozen=True) @@ -141,13 +146,21 @@ def evaluate_command(org_id: Optional[str], command: str) -> CommandVerdict: if states.denylist_enabled: for rule in deny_rules: if rule.compiled.search(command): - return CommandVerdict(allowed=False, rule_description=rule.description) + return CommandVerdict( + allowed=False, + rule_description=rule.description, + deny_rule_id=rule.id, + ) if states.allowlist_enabled: for rule in allow_rules: if rule.compiled.search(command): return CommandVerdict(allowed=True, rule_description=rule.description) - return CommandVerdict(allowed=False, rule_description="No matching allow rule") + return CommandVerdict( + allowed=False, + rule_description="No matching allow rule", + allowlist_exhausted=True, + ) return CommandVerdict(allowed=True) @@ -286,6 +299,117 @@ def validate_pattern(pattern: str) -> Optional[str]: return str(exc) +# Leading `VAR=value` env assignments the user didn't intend as the "command". +_ENV_ASSIGN_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=\S*$") + + +def derive_pattern_from_command(command: str) -> str: + """Propose a conservative allow-rule regex for *command*. + + Strips leading ``sudo`` and leading ``VAR=value`` env assignments, then + anchors on the CLI name plus its first non-flag subcommand. Falls back to + the CLI name alone when there is no subcommand. + + Examples: + "sudo kubectl delete pod foo" -> ^kubectl delete\\b + "KUBECONFIG=/x kubectl get pods" -> ^kubectl get\\b + "aws ec2 terminate-instances --id" -> ^aws ec2 terminate-instances\\b + "terraform apply -auto-approve" -> ^terraform apply\\b + """ + try: + tokens = shlex.split(command.strip(), posix=True) + except ValueError: + tokens = command.strip().split() + # Drop leading sudo / env assignments. + while tokens and (tokens[0] == "sudo" or _ENV_ASSIGN_RE.match(tokens[0])): + tokens.pop(0) + if not tokens: + return r"^" + re.escape(command.strip()) + r"\b" + parts = [tokens[0]] + for tok in tokens[1:]: + if tok.startswith("-"): + break + parts.append(tok) + # Stop after CLI + first subcommand to keep the pattern reasonably + # narrow. Callers (UI) can edit before applying. + if len(parts) >= 2: + break + return r"^" + r"\s+".join(re.escape(p) for p in parts) + r"\b" + + +@dataclass(frozen=True) +class PolicyChange: + """One mutation that Yes-Always will apply to org_command_policies.""" + action: str # "disable_deny_rule" | "add_allow_rule" + rule_id: Optional[int] = None + pattern: Optional[str] = None + description: Optional[str] = None + editable: bool = False + + +def plan_yes_always(verdict: CommandVerdict, command: str) -> List[PolicyChange]: + """Build the list of policy mutations implied by clicking Yes-Always. + + Pure function — returns the plan without touching the DB. The caller + renders this to the user and then calls :func:`apply_yes_always` with + (optionally edited) patterns. + """ + changes: List[PolicyChange] = [] + if verdict.deny_rule_id is not None: + changes.append(PolicyChange( + action="disable_deny_rule", + rule_id=verdict.deny_rule_id, + description=verdict.rule_description or "", + editable=False, + )) + if verdict.allowlist_exhausted: + changes.append(PolicyChange( + action="add_allow_rule", + pattern=derive_pattern_from_command(command), + description="Auto-approved from chat", + editable=True, + )) + return changes + + +def apply_yes_always( + org_id: str, + changes: List[PolicyChange], + user_id: str, +) -> None: + """Persist Yes-Always mutations atomically and invalidate the cache. + + Each change in ``changes`` is either ``disable_deny_rule`` (soft-disable the + referenced rule) or ``add_allow_rule`` (insert a new allow rule with the + user-confirmed pattern). Patterns must already be validated by the caller. + """ + if not changes: + return + from utils.db.connection_pool import db_pool + with db_pool.get_admin_connection() as conn: + with conn.cursor() as cur: + cur.execute("SET myapp.current_org_id = %s", (org_id,)) + for ch in changes: + if ch.action == "disable_deny_rule" and ch.rule_id is not None: + cur.execute( + "UPDATE org_command_policies " + "SET enabled = false, updated_at = NOW(), updated_by = %s " + "WHERE id = %s AND org_id = %s AND mode = 'deny'", + (user_id, ch.rule_id, org_id), + ) + elif ch.action == "add_allow_rule" and ch.pattern: + cur.execute( + "INSERT INTO org_command_policies " + "(org_id, mode, pattern, description, priority, updated_by, source) " + "VALUES (%s, 'allow', %s, %s, %s, %s, 'custom') " + "ON CONFLICT (org_id, mode, pattern, source) DO NOTHING", + (org_id, ch.pattern, ch.description or "Auto-approved from chat", + 50, user_id), + ) + conn.commit() + invalidate_cache(org_id) + + def get_policy_prompt_text(org_id: str) -> str: """Render active policy as system-prompt text for the LLM.""" allow_rules, deny_rules, states = _get_cached(org_id) diff --git a/server/utils/cloud/infrastructure_confirmation.py b/server/utils/cloud/infrastructure_confirmation.py index 901125b8a..f42daf188 100644 --- a/server/utils/cloud/infrastructure_confirmation.py +++ b/server/utils/cloud/infrastructure_confirmation.py @@ -5,7 +5,7 @@ import logging import time import uuid -from typing import Dict, Any +from typing import Dict, Any, Optional # Lazy imports to avoid circular dependency with cloud_tools.py logger = logging.getLogger(__name__) @@ -54,6 +54,10 @@ async def handle_websocket_confirmation_response(data: dict): if confirmation_id in _pending_confirmations: confirmation_data = _pending_confirmations[confirmation_id] confirmation_data['result'] = decision + # Carry optional per-change edited patterns for Yes-Always flows. + edited = data.get('edited_patterns') + if isinstance(edited, dict): + confirmation_data['edited_patterns'] = edited logger.debug(f"WEBSOCKET: Confirmation {confirmation_id} resolved with decision: {decision}") else: logger.warning(f"WEBSOCKET: No pending confirmation found for ID: {confirmation_id}") @@ -196,6 +200,103 @@ def wait_for_user_confirmation( return decision == "execute" +def wait_for_user_confirmation_ex( + *, + user_id: str, + message: str, + tool_name: str, + session_id: Optional[str], + options: list, + extra: Optional[Dict[str, Any]] = None, + workflow_instance=None, + timeout_seconds: int = 300, +) -> Dict[str, Any]: + """Extended HITL helper that returns the full response payload. + + Unlike :func:`wait_for_user_confirmation` (which returns a bool), this + preserves the user's exact decision (e.g. ``execute`` vs ``execute_always`` + vs ``cancel``) plus any ``edited_patterns`` map the frontend sends back. + + Returns ``{"decision": str | None, "edited_patterns": dict}``. ``decision`` + is ``None`` only on timeout. ``edited_patterns`` maps stringified change + indices to user-edited regex patterns. + """ + # Background chats: no interactive user. + try: + from chat.backend.agent.tools.cloud_tools import get_state_context + state = get_state_context() + if state and getattr(state, 'is_background', False): + logger.warning(f"[BackgroundChat] Denying confirmation for {tool_name} -- no interactive user") + return {"decision": "cancel", "edited_patterns": {}} + except Exception as e: + logger.debug(f"Could not check background state: {e}") + + timestamp_ms = int(time.time() * 1000) + unique_id = str(uuid.uuid4())[:8] + confirmation_id = f"{timestamp_ms}:{unique_id}" + + payload: Dict[str, Any] = { + "type": "execution_confirmation", + "data": { + "message": message, + "status": "awaiting_confirmation", + "user_id": user_id, + "confirmation_id": confirmation_id, + "tool_name": tool_name, + "options": options, + }, + } + if extra: + payload["data"].update(extra) + if session_id: + payload["session_id"] = session_id + if user_id: + payload["user_id"] = user_id + + if not workflow_instance: + from chat.backend.agent.tools.cloud_tools import get_workflow_context + workflow_instance = get_workflow_context() + + if session_id and user_id and workflow_instance: + try: + workflow_instance._consolidate_message_chunks() + _save_ui_messages_with_confirmation_via_workflow( + workflow_instance, session_id, user_id, tool_name, message, confirmation_id + ) + except Exception as e: + logger.error(f"Error consolidating/saving UI messages before confirmation: {e}") + + _send_ws(payload, tool_name) + + _pending_confirmations[confirmation_id] = { + 'result': None, + 'user_id': user_id, + 'timestamp': time.time(), + } + logger.debug(f"WEBSOCKET: Waiting for confirmation_ex {confirmation_id}") + + decision: Optional[str] = None + edited: Dict[str, Any] = {} + try: + elapsed, poll_interval = 0.0, 1.0 + while elapsed < timeout_seconds: + data = _pending_confirmations.get(confirmation_id) + if data and data.get('result'): + decision = data['result'] + edited = data.get('edited_patterns') or {} + break + time.sleep(poll_interval) + elapsed += poll_interval + else: + logger.warning(f"WEBSOCKET: Timeout waiting for confirmation_ex {confirmation_id}") + except Exception as e: + logger.error(f"Error waiting for confirmation_ex: {e}") + finally: + _pending_confirmations.pop(confirmation_id, None) + + return {"decision": decision, "edited_patterns": edited} + + def _save_ui_messages_with_confirmation_via_workflow(workflow_instance, session_id: str, user_id: str, tool_name: str, message: str, confirmation_id: str) -> bool: """Save UI messages to database including any ongoing tool calls and confirmation marker. Uses workflow's methods to avoid code duplication.""" diff --git a/server/utils/terminal/terminal_run.py b/server/utils/terminal/terminal_run.py index 8361581de..abbf823e0 100644 --- a/server/utils/terminal/terminal_run.py +++ b/server/utils/terminal/terminal_run.py @@ -220,6 +220,20 @@ def _check_guardrails(args: Union[str, List[str]]) -> Optional[CompletedProcess] cmd = args if isinstance(args, str) else shlex.join(str(a) for a in args) + # Agent-path short-circuit: command_gate ran the same signature+judge check + # moments ago for this exact command. Skipping avoids a second LLM call + # (cheap but not cached) and eliminates the small risk of a divergent + # verdict on the second run. Direct callers (no contextvar set) still + # execute the full check and fail closed as before. + try: + from utils.auth.command_gate import guardrails_approved_hash + import hashlib + approved = guardrails_approved_hash() + if approved and approved == hashlib.sha256(cmd.encode("utf-8", errors="replace")).hexdigest(): + return None + except Exception: + logger.debug("[Guardrails] command_gate bypass unavailable", exc_info=True) + uid, sid = None, None try: ctx = get_user_context() From 8a140d07591e0d6a4b244c799c5b01ac95fa5ef0 Mon Sep 17 00:00:00 2001 From: Damian Loch Date: Mon, 27 Apr 2026 17:55:57 -0400 Subject: [PATCH 02/13] edit to dynamic change of denylist/allowlist --- server/utils/auth/command_policy.py | 39 +++++++++++++++++++++-------- 1 file changed, 28 insertions(+), 11 deletions(-) diff --git a/server/utils/auth/command_policy.py b/server/utils/auth/command_policy.py index ad61827ce..86ab75dd1 100644 --- a/server/utils/auth/command_policy.py +++ b/server/utils/auth/command_policy.py @@ -302,29 +302,46 @@ def validate_pattern(pattern: str) -> Optional[str]: # Leading `VAR=value` env assignments the user didn't intend as the "command". _ENV_ASSIGN_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=\S*$") +# Shell keywords / operators that mean "the first token is not a CLI name". +# When we see one of these at position 0 (after stripping sudo / env assigns), +# the command is a shell compound statement, pipeline, subshell, or similar +# construct whose first token is not a useful anchor for an allow rule. In +# that case we fall back to the full regex-escaped command so the proposed +# rule only matches this exact invocation -- the user can relax it in the +# editable UI field. +_SHELL_NON_CLI_LEADERS = frozenset({ + "for", "while", "until", "if", "case", "select", "time", "function", + "{", "(", "[", "[[", "!", "coproc", +}) + def derive_pattern_from_command(command: str) -> str: """Propose a conservative allow-rule regex for *command*. Strips leading ``sudo`` and leading ``VAR=value`` env assignments, then - anchors on the CLI name plus its first non-flag subcommand. Falls back to - the CLI name alone when there is no subcommand. + anchors on the CLI name plus its first non-flag subcommand. + + When the leading token after stripping is a shell keyword (``for``, + ``while``, ``if``, subshell ``(``, brace group ``{`` ...) the first token + is not a CLI name, so instead we anchor on the full regex-escaped command. + That pattern only matches the exact invocation; it is intentionally narrow + because the user can loosen it in the edit box before confirming. Examples: - "sudo kubectl delete pod foo" -> ^kubectl delete\\b - "KUBECONFIG=/x kubectl get pods" -> ^kubectl get\\b - "aws ec2 terminate-instances --id" -> ^aws ec2 terminate-instances\\b - "terraform apply -auto-approve" -> ^terraform apply\\b + "sudo kubectl delete pod foo" -> ^kubectl delete\\b + "KUBECONFIG=/x kubectl get pods" -> ^kubectl get\\b + "aws ec2 terminate-instances --id" -> ^aws ec2 terminate-instances\\b + "for i in {1..10}; do echo $i; done" -> ^for i in \\{1\\.\\.10\\}; do echo \\$i; done$ """ + stripped = command.strip() try: - tokens = shlex.split(command.strip(), posix=True) + tokens = shlex.split(stripped, posix=True) except ValueError: - tokens = command.strip().split() - # Drop leading sudo / env assignments. + tokens = stripped.split() while tokens and (tokens[0] == "sudo" or _ENV_ASSIGN_RE.match(tokens[0])): tokens.pop(0) - if not tokens: - return r"^" + re.escape(command.strip()) + r"\b" + if not tokens or tokens[0] in _SHELL_NON_CLI_LEADERS: + return r"^" + re.escape(stripped) + r"$" parts = [tokens[0]] for tok in tokens[1:]: if tok.startswith("-"): From fab31cc44237c471b23dd48626cccc26956ed3ad Mon Sep 17 00:00:00 2001 From: Damian Loch Date: Tue, 28 Apr 2026 09:13:43 -0400 Subject: [PATCH 03/13] edits --- .../tool-calls/ToolExecutionWidget.tsx | 179 +++++++++++------- client/src/hooks/useMessageHandler.ts | 1 + server/utils/auth/command_gate.py | 6 +- 3 files changed, 116 insertions(+), 70 deletions(-) diff --git a/client/src/components/tool-calls/ToolExecutionWidget.tsx b/client/src/components/tool-calls/ToolExecutionWidget.tsx index e70af71d5..af3ce4efc 100644 --- a/client/src/components/tool-calls/ToolExecutionWidget.tsx +++ b/client/src/components/tool-calls/ToolExecutionWidget.tsx @@ -5,8 +5,9 @@ import type { JSX } from "react" import { Card } from "@/components/ui/card" import { Skeleton } from "@/components/ui/skeleton" import { Button } from "@/components/ui/button" +import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover" import { cn } from "@/lib/utils" -import { ChevronDown, ChevronUp, X, Check, AlertCircle } from "lucide-react" +import { ChevronDown, ChevronUp, AlertCircle, Settings2 } from "lucide-react" import CommandLogo from "./CommandLogo" import { useTheme } from "next-themes" import { GitHubCommitTool } from "@/components/GitHubCommitTool" @@ -233,7 +234,13 @@ const ToolExecutionWidget = ({ tool, className, sendMessage, sendRaw, onToolUpda // cloud_exec parsing else if (tool.tool_name === "cloud_exec") { const parsed = parseCloudExecCommand(normalizedInput, tool.output, defaultCliCommand) - command = parsed.command + // Prefer the authoritative command from the gate's confirmation payload + // when the tool is paused for approval -- parseCloudExecCommand returns a + // placeholder ("cloud exec") when there's no final_command in the output + // yet, which is exactly the state we're in while awaiting confirmation. + command = tool.command && parsed.command === defaultCliCommand + ? tool.command + : parsed.command } // GitHub MCP tools parsing else if (tool.tool_name.startsWith("mcp_") && typeof command === "string" && command.trim().startsWith("{")) { @@ -468,6 +475,7 @@ const ToolExecutionWidget = ({ tool, className, sendMessage, sendRaw, onToolUpda {tool.status === "awaiting_confirmation" && !tool.output && !tool.error && ( boolean onToolUpdate?: (updatedTool: Partial) => void } -const ConfirmationPanel = ({ tool, userId, sessionId, sendRaw, onToolUpdate }: ConfirmationPanelProps) => { +// Compact human summary of a shell command: CLI name + first non-flag +// subcommand (e.g. "aws ec2 describe-instances --query ..." -> "aws ec2 +// describe-instances"). Falls back to the raw command if it doesn't parse. +const summarizeCommand = (cmd: string): string => { + if (!cmd) return "this command" + const trimmed = cmd.trim() + const tokens = trimmed.split(/\s+/) + const parts: string[] = [] + for (const tok of tokens) { + if (tok.startsWith("-")) break + parts.push(tok) + if (parts.length >= 3) break + } + return parts.join(" ") || trimmed.slice(0, 40) +} + +const ConfirmationPanel = ({ tool, command, userId, sessionId, sendRaw, onToolUpdate }: ConfirmationPanelProps) => { const effect = tool.yes_always_effect const allowYesAlways = !!(effect && effect.changes.length > 0) + const summary = summarizeCommand(command) - // Pre-fill editable change patterns. Map is keyed by change index (as string - // to match the backend's JSON-parsed dict key). const [edited, setEdited] = React.useState>(() => { const m: Record = {} if (effect) { @@ -567,6 +591,7 @@ const ConfirmationPanel = ({ tool, userId, sessionId, sendRaw, onToolUpdate }: C } return m }) + const [alwaysOpen, setAlwaysOpen] = React.useState(false) const respond = (decision: 'execute' | 'cancel' | 'execute_always') => { const confirmationId = tool.confirmation_id @@ -588,71 +613,89 @@ const ConfirmationPanel = ({ tool, userId, sessionId, sendRaw, onToolUpdate }: C } return ( -
-
- - {tool.confirmation_message || "This action requires confirmation"} +
+
+ + + Approval needed for {summary} -
- - -
- - {allowYesAlways && effect && ( -
-
{effect.summary}
-
    - {effect.changes.map((c, i) => ( -
  • - {c.action === 'disable_deny_rule' ? ( - - Disable deny rule: {c.description || 'rule'} - {c.pattern ? <> (pattern: {c.pattern}) : null} - - ) : ( - <> - Add allow rule (you can edit the pattern): - setEdited(prev => ({ ...prev, [String(i)]: e.target.value }))} - spellCheck={false} - aria-label="Allow-rule regex pattern" - /> - - )} -
    • - ))} -
-
- -
-
- )} +
+ + + {allowYesAlways && effect && ( + + + + + +
{effect.summary}
+
    + {effect.changes.map((c, i) => ( +
  • + {c.action === 'disable_deny_rule' ? ( + <> +
    + Disable deny rule: {c.description || 'rule'} +
    + {c.pattern && ( + + {c.pattern} + + )} + + ) : ( + <> + + setEdited(prev => ({ ...prev, [String(i)]: e.target.value }))} + spellCheck={false} + aria-label="Allow-rule regex pattern" + /> + + )} +
    • + ))} +
+
+ +
+ + + )} +
) } diff --git a/client/src/hooks/useMessageHandler.ts b/client/src/hooks/useMessageHandler.ts index 00fa5787d..c332b4ccd 100644 --- a/client/src/hooks/useMessageHandler.ts +++ b/client/src/hooks/useMessageHandler.ts @@ -116,6 +116,7 @@ export const useMessageHandler = ({ status: 'awaiting_confirmation' as const, confirmation_id: confirmationId, confirmation_message: message.data.message, + command: message.data.command ?? tc.command, block_layer: message.data.block_layer, yes_always_effect: message.data.yes_always_effect, }; diff --git a/server/utils/auth/command_gate.py b/server/utils/auth/command_gate.py index a7a831677..c76b840ac 100644 --- a/server/utils/auth/command_gate.py +++ b/server/utils/auth/command_gate.py @@ -260,8 +260,10 @@ def _prompt_user( ], } - # A single-line human prompt shown in the UI alongside the options. - message = f"{tool_name}: {block_reason}. Command: {command[:500]}" + # Compact label for the UI. The full reason and command ride along in + # ``extra`` (``block_reason`` / ``command``) for tooltips, logging, and + # the pattern derivation in the Yes-Always popover. + message = "Approval needed" result = wait_for_user_confirmation_ex( user_id=user_id, message=message, From ff48309bd2882fb5de70ff2d24112a6b1e702f42 Mon Sep 17 00:00:00 2001 From: Damian Loch Date: Tue, 28 Apr 2026 09:50:17 -0400 Subject: [PATCH 04/13] fixed ui state --- client/src/components/SecuritySettings.tsx | 3 +- server/routes/command_policies.py | 4 + server/utils/auth/command_policy.py | 44 +++++--- .../cloud/infrastructure_confirmation.py | 102 ++++++++++++++++-- 4 files changed, 131 insertions(+), 22 deletions(-) diff --git a/client/src/components/SecuritySettings.tsx b/client/src/components/SecuritySettings.tsx index 6fa605793..ecb5f2d85 100644 --- a/client/src/components/SecuritySettings.tsx +++ b/client/src/components/SecuritySettings.tsx @@ -104,7 +104,8 @@ function AddRuleForm({ diff --git a/server/routes/command_policies.py b/server/routes/command_policies.py index 27003eb79..47cf1ceed 100644 --- a/server/routes/command_policies.py +++ b/server/routes/command_policies.py @@ -102,6 +102,8 @@ def create_policy(user_id): return jsonify({"error": "mode must be 'allow' or 'deny'"}), 400 if not pattern: return jsonify({"error": "pattern is required"}), 400 + if not description: + return jsonify({"error": "description is required"}), 400 err = validate_pattern(pattern) if err: @@ -150,6 +152,8 @@ def update_policy(user_id, rule_id): err = validate_pattern(data[field]) if err: return jsonify({"error": "Invalid regex pattern"}), 400 + if field == "description" and not str(data[field] or "").strip(): + return jsonify({"error": "description is required"}), 400 updates.append(f"{col} = %s") params.append(data[field]) diff --git a/server/utils/auth/command_policy.py b/server/utils/auth/command_policy.py index 86ab75dd1..bed4220c8 100644 --- a/server/utils/auth/command_policy.py +++ b/server/utils/auth/command_policy.py @@ -133,7 +133,14 @@ def _get_cached(org_id: str) -> Tuple[List[PolicyRule], List[PolicyRule], ListSt def evaluate_command(org_id: Optional[str], command: str) -> CommandVerdict: - """Core gate. Returns whether *command* is allowed for *org_id*.""" + """Core gate. Returns whether *command* is allowed for *org_id*. + + Both lists are evaluated independently so the verdict carries the full + picture: a deny-list hit still reports whether the allowlist would have + matched, which is what the Yes-Always planner needs to decide whether + clicking ``Always`` should also add an allow rule alongside disabling + the deny rule. + """ if not org_id: logger.info("policy_check_skipped reason=no_org_context func=evaluate_command") return CommandVerdict(allowed=True) @@ -143,26 +150,35 @@ def evaluate_command(org_id: Optional[str], command: str) -> CommandVerdict: if not states.denylist_enabled and not states.allowlist_enabled: return CommandVerdict(allowed=True, rule_description="Policy lists are disabled") - if states.denylist_enabled: - for rule in deny_rules: - if rule.compiled.search(command): - return CommandVerdict( - allowed=False, - rule_description=rule.description, - deny_rule_id=rule.id, - ) + deny_hit = next( + (r for r in deny_rules if r.compiled.search(command)), + None, + ) if states.denylist_enabled else None + + allow_hit = next( + (r for r in allow_rules if r.compiled.search(command)), + None, + ) if states.allowlist_enabled else None + + if deny_hit: + return CommandVerdict( + allowed=False, + rule_description=deny_hit.description, + deny_rule_id=deny_hit.id, + allowlist_exhausted=(states.allowlist_enabled and allow_hit is None), + ) - if states.allowlist_enabled: - for rule in allow_rules: - if rule.compiled.search(command): - return CommandVerdict(allowed=True, rule_description=rule.description) + if states.allowlist_enabled and allow_hit is None: return CommandVerdict( allowed=False, rule_description="No matching allow rule", allowlist_exhausted=True, ) - return CommandVerdict(allowed=True) + return CommandVerdict( + allowed=True, + rule_description=allow_hit.description if allow_hit else None, + ) _UNSPLITTABLE_SHELL_RE = re.compile(r"<<-?\s*\w+|<\(|>\(") diff --git a/server/utils/cloud/infrastructure_confirmation.py b/server/utils/cloud/infrastructure_confirmation.py index f42daf188..2c9c708d3 100644 --- a/server/utils/cloud/infrastructure_confirmation.py +++ b/server/utils/cloud/infrastructure_confirmation.py @@ -5,7 +5,7 @@ import logging import time import uuid -from typing import Dict, Any, Optional +from typing import Dict, Any, List, Optional # Lazy imports to avoid circular dependency with cloud_tools.py logger = logging.getLogger(__name__) @@ -261,7 +261,8 @@ def wait_for_user_confirmation_ex( try: workflow_instance._consolidate_message_chunks() _save_ui_messages_with_confirmation_via_workflow( - workflow_instance, session_id, user_id, tool_name, message, confirmation_id + workflow_instance, session_id, user_id, tool_name, message, confirmation_id, + extra=extra, ) except Exception as e: logger.error(f"Error consolidating/saving UI messages before confirmation: {e}") @@ -297,9 +298,24 @@ def wait_for_user_confirmation_ex( return {"decision": decision, "edited_patterns": edited} -def _save_ui_messages_with_confirmation_via_workflow(workflow_instance, session_id: str, user_id: str, tool_name: str, message: str, confirmation_id: str) -> bool: +def _save_ui_messages_with_confirmation_via_workflow( + workflow_instance, + session_id: str, + user_id: str, + tool_name: str, + message: str, + confirmation_id: str, + extra: Optional[Dict[str, Any]] = None, +) -> bool: """Save UI messages to database including any ongoing tool calls and confirmation marker. - Uses workflow's methods to avoid code duplication.""" + + When *extra* is provided (gate-driven confirmations), the awaiting tool call + is enriched with ``command`` / ``block_layer`` / ``yes_always_effect`` so the + Yes-Always popover can rehydrate correctly after a page reload. We also + merge-preserve any previously saved user messages that the rebuilt UI list + happens to omit (e.g. when the LLM state snapshot was taken mid-turn), so + navigating away does not silently drop the user's prompt. + """ try: # Get LLM messages from workflow state if workflow_instance._last_state: @@ -342,9 +358,17 @@ def _save_ui_messages_with_confirmation_via_workflow(workflow_instance, session_ for tool_call in ui_msg.get('toolCalls', []): # Match by tool_call_id (call_xxx format) if tool_call.get('id') == target_tool_call: - # Update the tool call to show it's awaiting confirmation tool_call['status'] = 'awaiting_confirmation' tool_call['confirmation_id'] = confirmation_id + tool_call['confirmation_message'] = message + if extra: + # Persist only the fields the confirmation UI + # needs on rehydrate; avoid leaking internal + # fields into the saved record. + for key in ("command", "block_layer", "yes_always_effect"): + val = extra.get(key) + if val is not None: + tool_call[key] = val logger.debug(f"Updated tool call {target_tool_call} to awaiting_confirmation status") tool_call_updated = True break @@ -356,11 +380,75 @@ def _save_ui_messages_with_confirmation_via_workflow(workflow_instance, session_ else: logger.warning(f"No tool call found for {tool_name} in current tool calls") logger.debug(f"Available tool calls: {list(tool_capture.current_tool_calls.keys())}") - + + # Preserve any previously-saved user messages the rebuild dropped. The + # LLM-state snapshot may not include the current turn's HumanMessage + # if the gate fires before LangGraph has persisted it; without this + # merge the unconditional overwrite would wipe the user's prompt. + ui_messages = _merge_missing_user_messages(session_id, user_id, ui_messages) + # Use workflow's _save_ui_messages method to save the updated messages logger.debug(f"Saving {len(ui_messages)} UI messages") return workflow_instance._save_ui_messages(session_id, user_id, ui_messages) except Exception as e: logger.error(f"Error saving UI messages with confirmation: {e}") - return False \ No newline at end of file + return False + + +def _merge_missing_user_messages( + session_id: str, user_id: str, rebuilt: List[Dict[str, Any]], +) -> List[Dict[str, Any]]: + """Re-insert persisted user messages that ``rebuilt`` is missing. + + The LLM-state snapshot may not include the current turn's HumanMessage, + in which case ``rebuilt`` starts with the assistant/tool response and + drops the user's prompt. We splice those missing user messages back in + at their original chronological position from the persisted row so the + UI renders the prompt above its response. Fails open on any DB error. + """ + try: + from utils.db.connection_pool import db_pool + with db_pool.get_user_connection() as conn: + with conn.cursor() as cur: + cur.execute( + "SELECT messages FROM chat_sessions WHERE id = %s AND user_id = %s", + (session_id, user_id), + ) + row = cur.fetchone() + existing = row[0] if row and row[0] else [] + if not isinstance(existing, list): + return rebuilt + except Exception as e: + logger.warning(f"merge_missing_user_messages: could not read existing messages: {e}") + return rebuilt + + rebuilt_user_texts = { + m.get('text') for m in rebuilt if m.get('sender') == 'user' and m.get('text') + } + missing_any = any( + m.get('sender') == 'user' and m.get('text') and m.get('text') not in rebuilt_user_texts + for m in existing + ) + if not missing_any: + return rebuilt + + # Walk ``existing`` in order. For each user message not present in the + # rebuilt set, insert a copy into ``merged`` at the position we reach in + # ``existing``; for every other slot we advance through ``rebuilt``. + merged: List[Dict[str, Any]] = [] + rebuilt_idx = 0 + for em in existing: + if (em.get('sender') == 'user' and em.get('text') + and em.get('text') not in rebuilt_user_texts): + merged.append(em) + continue + if rebuilt_idx < len(rebuilt): + merged.append(rebuilt[rebuilt_idx]) + rebuilt_idx += 1 + # Anything left in ``rebuilt`` is newer than ``existing`` — append as-is. + merged.extend(rebuilt[rebuilt_idx:]) + + for i, m in enumerate(merged, start=1): + m['message_number'] = i + return merged \ No newline at end of file From f78b28c581a670611600882549d7a0b6ac2e372f Mon Sep 17 00:00:00 2001 From: Damian Loch Date: Tue, 28 Apr 2026 10:05:22 -0400 Subject: [PATCH 05/13] update UI --- client/src/hooks/useChatHistory.ts | 35 +- server/routes/chat_routes.py | 4 +- .../cloud/infrastructure_confirmation.py | 356 ++++++------------ server/utils/db/db_utils.py | 17 + 4 files changed, 166 insertions(+), 246 deletions(-) diff --git a/client/src/hooks/useChatHistory.ts b/client/src/hooks/useChatHistory.ts index dc6ded63e..028aa72ed 100644 --- a/client/src/hooks/useChatHistory.ts +++ b/client/src/hooks/useChatHistory.ts @@ -51,6 +51,11 @@ export interface ChatMessage { error?: string | null; status: 'running' | 'completed' | 'error' | 'cancelled' | 'awaiting_confirmation'; timestamp: string; + confirmation_id?: string; + confirmation_message?: string; + command?: string; + block_layer?: string; + yes_always_effect?: unknown; }>; images?: Array<{ data: string; @@ -223,8 +228,36 @@ export function useChatHistory(): UseChatHistoryReturn { const rawMessages = data.messages || []; const cleanedMessages = cleanupStaleToolCalls(rawMessages, data.updated_at); const cleanupTime = performance.now() - cleanupStart; + // If the server has a live HITL confirmation pending for this session, + // append a synthetic bot message with an awaiting_confirmation tool + // call so the user sees the prompt immediately on reload. The card is + // driven entirely by chat_sessions.pending_turn -- history remains + // append-only and contains no mid-turn snapshot. + const pending = data.pending_turn; + const messagesWithPending: ChatMessage[] = pending && pending.confirmation_id + ? [ + ...(cleanedMessages as ChatMessage[]), + { + id: Date.now(), + sender: 'bot', + text: '', + toolCalls: [{ + id: `pending-${pending.confirmation_id}`, + tool_name: pending.tool_name || 'action', + input: '', + status: 'awaiting_confirmation', + timestamp: pending.created_at || new Date().toISOString(), + confirmation_id: pending.confirmation_id, + confirmation_message: pending.message, + command: pending.command, + block_layer: pending.block_layer, + yes_always_effect: pending.yes_always_effect, + }], + } as ChatMessage, + ] + : (cleanedMessages as ChatMessage[]); // DON'T refresh sessions when just loading - this prevents sessions from moving to top - return { messages: cleanedMessages as ChatMessage[], uiState, incidentId: data.incident_id || null, status: data.status || null }; + return { messages: messagesWithPending, uiState, incidentId: data.incident_id || null, status: data.status || null }; } catch (err) { console.error('[useChatHistory] Error loading chat session:', err); setCrudError(err instanceof Error ? err.message : 'Failed to load chat session'); diff --git a/server/routes/chat_routes.py b/server/routes/chat_routes.py index c0ff40ac3..ad8ab274f 100644 --- a/server/routes/chat_routes.py +++ b/server/routes/chat_routes.py @@ -199,7 +199,8 @@ def get_chat_session(user_id, session_id): CASE WHEN ui_state IS NULL THEN '{}'::jsonb ELSE ui_state END as ui_state, COALESCE(status, 'active') as status, user_id, - incident_id + incident_id, + pending_turn FROM chat_sessions WHERE id = %s AND org_id = %s AND is_active = true AND status != 'cancelled' """, (session_id, org_id)) @@ -276,6 +277,7 @@ def get_chat_session(user_id, session_id): 'user_id': session_data[7], 'is_own': session_data[7] == user_id, 'incident_id': str(session_data[8]) if session_data[8] else None, + 'pending_turn': session_data[9] if session_data[9] else None, } return jsonify(result), 200 diff --git a/server/utils/cloud/infrastructure_confirmation.py b/server/utils/cloud/infrastructure_confirmation.py index 2c9c708d3..c283e2c0e 100644 --- a/server/utils/cloud/infrastructure_confirmation.py +++ b/server/utils/cloud/infrastructure_confirmation.py @@ -1,12 +1,27 @@ from __future__ import annotations -"""Shared helper for interactive user confirmation before running destructive cloud actions.""" +"""Shared helper for interactive user confirmation before running destructive cloud actions. +State model: + +* ``chat_sessions.messages`` — append-only history. Written only by + ``handle_immediate_save`` (user turn) and ``_append_new_turn_ui_messages`` + (assistant/tool turn). This module never touches it. +* ``chat_sessions.pending_turn`` — ephemeral live state for an in-flight HITL + confirmation. Set here when we prompt, cleared here when the user responds + (or on timeout). The frontend reads it on session load and renders a + synthetic awaiting-confirmation tool card at the tail. + +This separation is what prevents the duplicate/stale-tail-card class of bugs: +there is exactly one writer per slot, no reads-then-writes, no merging. +""" + +import json import logging import time import uuid -from typing import Dict, Any, List, Optional -# Lazy imports to avoid circular dependency with cloud_tools.py +from datetime import datetime +from typing import Dict, Any, Optional logger = logging.getLogger(__name__) @@ -18,7 +33,6 @@ def _send_ws(payload: dict, tool_name: str): try: logger.debug(f"WEBSOCKET: Sending confirmation via WebSocket: {payload.get('data', {}).get('confirmation_id')} for tool {tool_name}") - # Lazy import to avoid circular dependency from chat.backend.agent.tools.cloud_tools import send_websocket_message send_websocket_message( payload, @@ -29,6 +43,57 @@ def _send_ws(payload: dict, tool_name: str): logger.error(f"Failed to send WebSocket confirmation: {e}") +def _set_pending_turn(session_id: str, user_id: str, payload: Dict[str, Any]) -> None: + """Write the live HITL confirmation state to ``chat_sessions.pending_turn``. + + Fails open: if the DB write errors out we still fall through to the + WebSocket prompt, because a live tab will render the confirmation via the + WS message regardless. Only reload-into-pending requires the DB write. + """ + try: + from utils.db.connection_pool import db_pool + from utils.auth.stateless_auth import set_rls_context + with db_pool.get_user_connection() as conn: + cursor = conn.cursor() + if not set_rls_context(cursor, conn, user_id, log_prefix="[PendingTurn:Set]"): + return + cursor.execute( + """ + UPDATE chat_sessions + SET pending_turn = %s, updated_at = %s + WHERE id = %s AND user_id = %s + """, + (json.dumps(payload), datetime.now(), session_id, user_id), + ) + conn.commit() + except Exception as e: + logger.warning(f"Failed to set pending_turn for session {session_id}: {e}") + + +def _clear_pending_turn(session_id: Optional[str], user_id: Optional[str]) -> None: + """Clear ``chat_sessions.pending_turn``. Idempotent; safe to call twice.""" + if not session_id or not user_id: + return + try: + from utils.db.connection_pool import db_pool + from utils.auth.stateless_auth import set_rls_context + with db_pool.get_user_connection() as conn: + cursor = conn.cursor() + if not set_rls_context(cursor, conn, user_id, log_prefix="[PendingTurn:Clear]"): + return + cursor.execute( + """ + UPDATE chat_sessions + SET pending_turn = NULL, updated_at = %s + WHERE id = %s AND user_id = %s + """, + (datetime.now(), session_id, user_id), + ) + conn.commit() + except Exception as e: + logger.warning(f"Failed to clear pending_turn for session {session_id}: {e}") + + async def handle_websocket_confirmation_response(data: dict): """Handle incoming WebSocket confirmation responses from the frontend.""" try: @@ -36,53 +101,52 @@ async def handle_websocket_confirmation_response(data: dict): decision = data.get('decision') user_id = data.get('user_id') session_id = data.get('session_id') - + if not confirmation_id or not decision or not user_id: logger.error(f"Invalid confirmation response: {data}") return - + logger.debug(f"WEBSOCKET: Received confirmation response: user={user_id}, session={session_id}, confirmation_id={confirmation_id}, decision={decision}") - + # CRITICAL: Update the workflow's WebSocket context with the current active connection # This handles the case where the user reconnected and the workflow is still using the old connection if session_id: - # Lazy import to avoid circular dependency from chat.backend.agent.tools.cloud_tools import update_workflow_websocket_context - update_workflow_websocket_context(user_id, session_id) # Logging is handled by the function - - # Find the pending confirmation + update_workflow_websocket_context(user_id, session_id) + + # Resolve the waiter. if confirmation_id in _pending_confirmations: confirmation_data = _pending_confirmations[confirmation_id] confirmation_data['result'] = decision - # Carry optional per-change edited patterns for Yes-Always flows. edited = data.get('edited_patterns') if isinstance(edited, dict): confirmation_data['edited_patterns'] = edited logger.debug(f"WEBSOCKET: Confirmation {confirmation_id} resolved with decision: {decision}") else: logger.warning(f"WEBSOCKET: No pending confirmation found for ID: {confirmation_id}") - + + # Clear the durable live-state slot immediately on user response so a + # concurrent reload in another tab sees the resolved state rather + # than the stale prompt. The waiter's finally block also clears it, + # which makes this a safe no-op if it ran first. + _clear_pending_turn(session_id, user_id) + except Exception as e: logger.error(f"Error handling WebSocket confirmation response: {e}") def cancel_pending_confirmations_for_session(session_id: str) -> int: - """Cancel all pending confirmations for a given session. - - Called when user cancels from chat input to unblock waiting confirmation threads. - Returns the number of confirmations cancelled. - """ + """Cancel all pending confirmations for a given session.""" if not session_id: return 0 - + cancelled_count = 0 for confirmation_id, confirmation_data in _pending_confirmations.items(): - # Only cancel if not already resolved if confirmation_data.get('result') is None: confirmation_data['result'] = 'cancel' cancelled_count += 1 logger.info(f"Cancelled pending confirmation {confirmation_id} for session {session_id}") - + return cancelled_count @@ -93,15 +157,11 @@ def wait_for_user_confirmation( session_id: str = None, workflow_instance = None, ) -> bool: + """Legacy HITL helper (returns bool). Used by non-gate flows that do not + carry Yes-Always semantics. Does not write ``pending_turn`` because these + flows predate the rehydration model; they keep their existing in-memory + UI flow untouched. """ - Wait for user confirmation via WebSocket. - Uses simple polling to check for responses. - - For background chats (is_background=True), confirmations are denied - since there is no interactive user to approve destructive operations. - """ - # Background chats have no user interaction channel, so destructive operations - # must be denied rather than auto-approved. try: from chat.backend.agent.tools.cloud_tools import get_state_context state = get_state_context() @@ -110,8 +170,7 @@ def wait_for_user_confirmation( return False except Exception as e: logger.debug(f"Could not check background state: {e}") - - # Generate a unique confirmation ID for this specific request + timestamp_ms = int(time.time() * 1000) unique_id = str(uuid.uuid4())[:8] confirmation_id = f"{timestamp_ms}:{unique_id}" @@ -130,39 +189,13 @@ def wait_for_user_confirmation( ], }, } - - # Add session and user information at the top level for filtering if session_id: payload["session_id"] = session_id if user_id: payload["user_id"] = user_id - # Save UI messages to database before waiting for confirmation - # Get workflow instance from context if not provided - if not workflow_instance: - # Lazy import to avoid circular dependency - from chat.backend.agent.tools.cloud_tools import get_workflow_context - workflow_instance = get_workflow_context() - - if session_id and user_id and workflow_instance: - logger.debug(f"Consolidating and saving UI messages before confirmation for session {session_id}") - try: - # First consolidate message chunks (same as workflow pattern) - workflow_instance._consolidate_message_chunks() - - # Then save UI messages with confirmation - _save_ui_messages_with_confirmation_via_workflow(workflow_instance, session_id, user_id, tool_name, message, confirmation_id) - - except Exception as e: - logger.error(f"Error consolidating or saving UI messages before confirmation: {e}") - # Continue with confirmation even if saving fails - elif session_id and user_id: - logger.warning(f"No workflow instance available for UI message saving - skipping confirmation UI update") - - # Send confirmation prompt via WebSocket _send_ws(payload, tool_name) - # Store the confirmation request (without asyncio.Event for simplicity) _pending_confirmations[confirmation_id] = { 'result': None, 'user_id': user_id, @@ -172,29 +205,24 @@ def wait_for_user_confirmation( logger.debug(f"WEBSOCKET: Waiting for user confirmation via WebSocket with ID: {confirmation_id}") try: - # Simple polling approach - check every 2 seconds (can be reduced for higher responsiveness) elapsed = 0.0 poll_interval = 1.0 - - while elapsed < 300: # 5 minutes (Subjective value can be changed) + while elapsed < 300: confirmation_data = _pending_confirmations.get(confirmation_id) if confirmation_data and confirmation_data.get('result'): decision = confirmation_data['result'] logger.debug(f"WEBSOCKET: Received decision for {confirmation_id}: {decision}") break - time.sleep(poll_interval) elapsed += poll_interval - else: # Timeout occurred + else: decision = None logger.warning(f"WEBSOCKET: Timeout waiting for confirmation {confirmation_id}") - except Exception as e: logger.error(f"Error waiting for confirmation: {e}") decision = None - finally: # Clean up - if confirmation_id in _pending_confirmations: - del _pending_confirmations[confirmation_id] + finally: + _pending_confirmations.pop(confirmation_id, None) logger.debug(f"WEBSOCKET: Decision for {confirmation_id}: {decision}") return decision == "execute" @@ -211,17 +239,13 @@ def wait_for_user_confirmation_ex( workflow_instance=None, timeout_seconds: int = 300, ) -> Dict[str, Any]: - """Extended HITL helper that returns the full response payload. + """Extended HITL helper used exclusively by the command gate. - Unlike :func:`wait_for_user_confirmation` (which returns a bool), this - preserves the user's exact decision (e.g. ``execute`` vs ``execute_always`` - vs ``cancel``) plus any ``edited_patterns`` map the frontend sends back. - - Returns ``{"decision": str | None, "edited_patterns": dict}``. ``decision`` - is ``None`` only on timeout. ``edited_patterns`` maps stringified change - indices to user-edited regex patterns. + Returns ``{"decision": str | None, "edited_patterns": dict}``. Persists the + live prompt to ``chat_sessions.pending_turn`` so the UI can rehydrate the + awaiting-confirmation card after a page reload, and always clears it on + return (user response or timeout). """ - # Background chats: no interactive user. try: from chat.backend.agent.tools.cloud_tools import get_state_context state = get_state_context() @@ -253,19 +277,18 @@ def wait_for_user_confirmation_ex( if user_id: payload["user_id"] = user_id - if not workflow_instance: - from chat.backend.agent.tools.cloud_tools import get_workflow_context - workflow_instance = get_workflow_context() - - if session_id and user_id and workflow_instance: - try: - workflow_instance._consolidate_message_chunks() - _save_ui_messages_with_confirmation_via_workflow( - workflow_instance, session_id, user_id, tool_name, message, confirmation_id, - extra=extra, - ) - except Exception as e: - logger.error(f"Error consolidating/saving UI messages before confirmation: {e}") + # Durable live-state slot for reload-into-pending rehydration. Mirrors + # the WS payload's ``data`` field so the frontend can build the synthetic + # tool card with identical semantics to the live WS path. + if session_id: + _set_pending_turn(session_id, user_id, { + "confirmation_id": confirmation_id, + "tool_name": tool_name, + "message": message, + "options": options, + **(extra or {}), + "created_at": datetime.now().isoformat(), + }) _send_ws(payload, tool_name) @@ -294,161 +317,6 @@ def wait_for_user_confirmation_ex( logger.error(f"Error waiting for confirmation_ex: {e}") finally: _pending_confirmations.pop(confirmation_id, None) + _clear_pending_turn(session_id, user_id) return {"decision": decision, "edited_patterns": edited} - - -def _save_ui_messages_with_confirmation_via_workflow( - workflow_instance, - session_id: str, - user_id: str, - tool_name: str, - message: str, - confirmation_id: str, - extra: Optional[Dict[str, Any]] = None, -) -> bool: - """Save UI messages to database including any ongoing tool calls and confirmation marker. - - When *extra* is provided (gate-driven confirmations), the awaiting tool call - is enriched with ``command`` / ``block_layer`` / ``yes_always_effect`` so the - Yes-Always popover can rehydrate correctly after a page reload. We also - merge-preserve any previously saved user messages that the rebuilt UI list - happens to omit (e.g. when the LLM state snapshot was taken mid-turn), so - navigating away does not silently drop the user's prompt. - """ - try: - # Get LLM messages from workflow state - if workflow_instance._last_state: - llm_messages = ( - workflow_instance._last_state.get('messages', []) - if hasattr(workflow_instance._last_state, 'get') - else getattr(workflow_instance._last_state, 'messages', []) - ) - logger.debug(f"Found {len(llm_messages)} LLM messages in workflow state") - else: - logger.warning("No last state found for workflow instance") - llm_messages = [] - - # Get the current tool capture to include ongoing tool calls - # Lazy import to avoid circular dependency - from chat.backend.agent.tools.cloud_tools import get_tool_capture - tool_capture = get_tool_capture() - if not tool_capture or not hasattr(tool_capture, 'current_tool_calls'): - logger.warning("No tool capture found for confirmation UI update") - return False - - # Convert LLM messages to UI format using workflow's method - ui_messages = workflow_instance._convert_to_ui_messages(llm_messages, tool_capture) # Logging is handled by the function - - # Find the tool call for the requesting tool and update its status - target_tool_call = None - for _, tool_info in tool_capture.current_tool_calls.items(): - current_tool_name = tool_info.get('tool_name') - if current_tool_name == tool_name: - target_tool_call = tool_info.get('call_id') # Use call_id (call_xxx), not run_id (run-xxx) - break - - if target_tool_call: - logger.debug(f"Found target tool call: {target_tool_call} for tool {tool_name}") - - # Update the specific tool call to have awaiting_confirmation status - tool_call_updated = False - for ui_msg in ui_messages: - if ui_msg.get('sender') == 'bot' and ui_msg.get('toolCalls'): - for tool_call in ui_msg.get('toolCalls', []): - # Match by tool_call_id (call_xxx format) - if tool_call.get('id') == target_tool_call: - tool_call['status'] = 'awaiting_confirmation' - tool_call['confirmation_id'] = confirmation_id - tool_call['confirmation_message'] = message - if extra: - # Persist only the fields the confirmation UI - # needs on rehydrate; avoid leaking internal - # fields into the saved record. - for key in ("command", "block_layer", "yes_always_effect"): - val = extra.get(key) - if val is not None: - tool_call[key] = val - logger.debug(f"Updated tool call {target_tool_call} to awaiting_confirmation status") - tool_call_updated = True - break - if tool_call_updated: - break - - if not tool_call_updated: - logger.warning(f"Could not find tool call {target_tool_call} in UI messages to update confirmation status") - else: - logger.warning(f"No tool call found for {tool_name} in current tool calls") - logger.debug(f"Available tool calls: {list(tool_capture.current_tool_calls.keys())}") - - # Preserve any previously-saved user messages the rebuild dropped. The - # LLM-state snapshot may not include the current turn's HumanMessage - # if the gate fires before LangGraph has persisted it; without this - # merge the unconditional overwrite would wipe the user's prompt. - ui_messages = _merge_missing_user_messages(session_id, user_id, ui_messages) - - # Use workflow's _save_ui_messages method to save the updated messages - logger.debug(f"Saving {len(ui_messages)} UI messages") - return workflow_instance._save_ui_messages(session_id, user_id, ui_messages) - - except Exception as e: - logger.error(f"Error saving UI messages with confirmation: {e}") - return False - - -def _merge_missing_user_messages( - session_id: str, user_id: str, rebuilt: List[Dict[str, Any]], -) -> List[Dict[str, Any]]: - """Re-insert persisted user messages that ``rebuilt`` is missing. - - The LLM-state snapshot may not include the current turn's HumanMessage, - in which case ``rebuilt`` starts with the assistant/tool response and - drops the user's prompt. We splice those missing user messages back in - at their original chronological position from the persisted row so the - UI renders the prompt above its response. Fails open on any DB error. - """ - try: - from utils.db.connection_pool import db_pool - with db_pool.get_user_connection() as conn: - with conn.cursor() as cur: - cur.execute( - "SELECT messages FROM chat_sessions WHERE id = %s AND user_id = %s", - (session_id, user_id), - ) - row = cur.fetchone() - existing = row[0] if row and row[0] else [] - if not isinstance(existing, list): - return rebuilt - except Exception as e: - logger.warning(f"merge_missing_user_messages: could not read existing messages: {e}") - return rebuilt - - rebuilt_user_texts = { - m.get('text') for m in rebuilt if m.get('sender') == 'user' and m.get('text') - } - missing_any = any( - m.get('sender') == 'user' and m.get('text') and m.get('text') not in rebuilt_user_texts - for m in existing - ) - if not missing_any: - return rebuilt - - # Walk ``existing`` in order. For each user message not present in the - # rebuilt set, insert a copy into ``merged`` at the position we reach in - # ``existing``; for every other slot we advance through ``rebuilt``. - merged: List[Dict[str, Any]] = [] - rebuilt_idx = 0 - for em in existing: - if (em.get('sender') == 'user' and em.get('text') - and em.get('text') not in rebuilt_user_texts): - merged.append(em) - continue - if rebuilt_idx < len(rebuilt): - merged.append(rebuilt[rebuilt_idx]) - rebuilt_idx += 1 - # Anything left in ``rebuilt`` is newer than ``existing`` — append as-is. - merged.extend(rebuilt[rebuilt_idx:]) - - for i, m in enumerate(merged, start=1): - m['message_number'] = i - return merged \ No newline at end of file diff --git a/server/utils/db/db_utils.py b/server/utils/db/db_utils.py index 3a07dff4e..320768c20 100644 --- a/server/utils/db/db_utils.py +++ b/server/utils/db/db_utils.py @@ -1461,6 +1461,23 @@ def initialize_tables(): logging.warning(f"Error adding incident_id index: {e}") conn.rollback() + # Migration: Add pending_turn column (live HITL state, separate from + # the append-only messages history). Cleared by the command gate + # once the user resolves the confirmation. Rehydrated as a + # synthetic tail card on session load. + try: + cursor.execute(""" + ALTER TABLE chat_sessions + ADD COLUMN IF NOT EXISTS pending_turn JSONB; + """) + logging.info( + "Added pending_turn column to chat_sessions table (if not exists)." + ) + conn.commit() + except Exception as e: + logging.warning(f"Error adding pending_turn column: {e}") + conn.rollback() + # Migration: Add surcharge fields to llm_usage_tracking table if they don't exist try: cursor.execute(""" From 0e2772063093ff17fb2bdffc14ee73a9e4e590ac Mon Sep 17 00:00:00 2001 From: Damian Loch Date: Tue, 28 Apr 2026 10:37:31 -0400 Subject: [PATCH 06/13] tainting sessinos originally blocked by nemo --- server/chat/backend/agent/workflow.py | 15 ++- server/utils/auth/command_gate.py | 138 +++++++++++++++++--------- server/utils/db/db_utils.py | 17 ++++ 3 files changed, 123 insertions(+), 47 deletions(-) diff --git a/server/chat/backend/agent/workflow.py b/server/chat/backend/agent/workflow.py index ea23b8b12..989fb5bd9 100644 --- a/server/chat/backend/agent/workflow.py +++ b/server/chat/backend/agent/workflow.py @@ -979,8 +979,19 @@ async def stream(self, input_state: State): reason=rail_result.reason, latency_ms=rail_result.latency_ms, ) - yield ("token", "Your message was blocked by our safety system. Please rephrase your request.") - return + # Background chats have no interactive user: hard block stays. + # Foreground chats: taint the session so every subsequent tool + # call goes through the command gate's Yes/No prompt. The user + # gets no special UI for this block -- it surfaces only when + # the agent next tries to run a tool. + if getattr(input_state, "is_background", False): + yield ("token", "Your message was blocked by our safety system. Please rephrase your request.") + return + from utils.auth.command_gate import mark_session_tainted + mark_session_tainted( + getattr(input_state, "session_id", None), + getattr(input_state, "user_id", None), + ) # Log initial state logger.info(f"Starting workflow with session_id={input_state.session_id}, user_id={input_state.user_id}") diff --git a/server/utils/auth/command_gate.py b/server/utils/auth/command_gate.py index c76b840ac..4d3ff09e3 100644 --- a/server/utils/auth/command_gate.py +++ b/server/utils/auth/command_gate.py @@ -135,10 +135,54 @@ def gate_command( _gate_inflight_command.reset(token) +def is_session_tainted(session_id: Optional[str]) -> bool: + """Return True iff ``session_id`` has been marked tainted (NeMo input-rail + hit on the opening user message of this foreground chat). + + Tainted sessions force every command through user confirmation, even when + all guardrail layers pass. Fails closed on DB errors (treats as untainted) + since the gate's other layers already provide defense-in-depth. + """ + if not session_id: + return False + try: + from utils.db.connection_pool import db_pool + with db_pool.get_user_connection() as conn: + cursor = conn.cursor() + cursor.execute( + "SELECT security_tainted FROM chat_sessions WHERE id = %s", + (session_id,), + ) + row = cursor.fetchone() + return bool(row and row[0]) + except Exception as e: + logger.warning(f"[CommandGate] taint lookup failed for {session_id}: {e}") + return False + + +def mark_session_tainted(session_id: Optional[str], user_id: Optional[str]) -> None: + """Flip ``security_tainted`` to true for this session. Idempotent.""" + if not session_id or not user_id: + return + try: + from utils.db.connection_pool import db_pool + from utils.auth.stateless_auth import set_rls_context + with db_pool.get_user_connection() as conn: + cursor = conn.cursor() + if not set_rls_context(cursor, conn, user_id, log_prefix="[CommandGate:Taint]"): + return + cursor.execute( + "UPDATE chat_sessions SET security_tainted = true WHERE id = %s AND user_id = %s", + (session_id, user_id), + ) + conn.commit() + except Exception as e: + logger.warning(f"[CommandGate] failed to mark session {session_id} tainted: {e}") + + def _gate_impl(*, user_id: str, tool_name: str, command: str, cmd_hash: str) -> GateDecision: from utils.auth.command_policy import ( evaluate_compound_command, CommandVerdict, plan_yes_always, - apply_yes_always, validate_pattern, PolicyChange, ) from utils.auth.stateless_auth import get_org_id_for_user from utils.security.command_safety import evaluate_command as safety_evaluate @@ -147,71 +191,75 @@ def _gate_impl(*, user_id: str, tool_name: str, command: str, cmd_hash: str) -> foreground = _is_foreground() session_id = _session_id() - # --- Layer 1 & 3: signature_match + LLM judge ---------------------------- - # command_safety.evaluate_command bundles both: signature first (fast, - # cached), LLM judge second (network call). We read decision.layer to know - # which fired. + # Evaluate all layers unconditionally so we can report the combined + # block state to the user. The previous short-circuit (return on first + # safety block) prevented Always from showing when the policy layer + # would have also fired. safety_decision = safety_evaluate( command, tool=tool_name, user_id=user_id, session_id=session_id, ) - if safety_decision.blocked: - layer = safety_decision.layer or "llm_judge" - code = "SIGNATURE_MATCHED" if layer == "signature_match" else "SAFETY_BLOCKED" - reason = f"Command blocked by safety guardrail: {safety_decision.reason}" - if not foreground: - return _block(code, reason) - # Foreground: ask user. No Yes-Always for this layer because there is - # no policy slot to flip. - return _prompt_user( - user_id=user_id, - session_id=session_id, - tool_name=tool_name, - command=command, - block_code=code, - block_reason=reason, - block_layer=layer, - allow_yes_always=False, - yes_always_changes=[], - org_id=None, - cmd_hash=cmd_hash, - ) - - # --- Layer 2: org allow/deny -------------------------------------------- policy_verdict: CommandVerdict = evaluate_compound_command(org_id, command) - if policy_verdict.allowed: + + safety_blocked = safety_decision.blocked + policy_blocked = not policy_verdict.allowed + tainted = foreground and is_session_tainted(session_id) + + if not (safety_blocked or policy_blocked or tainted): # Tell terminal_run._check_guardrails it may skip re-running # signature+judge for this command on the same invocation. _guardrails_approved_command.set(cmd_hash) return _ALLOWED - reason = (policy_verdict.rule_description or "Matched organization policy")[:200] - block_reason = f"Command blocked by organization policy: {reason}" - layer = ( - "policy_both" if policy_verdict.deny_rule_id and policy_verdict.allowlist_exhausted - else "policy_deny" if policy_verdict.deny_rule_id - else "policy_allow_exhausted" + # Compose the block code/reason/layer from whichever layers fired. + safety_layer = safety_decision.layer if safety_blocked else None + safety_code = ( + "SIGNATURE_MATCHED" if safety_layer == "signature_match" + else "SAFETY_BLOCKED" if safety_blocked else None + ) + policy_layer = ( + "policy_both" if policy_blocked and policy_verdict.deny_rule_id + and policy_verdict.allowlist_exhausted + else "policy_deny" if policy_blocked and policy_verdict.deny_rule_id + else "policy_allow_exhausted" if policy_blocked + else None ) - if not foreground or not org_id: - return _block("POLICY_DENIED", block_reason) + layers = [l for l in (safety_layer, policy_layer) if l] + if not layers and tainted: + layers = ["session_taint"] + block_layer = "+".join(layers) + + code = safety_code or ("POLICY_DENIED" if policy_blocked else "SESSION_TAINTED") + reasons = [] + if safety_blocked: + reasons.append(f"safety guardrail: {safety_decision.reason}") + if policy_blocked: + reasons.append( + "organization policy: " + + (policy_verdict.rule_description or "matched organization policy")[:200] + ) + if not reasons and tainted: + reasons.append("session flagged by input safety check; approval required") + block_reason = "Command blocked by " + "; ".join(reasons) + + if not foreground: + return _block(code, block_reason) - changes = plan_yes_always(policy_verdict, command) + # Always is offered iff the policy layer fired with a real mutation to + # propose. Safety-only or taint-only blocks show Yes/No. + changes = plan_yes_always(policy_verdict, command) if policy_blocked else [] decision = _prompt_user( user_id=user_id, session_id=session_id, tool_name=tool_name, command=command, - block_code="POLICY_DENIED", + block_code=code, block_reason=block_reason, - block_layer=layer, + block_layer=block_layer or "unknown", allow_yes_always=bool(changes), yes_always_changes=changes, - org_id=org_id, + org_id=org_id if policy_blocked else None, cmd_hash=cmd_hash, ) - # On Yes / Yes-Always, also clear the guardrails-approved marker so the - # downstream signature+judge re-check (if any) still runs normally for - # commands that bypassed the policy via user consent — the safety layers - # already passed above. if decision.allowed: _guardrails_approved_command.set(cmd_hash) return decision diff --git a/server/utils/db/db_utils.py b/server/utils/db/db_utils.py index 320768c20..34fe5c682 100644 --- a/server/utils/db/db_utils.py +++ b/server/utils/db/db_utils.py @@ -1478,6 +1478,23 @@ def initialize_tables(): logging.warning(f"Error adding pending_turn column: {e}") conn.rollback() + # Migration: Add security_tainted column. When NeMo's input rail + # blocks the opening user message in a foreground chat, we mark + # the session tainted instead of hard-failing; every subsequent + # tool call then requires user approval via the command gate. + try: + cursor.execute(""" + ALTER TABLE chat_sessions + ADD COLUMN IF NOT EXISTS security_tainted BOOLEAN NOT NULL DEFAULT false; + """) + logging.info( + "Added security_tainted column to chat_sessions table (if not exists)." + ) + conn.commit() + except Exception as e: + logging.warning(f"Error adding security_tainted column: {e}") + conn.rollback() + # Migration: Add surcharge fields to llm_usage_tracking table if they don't exist try: cursor.execute(""" From 0910cc238668058a6888329bdf288c4bc4db64c9 Mon Sep 17 00:00:00 2001 From: Damian Loch Date: Tue, 28 Apr 2026 10:58:20 -0400 Subject: [PATCH 07/13] address coderabbit review comments - ToolExecutionWidget: only transition card after sendRaw succeeds - is_session_tainted: set RLS context + scope by user_id (mirrors mark_session_tainted) - derive_pattern_from_command: keep up to 3 tokens so "aws ec2 terminate-instances" matches docstring - command_gate: require org admin (Casbin admin/access) before offering/applying Yes-Always - infrastructure_confirmation: only clear pending_turn when confirmation id matched - infrastructure_confirmation: register waiter before WS send to avoid fast-response race Made-with: Cursor --- .../tool-calls/ToolExecutionWidget.tsx | 3 +- server/utils/auth/command_gate.py | 46 +++++++++++++++---- server/utils/auth/command_policy.py | 12 +++-- .../cloud/infrastructure_confirmation.py | 23 ++++++---- 4 files changed, 60 insertions(+), 24 deletions(-) diff --git a/client/src/components/tool-calls/ToolExecutionWidget.tsx b/client/src/components/tool-calls/ToolExecutionWidget.tsx index af3ce4efc..29940207e 100644 --- a/client/src/components/tool-calls/ToolExecutionWidget.tsx +++ b/client/src/components/tool-calls/ToolExecutionWidget.tsx @@ -604,7 +604,8 @@ const ConfirmationPanel = ({ tool, command, userId, sessionId, sendRaw, onToolUp session_id: sessionId, } if (decision === 'execute_always') payload.edited_patterns = edited - sendRaw(JSON.stringify(payload)) + const sent = sendRaw(JSON.stringify(payload)) + if (!sent) return if (decision === 'cancel') { onToolUpdate?.({ status: 'completed', output: 'Operation cancelled by user' }) } else { diff --git a/server/utils/auth/command_gate.py b/server/utils/auth/command_gate.py index 4d3ff09e3..d2bec4ddb 100644 --- a/server/utils/auth/command_gate.py +++ b/server/utils/auth/command_gate.py @@ -135,23 +135,27 @@ def gate_command( _gate_inflight_command.reset(token) -def is_session_tainted(session_id: Optional[str]) -> bool: +def is_session_tainted(session_id: Optional[str], user_id: Optional[str]) -> bool: """Return True iff ``session_id`` has been marked tainted (NeMo input-rail hit on the opening user message of this foreground chat). Tainted sessions force every command through user confirmation, even when - all guardrail layers pass. Fails closed on DB errors (treats as untainted) + all guardrail layers pass. Reads under the caller's RLS context (matches + ``mark_session_tainted``); fails closed (treats as untainted) on DB error since the gate's other layers already provide defense-in-depth. """ - if not session_id: + if not session_id or not user_id: return False try: from utils.db.connection_pool import db_pool + from utils.auth.stateless_auth import set_rls_context with db_pool.get_user_connection() as conn: cursor = conn.cursor() + if not set_rls_context(cursor, conn, user_id, log_prefix="[CommandGate:TaintRead]"): + return False cursor.execute( - "SELECT security_tainted FROM chat_sessions WHERE id = %s", - (session_id,), + "SELECT security_tainted FROM chat_sessions WHERE id = %s AND user_id = %s", + (session_id, user_id), ) row = cursor.fetchone() return bool(row and row[0]) @@ -202,7 +206,7 @@ def _gate_impl(*, user_id: str, tool_name: str, command: str, cmd_hash: str) -> safety_blocked = safety_decision.blocked policy_blocked = not policy_verdict.allowed - tainted = foreground and is_session_tainted(session_id) + tainted = foreground and is_session_tainted(session_id, user_id) if not (safety_blocked or policy_blocked or tainted): # Tell terminal_run._check_guardrails it may skip re-running @@ -265,6 +269,24 @@ def _gate_impl(*, user_id: str, tool_name: str, command: str, cmd_hash: str) -> return decision +def _user_is_org_admin(user_id: str, org_id: Optional[str]) -> bool: + """Return True iff *user_id* has admin access in *org_id*. + + Yes-Always mutates ``org_command_policies``, which the HTTP routes + gate behind ``require_permission("admin", "access")``. We mirror that + check here so a non-admin cannot rewrite org policy by clicking a + chat button. Fails closed on any error. + """ + if not user_id or not org_id: + return False + try: + from utils.auth.enforcer import get_enforcer + return bool(get_enforcer().enforce(user_id, org_id, "admin", "access")) + except Exception as e: + logger.warning(f"[CommandGate] admin check failed for {user_id}/{org_id}: {e}") + return False + + def _prompt_user( *, user_id: str, @@ -292,7 +314,15 @@ def _prompt_user( "block_reason": block_reason, "command": command, } - if allow_yes_always and yes_always_changes: + # Yes-Always rewrites org_command_policies, which the HTTP routes gate + # behind admin. Non-admins get Yes/No only; their approval is + # session-scoped and no policy row changes. + offer_always = ( + allow_yes_always + and bool(yes_always_changes) + and _user_is_org_admin(user_id, org_id) + ) + if offer_always: options.append({"text": "Yes, Always", "value": "execute_always"}) extra["yes_always_effect"] = { "summary": "This will modify your organization's command policy:", @@ -324,7 +354,7 @@ def _prompt_user( if decision == "execute": return _ALLOWED - if decision == "execute_always" and allow_yes_always and org_id: + if decision == "execute_always" and offer_always and org_id: edited = result.get("edited_patterns") or {} applied = [] for idx, ch in enumerate(yes_always_changes): diff --git a/server/utils/auth/command_policy.py b/server/utils/auth/command_policy.py index bed4220c8..3ee00dd0a 100644 --- a/server/utils/auth/command_policy.py +++ b/server/utils/auth/command_policy.py @@ -344,8 +344,8 @@ def derive_pattern_from_command(command: str) -> str: because the user can loosen it in the edit box before confirming. Examples: - "sudo kubectl delete pod foo" -> ^kubectl delete\\b - "KUBECONFIG=/x kubectl get pods" -> ^kubectl get\\b + "sudo kubectl delete pod foo" -> ^kubectl delete pod\\b + "KUBECONFIG=/x kubectl get pods" -> ^kubectl get pods\\b "aws ec2 terminate-instances --id" -> ^aws ec2 terminate-instances\\b "for i in {1..10}; do echo $i; done" -> ^for i in \\{1\\.\\.10\\}; do echo \\$i; done$ """ @@ -363,9 +363,11 @@ def derive_pattern_from_command(command: str) -> str: if tok.startswith("-"): break parts.append(tok) - # Stop after CLI + first subcommand to keep the pattern reasonably - # narrow. Callers (UI) can edit before applying. - if len(parts) >= 2: + # Stop after CLI + up to two subcommand tokens so multi-word + # subcommands like "aws ec2 terminate-instances" match the docstring + # example rather than collapsing to "^aws\s+ec2\b". Users can still + # tighten or loosen in the editable UI field. + if len(parts) >= 3: break return r"^" + r"\s+".join(re.escape(p) for p in parts) + r"\b" diff --git a/server/utils/cloud/infrastructure_confirmation.py b/server/utils/cloud/infrastructure_confirmation.py index c283e2c0e..1ac1e4d35 100644 --- a/server/utils/cloud/infrastructure_confirmation.py +++ b/server/utils/cloud/infrastructure_confirmation.py @@ -122,15 +122,17 @@ async def handle_websocket_confirmation_response(data: dict): if isinstance(edited, dict): confirmation_data['edited_patterns'] = edited logger.debug(f"WEBSOCKET: Confirmation {confirmation_id} resolved with decision: {decision}") + # Clear the durable live-state slot immediately on user response + # so a concurrent reload in another tab sees the resolved state + # rather than the stale prompt. The waiter's finally block also + # clears it, which makes this a safe no-op if it ran first. + # Only clear when the id actually matched a pending waiter -- + # otherwise a stale/late response for a superseded prompt would + # erase the currently active slot. + _clear_pending_turn(session_id, user_id) else: logger.warning(f"WEBSOCKET: No pending confirmation found for ID: {confirmation_id}") - # Clear the durable live-state slot immediately on user response so a - # concurrent reload in another tab sees the resolved state rather - # than the stale prompt. The waiter's finally block also clears it, - # which makes this a safe no-op if it ran first. - _clear_pending_turn(session_id, user_id) - except Exception as e: logger.error(f"Error handling WebSocket confirmation response: {e}") @@ -194,13 +196,12 @@ def wait_for_user_confirmation( if user_id: payload["user_id"] = user_id - _send_ws(payload, tool_name) - _pending_confirmations[confirmation_id] = { 'result': None, 'user_id': user_id, 'timestamp': time.time() } + _send_ws(payload, tool_name) logger.debug(f"WEBSOCKET: Waiting for user confirmation via WebSocket with ID: {confirmation_id}") @@ -290,13 +291,15 @@ def wait_for_user_confirmation_ex( "created_at": datetime.now().isoformat(), }) - _send_ws(payload, tool_name) - _pending_confirmations[confirmation_id] = { 'result': None, 'user_id': user_id, 'timestamp': time.time(), } + # Register the waiter before sending so a fast client response can never + # arrive before _pending_confirmations has the entry + # (_send_ws dispatches on a daemon thread, see cloud_tools). + _send_ws(payload, tool_name) logger.debug(f"WEBSOCKET: Waiting for confirmation_ex {confirmation_id}") decision: Optional[str] = None From 1757f39d7d4896c6183db5e13092e5129032ef4d Mon Sep 17 00:00:00 2001 From: Damian Loch Date: Thu, 30 Apr 2026 15:17:02 -0400 Subject: [PATCH 08/13] consolidate legacy confirmation prompts into unified command gate The pre-gate codebase had two parallel confirmation paths: the new command gate (shell commands, policy-aware, with Yes/No/Yes-Always) and a legacy wait_for_user_confirmation helper used by structured actions (Terraform apply/destroy, Bitbucket merges/deletes, Notion column deletes, Spinnaker pipeline triggers, destructive MCP tools, multi- account cloud_exec). Two code paths, two UIs, two pieces of plumbing for the same "are you sure?" question. This commit collapses them into one: - Add gate_action() to command_gate.py. Same _prompt_user surface as gate_command() but Yes/No only (no regex means no Yes-Always), denies in background (matches the previous legacy-helper behavior). - Delete the two cloud_exec legacy prompts and the now-orphaned summarize_cloud_command helper and its verb sets. Org policy is the source of truth for which cloud CLI verbs need approval; the unified gate at the top of cloud_exec already enforces that. Gate the multi-account fan-out once up front to keep the single-prompt UX. - Migrate iac_commands_tool (apply/destroy), bitbucket confirm_or_cancel, notion_export_postmortem, notion_update_database_properties (deletes only), spinnaker trigger_pipeline, and mcp_tools destructive-tool path to gate_action. - Delete wait_for_user_confirmation from infrastructure_confirmation.py now that nothing calls it. Net result: one gate, one confirmation UI, one WS protocol, one taint model. No behavior change for users beyond the visual unification (same Yes/No buttons everywhere, same cancel semantics, same background-chat deny behavior). Made-with: Cursor --- .../backend/agent/tools/bitbucket/utils.py | 31 ++-- .../backend/agent/tools/cloud_exec_tool.py | 146 +++--------------- .../agent/tools/iac/iac_commands_tool.py | 38 ++--- server/chat/backend/agent/tools/mcp_tools.py | 18 +-- .../backend/agent/tools/notion/postmortem.py | 12 +- .../backend/agent/tools/notion/structured.py | 12 +- .../backend/agent/tools/spinnaker_rca_tool.py | 13 +- server/utils/auth/command_gate.py | 57 ++++++- .../cloud/infrastructure_confirmation.py | 77 --------- 9 files changed, 112 insertions(+), 292 deletions(-) diff --git a/server/chat/backend/agent/tools/bitbucket/utils.py b/server/chat/backend/agent/tools/bitbucket/utils.py index ad04fc332..f22628314 100644 --- a/server/chat/backend/agent/tools/bitbucket/utils.py +++ b/server/chat/backend/agent/tools/bitbucket/utils.py @@ -145,13 +145,6 @@ def build_success_response(**kwargs) -> str: return json.dumps(result, default=str) -def get_session_id(): - """Get session_id from the current agent state context.""" - from chat.backend.agent.tools.cloud_tools import get_state_context - - return getattr(get_state_context(), "session_id", None) - - def build_cancelled_response() -> str: """Build the standard cancellation response for a rejected confirmation.""" return build_success_response(message="Operation cancelled by user", cancelled=True) @@ -160,19 +153,13 @@ def build_cancelled_response() -> str: def confirm_or_cancel(user_id: str, message: str, tool_name: str) -> Optional[str]: """Request human approval for a destructive action. - Retrieves the session ID automatically and prompts the user for - confirmation. Returns ``None`` if approved, or a JSON cancellation - response string if the user declines. + Returns ``None`` if approved, or a JSON cancellation response string + if the user declines. Delegates to the unified command gate so + Bitbucket confirmations share the same UI/WS/taint plumbing as the + shell-command gate. """ - from utils.cloud.infrastructure_confirmation import wait_for_user_confirmation - - session_id = get_session_id() - approved = wait_for_user_confirmation( - user_id=user_id, - message=message, - tool_name=tool_name, - session_id=session_id, - ) - if not approved: - return build_cancelled_response() - return None + from utils.auth.command_gate import gate_action + + if gate_action(user_id=user_id, tool_name=tool_name, summary=message).allowed: + return None + return build_cancelled_response() diff --git a/server/chat/backend/agent/tools/cloud_exec_tool.py b/server/chat/backend/agent/tools/cloud_exec_tool.py index c6e11b5a8..559ab38fb 100644 --- a/server/chat/backend/agent/tools/cloud_exec_tool.py +++ b/server/chat/backend/agent/tools/cloud_exec_tool.py @@ -16,7 +16,6 @@ from utils.auth.cloud_auth import generate_contextual_access_token from utils.auth.cloud_auth import generate_azure_access_token from .output_sanitizer import sanitize_command_output, filter_error_messages, truncate_json_fields -from utils.cloud.infrastructure_confirmation import wait_for_user_confirmation from .cloud_provider_utils import determine_target_provider_from_context from chat.backend.agent.prompt.prompt_builder import CLOUD_EXEC_PROVIDERS from chat.backend.agent.access import ModeAccessController @@ -45,77 +44,6 @@ def count_tokens(text: str, model: str = "gpt-4o") -> int: """Count tokens in text using LLMUsageTracker (for context management, not billing).""" return LLMUsageTracker.count_tokens(text, model) -# -------------------------------------------------------------------------------------- -# Common verb sets used across providers -# -------------------------------------------------------------------------------------- -_READ_ONLY_VERBS: set[str] = { - "list", "describe", "get", "show", "config", "version", "info", "view", "read", "status", -} -# A non-exhaustive but conservative set of verbs considered to CHANGE state -_ACTION_VERBS: set[str] = { - "create", "delete", "update", "apply", "destroy", "terminate", "start", "stop", - "restart", "attach", "detach", "enable", "disable", "put", "remove", -} - - -def summarize_cloud_command(cmd: str) -> str: - """Return a short description of what the CLI command (gcloud/az/aws) intends to do.""" - - try: - tokens = shlex.split(cmd) - except Exception: - tokens = cmd.split() - - action = None - resource_type = None - resource_name = None - zone_or_region = None - - # Tokens representing the CLI executable that we should skip when searching - cli_tokens = {"gcloud", "az", "aws", "kubectl", "gsutil", "bq", "scw", "ovh"} - - for i, tok in enumerate(tokens): - low = tok.lower() - - # Detect action verb - if low in _ACTION_VERBS.union(_READ_ONLY_VERBS): - action = low - - # Try to infer resource type – walk backwards until we find a token - # that is not a flag or CLI executable. - j = i - 1 - while j >= 0 and (tokens[j].lower() in cli_tokens or tokens[j].startswith("-")): - j -= 1 - if j >= 0: - resource_type = tokens[j] - - # Resource name usually follows the action verb (unless it's a flag) - k = i + 1 - while k < len(tokens) and tokens[k].startswith("-"): - k += 1 - if k < len(tokens): - resource_name = tokens[k] - - # Capture location / region / zone flags for common CLIs - if low.startswith("--zone="): - zone_or_region = tok.split("=", 1)[1] - elif low == "--zone" and i + 1 < len(tokens): - zone_or_region = tokens[i + 1] - elif low.startswith("--region=") or low.startswith("--location="): - zone_or_region = tok.split("=", 1)[1] - elif low in {"--region", "--location", "-r", "-l"} and i + 1 < len(tokens): - zone_or_region = tokens[i + 1] - - parts: list[str] = [] - if action and resource_type: - parts.append(f"The command will {action} {resource_type}") - if resource_name and not resource_name.startswith("--"): - parts.append(f"named '{resource_name}'") - if zone_or_region: - parts.append(f"in {zone_or_region}") - - summary_core = " ".join(parts) if parts else cmd - return f"{summary_core}.\n\n" logger = logging.getLogger(__name__) @@ -1228,25 +1156,21 @@ def _cloud_exec_aws_multi_account( "provider": "aws", }) - if not is_read_only_command(command): - from utils.cloud.infrastructure_confirmation import wait_for_user_confirmation - from .cloud_tools import get_state_context - summary_msg = summarize_cloud_command(command) - state_context = get_state_context() - context_session_id = state_context.session_id if state_context and hasattr(state_context, 'session_id') else None - if not wait_for_user_confirmation( - user_id=user_id, - message=f"[ALL {len(connections)} accounts] {summary_msg}", - tool_name="cloud_exec", - session_id=context_session_id, - ): - return json.dumps({ - "success": False, - "error": "User declined multi-account command execution", - "multi_account": True, - "command": command, - "provider": "aws", - }) + # One gate check for the fan-out (signature + org policy + LLM judge + HITL). + # Each per-account invocation is the same command text, so gating once up + # front matches the previous single-prompt UX and avoids N prompts. + from utils.auth.command_gate import gate_command + _gated = command if command.strip().startswith("aws") else f"aws {command}" + _gate = gate_command(user_id=user_id, tool_name="cloud_exec", command=_gated) + if not _gate.allowed: + return json.dumps({ + "success": False, + "error": _gate.block_reason, + "code": _gate.code, + "multi_account": True, + "command": command, + "provider": "aws", + }) def _run_on_account(conn: dict) -> dict: account_id = conn.get("account_id", "unknown") @@ -1831,42 +1755,10 @@ def cloud_exec(provider: str, command: str, user_id: Optional[str] = None, sessi }) - # If command is potentially action, ask for confirmation first (after command processing) - if not is_read_only_command(command): - summary_msg = summarize_cloud_command(command) - # Get session_id from context for state saving - from .cloud_tools import get_state_context - state_context = get_state_context() - context_session_id = state_context.session_id if state_context and hasattr(state_context, 'session_id') else None - - if not wait_for_user_confirmation( - user_id=user_id, - message=summary_msg, - tool_name="cloud_exec", - session_id=context_session_id - ): - # Capture the cancellation result in the tool capture system - tool_capture = get_tool_capture() - current_tool_call_id = get_current_tool_call_id( - tool_name="cloud_exec", - tool_kwargs={'provider': original_provider, 'command': original_command} - ) - - cancellation_result = json.dumps({ - "status": "cancelled", - "message": "cloud_exec command cancelled by user", - "chat_output": "Command cancelled.", - "user_cancelled": True, - "final_command": command, # Now includes the processed command with all prefixes and flags - }) - - # Capture the cancellation result if we have a tool capture and current tool call ID - if tool_capture and current_tool_call_id: - logger.info(f"Capturing cancellation result for tool call {current_tool_call_id}") - tool_capture.capture_tool_end(current_tool_call_id, cancellation_result, is_error=False) - - return cancellation_result - + # Destructive-command confirmation is handled by the unified command + # gate earlier in this function (signature + org policy + LLM judge + + # HITL). Org policy is the source of truth for what requires approval. + # Check if CLI tool is available before attempting to execute if not check_cli_availability(cli_tool): logger.error(f"CLI tool '{cli_tool}' is not available") diff --git a/server/chat/backend/agent/tools/iac/iac_commands_tool.py b/server/chat/backend/agent/tools/iac/iac_commands_tool.py index f72252573..7bc4c4b06 100644 --- a/server/chat/backend/agent/tools/iac/iac_commands_tool.py +++ b/server/chat/backend/agent/tools/iac/iac_commands_tool.py @@ -8,7 +8,7 @@ import shlex from typing import Any, Dict, Optional -from utils.cloud.infrastructure_confirmation import wait_for_user_confirmation +from utils.auth.command_gate import gate_action # Import core execution utilities from .iac_execution_core import ( @@ -235,21 +235,11 @@ def iac_apply( plan_summary_msg = summarize_plan(plan_result.get("stdout", "")) - from ..cloud_tools import get_state_context - - state_context = get_state_context() - context_session_id = ( - state_context.session_id - if state_context and hasattr(state_context, "session_id") - else None - ) - - if not wait_for_user_confirmation( - user_id=user_id, - message=plan_summary_msg, + if not gate_action( + user_id=user_id or "", tool_name="iac_tool", - session_id=context_session_id, - ): + summary=plan_summary_msg, + ).allowed: tool_capture = get_tool_capture() current_tool_call_id = get_current_tool_call_id( tool_name="iac_tool", @@ -541,21 +531,11 @@ def iac_destroy( plan_summary_msg = summarize_plan(destroy_plan_result.get("stdout", "")) - from ..cloud_tools import get_state_context - - state_context = get_state_context() - context_session_id = ( - state_context.session_id - if state_context and hasattr(state_context, "session_id") - else None - ) - - if not wait_for_user_confirmation( - user_id=user_id, - message=plan_summary_msg, + if not gate_action( + user_id=user_id or "", tool_name="iac_tool", - session_id=context_session_id, - ): + summary=plan_summary_msg, + ).allowed: tool_capture = get_tool_capture() current_tool_call_id = get_current_tool_call_id( tool_name="iac_tool", diff --git a/server/chat/backend/agent/tools/mcp_tools.py b/server/chat/backend/agent/tools/mcp_tools.py index f85881498..ef4e96ab7 100644 --- a/server/chat/backend/agent/tools/mcp_tools.py +++ b/server/chat/backend/agent/tools/mcp_tools.py @@ -1252,25 +1252,19 @@ def mcp_tool_wrapper(**kwargs): # Check if this is a destructive MCP tool and ask for confirmation if is_destructive_mcp_tool(original_tool_name): try: - from utils.cloud.infrastructure_confirmation import wait_for_user_confirmation + from utils.auth.command_gate import gate_action from utils.cloud.cloud_utils import get_user_context - from chat.backend.agent.tools.cloud_tools import get_state_context - - # Get user context + context = get_user_context() user_id = context.get('user_id') if isinstance(context, dict) else context - state_context = get_state_context() - session_id = state_context.session_id if state_context and hasattr(state_context, 'session_id') else None - + if user_id: summary_msg = summarize_mcp_tool_action(original_tool_name, kwargs) - if not wait_for_user_confirmation( + if not gate_action( user_id=user_id, - message=summary_msg, tool_name=tool_name, - session_id=session_id - ): - # User cancelled the action + summary=summary_msg, + ).allowed: cancellation_result = f"MCP tool {original_tool_name} cancelled by user." try: if send_tool_completion: diff --git a/server/chat/backend/agent/tools/notion/postmortem.py b/server/chat/backend/agent/tools/notion/postmortem.py index 76e16ee30..c9fcf0d3b 100644 --- a/server/chat/backend/agent/tools/notion/postmortem.py +++ b/server/chat/backend/agent/tools/notion/postmortem.py @@ -516,19 +516,17 @@ def notion_export_postmortem( session_id: Optional[str] = None, ) -> str: """Export an incident's postmortem to a Notion database.""" + _ = session_id if not user_id: return notion_tool_error("user_id is required", code="missing_user") - from chat.backend.agent.tools.cloud_tools import get_state_context - from utils.cloud.infrastructure_confirmation import wait_for_user_confirmation + from utils.auth.command_gate import gate_action - sid = session_id or getattr(get_state_context(), "session_id", None) - if not wait_for_user_confirmation( + if not gate_action( user_id=user_id, - message=f"Export postmortem for incident {incident_id} to Notion database {database_id}?", tool_name="notion_export_postmortem", - session_id=sid, - ): + summary=f"Export postmortem for incident {incident_id} to Notion database {database_id}?", + ).allowed: return notion_tool_error("Operation cancelled by user.", code="cancelled") try: diff --git a/server/chat/backend/agent/tools/notion/structured.py b/server/chat/backend/agent/tools/notion/structured.py index 7f1008f02..b24666e38 100644 --- a/server/chat/backend/agent/tools/notion/structured.py +++ b/server/chat/backend/agent/tools/notion/structured.py @@ -161,19 +161,17 @@ def notion_update_database_properties( session_id: Optional[str] = None, ) -> str: """Update a database's schema (add, rename, change, or remove columns).""" + _ = session_id has_deletions = any(v is None for v in properties.values()) if has_deletions: - from chat.backend.agent.tools.cloud_tools import get_state_context - from utils.cloud.infrastructure_confirmation import wait_for_user_confirmation + from utils.auth.command_gate import gate_action - sid = session_id or getattr(get_state_context(), "session_id", None) cols = [k for k, v in properties.items() if v is None] - if not wait_for_user_confirmation( + if not gate_action( user_id=user_id or "", - message=f"This will permanently delete column(s): {', '.join(cols)}. Proceed?", tool_name="notion_update_database_properties", - session_id=sid, - ): + summary=f"This will permanently delete column(s): {', '.join(cols)}. Proceed?", + ).allowed: return notion_tool_error("Operation cancelled by user.", code="cancelled") def _do(client: Any) -> Dict[str, Any]: diff --git a/server/chat/backend/agent/tools/spinnaker_rca_tool.py b/server/chat/backend/agent/tools/spinnaker_rca_tool.py index 7dda249e2..a490f3835 100644 --- a/server/chat/backend/agent/tools/spinnaker_rca_tool.py +++ b/server/chat/backend/agent/tools/spinnaker_rca_tool.py @@ -225,22 +225,17 @@ def spinnaker_rca( return json.dumps({"error": "Spinnaker is not connected. Configure credentials in Settings > Connectors > Spinnaker."}) try: - from utils.cloud.infrastructure_confirmation import wait_for_user_confirmation - from chat.backend.agent.tools.cloud_tools import get_state_context - - session_id = getattr(get_state_context(), "session_id", None) + from utils.auth.command_gate import gate_action summary = f"Trigger pipeline '{pipeline_name}' for application '{application}'" if parameters: summary += f"\nParameters: {json.dumps(parameters)}" - confirmed = wait_for_user_confirmation( + if not gate_action( user_id=user_id, - message=summary, tool_name="spinnaker_rca", - session_id=session_id, - ) - if not confirmed: + summary=summary, + ).allowed: return json.dumps({"status": "cancelled", "message": "Pipeline trigger cancelled by user"}) except Exception as e: logger.error("[SPINNAKER_RCA] Confirmation flow failed, aborting trigger: %s", e) diff --git a/server/utils/auth/command_gate.py b/server/utils/auth/command_gate.py index d2bec4ddb..bad1cbad0 100644 --- a/server/utils/auth/command_gate.py +++ b/server/utils/auth/command_gate.py @@ -1,7 +1,20 @@ """Command-execution gate — unified policy + safety enforcement with HITL. -Centralizes the four-layer defense-in-depth check that every shell tool must -run before execution: +Two entry points, one confirmation surface: + +* :func:`gate_command` — shell-command path. Runs the four-layer + defense-in-depth check (signature, org allow/deny, LLM judge, session + taint) and prompts on any block. +* :func:`gate_action` — structured-action path for tools with no shell + command (Terraform apply/destroy, Bitbucket PR merges, Notion column + deletes, destructive MCP tools, etc). Always prompts in foreground, + denies in background. No policy/Yes-Always (there is no regex to + persist). + +Both funnel into the same ``_prompt_user`` helper, WS message, React +panel, and DB-backed live state. + +Layers evaluated by :func:`gate_command`: 1. Signature match (utils/security/signature_match.py via command_safety) 2. Org allow/deny (utils/auth/command_policy.py) @@ -135,6 +148,46 @@ def gate_command( _gate_inflight_command.reset(token) +def gate_action( + *, + user_id: Optional[str], + tool_name: str, + summary: str, +) -> GateDecision: + """Human-in-the-loop gate for structured tool actions (no shell command). + + Foreground: prompts the user Yes / No with *summary* as the rendered + action. Background: denies (no interactive user). There is no + Yes-Always here -- these actions are not regex-addressable, so we + cannot persist an allow rule; org policy does not apply. + + Returns the same :class:`GateDecision` shape as :func:`gate_command` + so callers can treat both gates uniformly. + """ + if not user_id: + # Preserve prior behavior of wait_for_user_confirmation helpers, + # which required a user and otherwise denied. + return _block("USER_DECLINED", "Tool call not allowed by user") + + if not _is_foreground(): + return _block("USER_DECLINED", "Tool call not allowed by user") + + decision = _prompt_user( + user_id=user_id, + session_id=_session_id(), + tool_name=tool_name, + command=summary, + block_code="ACTION_CONFIRM", + block_reason=summary, + block_layer="destructive_action", + allow_yes_always=False, + yes_always_changes=[], + org_id=None, + cmd_hash="", + ) + return decision + + def is_session_tainted(session_id: Optional[str], user_id: Optional[str]) -> bool: """Return True iff ``session_id`` has been marked tainted (NeMo input-rail hit on the opening user message of this foreground chat). diff --git a/server/utils/cloud/infrastructure_confirmation.py b/server/utils/cloud/infrastructure_confirmation.py index 1ac1e4d35..c7de6a9aa 100644 --- a/server/utils/cloud/infrastructure_confirmation.py +++ b/server/utils/cloud/infrastructure_confirmation.py @@ -152,83 +152,6 @@ def cancel_pending_confirmations_for_session(session_id: str) -> int: return cancelled_count -def wait_for_user_confirmation( - user_id: str, - message: str, - tool_name: str = "action", - session_id: str = None, - workflow_instance = None, -) -> bool: - """Legacy HITL helper (returns bool). Used by non-gate flows that do not - carry Yes-Always semantics. Does not write ``pending_turn`` because these - flows predate the rehydration model; they keep their existing in-memory - UI flow untouched. - """ - try: - from chat.backend.agent.tools.cloud_tools import get_state_context - state = get_state_context() - if state and getattr(state, 'is_background', False): - logger.warning(f"[BackgroundChat] Denying confirmation for {tool_name} -- no interactive user") - return False - except Exception as e: - logger.debug(f"Could not check background state: {e}") - - timestamp_ms = int(time.time() * 1000) - unique_id = str(uuid.uuid4())[:8] - confirmation_id = f"{timestamp_ms}:{unique_id}" - - payload = { - "type": "execution_confirmation", - "data": { - "message": message, - "status": "awaiting_confirmation", - "user_id": user_id, - "confirmation_id": confirmation_id, - "tool_name": tool_name, - "options": [ - {"text": "Execute", "value": "execute"}, - {"text": "Cancel", "value": "cancel"}, - ], - }, - } - if session_id: - payload["session_id"] = session_id - if user_id: - payload["user_id"] = user_id - - _pending_confirmations[confirmation_id] = { - 'result': None, - 'user_id': user_id, - 'timestamp': time.time() - } - _send_ws(payload, tool_name) - - logger.debug(f"WEBSOCKET: Waiting for user confirmation via WebSocket with ID: {confirmation_id}") - - try: - elapsed = 0.0 - poll_interval = 1.0 - while elapsed < 300: - confirmation_data = _pending_confirmations.get(confirmation_id) - if confirmation_data and confirmation_data.get('result'): - decision = confirmation_data['result'] - logger.debug(f"WEBSOCKET: Received decision for {confirmation_id}: {decision}") - break - time.sleep(poll_interval) - elapsed += poll_interval - else: - decision = None - logger.warning(f"WEBSOCKET: Timeout waiting for confirmation {confirmation_id}") - except Exception as e: - logger.error(f"Error waiting for confirmation: {e}") - decision = None - finally: - _pending_confirmations.pop(confirmation_id, None) - - logger.debug(f"WEBSOCKET: Decision for {confirmation_id}: {decision}") - return decision == "execute" - - def wait_for_user_confirmation_ex( *, user_id: str, From 40ccec1fed2a32feae76dc9b5f69e594aa9e282b Mon Sep 17 00:00:00 2001 From: Damian Loch Date: Fri, 1 May 2026 08:33:02 -0400 Subject: [PATCH 09/13] seed observability_only policy on new org creation Every new org now gets the Observability Only template applied automatically at registration time: both the allowlist and denylist are enabled and populated with the full read-only rule set (AWS/GCP/Azure/k8s read operations, git, system diagnostics, etc.) plus the 13 universal deny rules (rm -rf /, LD_PRELOAD, reverse shells, SUID, etc.). The seeding is idempotent (skips if org already has rows), non-fatal (org creation succeeds even if policy seeding fails), and uses its own admin connection (the org INSERT has already committed). Previously both lists were OFF by default, leaving new orgs completely ungated at the policy layer until an admin manually configured Security settings. Made-with: Cursor --- server/routes/auth_routes.py | 12 +++++++ server/utils/auth/command_policy.py | 54 ++++++++++++++++++++++++++++- 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/server/routes/auth_routes.py b/server/routes/auth_routes.py index 8f6b1576b..786df3dda 100644 --- a/server/routes/auth_routes.py +++ b/server/routes/auth_routes.py @@ -140,6 +140,12 @@ def register(): logging.info(f"New user registered: {email[:3]}***@*** (role=admin, org={org_id})") + try: + from utils.auth.command_policy import seed_default_command_policy + seed_default_command_policy(org_id, user_id) + except Exception as policy_err: + logging.warning(f"Failed to seed command policy for org {org_id}: {policy_err}") + record_audit_event(org_id, user_id, "register", "organization", org_id, {"email": email}, request) @@ -253,6 +259,12 @@ def setup_org(user_id): logging.info(f"User {user_id} created org {org_id} ({org_name})") + try: + from utils.auth.command_policy import seed_default_command_policy + seed_default_command_policy(org_id, user_id) + except Exception as policy_err: + logging.warning(f"Failed to seed command policy for org {org_id}: {policy_err}") + record_audit_event(org_id, user_id, "setup_org", "organization", org_id, {"org_name": org_name}, request) diff --git a/server/utils/auth/command_policy.py b/server/utils/auth/command_policy.py index 3ee00dd0a..31d8427f9 100644 --- a/server/utils/auth/command_policy.py +++ b/server/utils/auth/command_policy.py @@ -11,7 +11,10 @@ atomic command is evaluated independently. One denied sub-command blocks the entire expression. -Default for new orgs: both lists OFF (no enforcement until configured). +Default for new orgs: both lists enabled, seeded with the Observability Only +template (read-only cloud/k8s/git/system commands allowed; universal deny +rules for dangerous patterns). Seeding happens at org creation time so every +org is protected from day one without any admin action. Fail-open on DB error: if rules cannot be fetched, commands are allowed. """ @@ -786,3 +789,52 @@ def get_policy_templates() -> List[dict]: def invalidate_cache(org_id: str) -> None: _cache.pop(org_id, None) + + +def seed_default_command_policy(org_id: str, created_by: str) -> None: + """Insert the Observability Only template and enable both lists for a new org. + + Called immediately after org creation so every org is protected from day + one without requiring any admin action. Uses an independent admin + connection (the org INSERT has already committed by the time this runs). + Idempotent: skips insertion if the org already has policy rows (e.g. + retried registration). + """ + from utils.db.connection_pool import db_pool + from utils.auth.stateless_auth import store_org_preference + + tpl = get_policy_templates()[0] # observability_only + try: + with db_pool.get_admin_connection() as conn: + with conn.cursor() as cur: + # Skip if rules already exist (idempotency). + cur.execute( + "SELECT 1 FROM org_command_policies WHERE org_id = %s LIMIT 1", + (org_id,), + ) + if cur.fetchone(): + return + + cur.execute("SET myapp.current_org_id = %s", (org_id,)) + for mode_key in ("allow", "deny"): + for rule in tpl[mode_key]: + cur.execute( + "INSERT INTO org_command_policies " + "(org_id, mode, pattern, description, priority, updated_by, source) " + "VALUES (%s, %s, %s, %s, %s, %s, 'template')", + (org_id, mode_key, rule["pattern"], + rule["description"], rule["priority"], created_by), + ) + store_org_preference(org_id, "command_policy_allowlist", "on", cursor=cur) + store_org_preference(org_id, "command_policy_denylist", "on", cursor=cur) + store_org_preference(org_id, "command_policy_active_template", tpl["id"], cursor=cur) + conn.commit() + logger.info( + "Seeded default command policy (observability_only) for new org %s", org_id + ) + except Exception: + # Non-fatal: org was created successfully; policy can be applied + # manually via Settings > Security. Log and continue. + logger.exception( + "Failed to seed default command policy for org %s; org creation continues", org_id + ) From febddb5947ebc294de467a7fd94b40b2eb07eed1 Mon Sep 17 00:00:00 2001 From: Damian Loch Date: Fri, 1 May 2026 08:58:41 -0400 Subject: [PATCH 10/13] Address PR review comments from CodeRabbit and Olivier - Remove duplicate gate_command call in _cloud_exec_aws_multi_account - Add tailscale to _CLI_PREFIX so policy patterns match correctly - Use TOOL_NOT_ALLOWED / BACKGROUND_DENIED instead of USER_DECLINED for non-user-initiated blocks in gate_action - Always include session_taint in block_layer/block_reason when tainted - Merge _is_foreground/_session_id into single _get_context() helper - Reset _guardrails_approved_command in gate_command finally block to prevent stale hash reuse across retried commands on the same thread - Fix cancel_pending_confirmations_for_session to actually filter by session_id (was cancelling all pending confirmations across all sessions) - Guard pending_turn rendering by is_own so non-owners don't see interactive confirmation cards - Add ReDoS protection to validate_pattern: length cap + nested quantifier check before compiling client-supplied regex - Use logging.warning with exc_info for seed_default_command_policy errors - Resolve observability_only seed template by id not list position - Cap policy cache at 512 entries with LRU eviction via OrderedDict - Strip sudo flags/values in derive_pattern_from_command so patterns like sudo -u root aws ec2 ... anchor on aws not -u Made-with: Cursor --- client/src/hooks/useChatHistory.ts | 2 +- .../backend/agent/tools/cloud_exec_tool.py | 18 +------- server/routes/auth_routes.py | 8 ++-- server/utils/auth/command_gate.py | 44 +++++++++---------- server/utils/auth/command_policy.py | 40 ++++++++++++++--- .../cloud/infrastructure_confirmation.py | 5 ++- 6 files changed, 64 insertions(+), 53 deletions(-) diff --git a/client/src/hooks/useChatHistory.ts b/client/src/hooks/useChatHistory.ts index 028aa72ed..819a07154 100644 --- a/client/src/hooks/useChatHistory.ts +++ b/client/src/hooks/useChatHistory.ts @@ -233,7 +233,7 @@ export function useChatHistory(): UseChatHistoryReturn { // call so the user sees the prompt immediately on reload. The card is // driven entirely by chat_sessions.pending_turn -- history remains // append-only and contains no mid-turn snapshot. - const pending = data.pending_turn; + const pending = data.is_own ? data.pending_turn : null; const messagesWithPending: ChatMessage[] = pending && pending.confirmation_id ? [ ...(cleanedMessages as ChatMessage[]), diff --git a/server/chat/backend/agent/tools/cloud_exec_tool.py b/server/chat/backend/agent/tools/cloud_exec_tool.py index 559ab38fb..4a0ec3999 100644 --- a/server/chat/backend/agent/tools/cloud_exec_tool.py +++ b/server/chat/backend/agent/tools/cloud_exec_tool.py @@ -1156,22 +1156,6 @@ def _cloud_exec_aws_multi_account( "provider": "aws", }) - # One gate check for the fan-out (signature + org policy + LLM judge + HITL). - # Each per-account invocation is the same command text, so gating once up - # front matches the previous single-prompt UX and avoids N prompts. - from utils.auth.command_gate import gate_command - _gated = command if command.strip().startswith("aws") else f"aws {command}" - _gate = gate_command(user_id=user_id, tool_name="cloud_exec", command=_gated) - if not _gate.allowed: - return json.dumps({ - "success": False, - "error": _gate.block_reason, - "code": _gate.code, - "multi_account": True, - "command": command, - "provider": "aws", - }) - def _run_on_account(conn: dict) -> dict: account_id = conn.get("account_id", "unknown") region = conn.get("region") or "us-east-1" @@ -1337,7 +1321,7 @@ def cloud_exec(provider: str, command: str, user_id: Optional[str] = None, sessi # Prepend CLI prefix so patterns like ^aws\s+ match (cloud_exec receives # the subcommand without the provider prefix, e.g. "ecs list-clusters"). _CLI_PREFIX = {"aws": "aws", "gcp": "gcloud", "azure": "az", - "scaleway": "scw", "ovh": "ovhcloud"} + "scaleway": "scw", "ovh": "ovhcloud", "tailscale": "tailscale"} prefix = _CLI_PREFIX.get(provider.lower(), "") gated_cmd = f"{prefix} {command}" if prefix and not command.strip().startswith(prefix) else command from utils.auth.command_gate import gate_command diff --git a/server/routes/auth_routes.py b/server/routes/auth_routes.py index 786df3dda..b5276a1b5 100644 --- a/server/routes/auth_routes.py +++ b/server/routes/auth_routes.py @@ -144,7 +144,7 @@ def register(): from utils.auth.command_policy import seed_default_command_policy seed_default_command_policy(org_id, user_id) except Exception as policy_err: - logging.warning(f"Failed to seed command policy for org {org_id}: {policy_err}") + logging.warning("Failed to seed command policy for org %s", org_id, exc_info=policy_err) record_audit_event(org_id, user_id, "register", "organization", org_id, {"email": email}, request) @@ -257,13 +257,13 @@ def setup_org(user_id): except Exception as casbin_err: logging.warning(f"Failed to assign Casbin role for {user_id}: {casbin_err}") - logging.info(f"User {user_id} created org {org_id} ({org_name})") - try: from utils.auth.command_policy import seed_default_command_policy seed_default_command_policy(org_id, user_id) except Exception as policy_err: - logging.warning(f"Failed to seed command policy for org {org_id}: {policy_err}") + logging.warning("Failed to seed command policy for org %s", org_id, exc_info=policy_err) + + logging.info(f"User {user_id} created org {org_id} ({org_name})") record_audit_event(org_id, user_id, "setup_org", "organization", org_id, {"org_name": org_name}, request) diff --git a/server/utils/auth/command_gate.py b/server/utils/auth/command_gate.py index bad1cbad0..0b78d97fa 100644 --- a/server/utils/auth/command_gate.py +++ b/server/utils/auth/command_gate.py @@ -85,7 +85,8 @@ def guardrails_approved_hash() -> Optional[str]: class GateDecision: allowed: bool code: str = "" # "" on allow, otherwise POLICY_DENIED / SAFETY_BLOCKED / - # SIGNATURE_MATCHED / USER_DECLINED + # SIGNATURE_MATCHED / USER_DECLINED / BACKGROUND_DENIED / + # TOOL_NOT_ALLOWED block_reason: str = "" @@ -96,25 +97,16 @@ def _block(code: str, reason: str) -> GateDecision: return GateDecision(allowed=False, code=code, block_reason=reason) -def _is_foreground() -> bool: - """Return True iff the current execution is a foreground (interactive) chat.""" +def _get_context() -> tuple[bool, Optional[str]]: + """Return (is_foreground, session_id) from the current execution state.""" try: from utils.cloud.cloud_utils import get_state_context state = get_state_context() if state is None: - return False - return not bool(getattr(state, "is_background", False)) + return False, None + return not bool(getattr(state, "is_background", False)), getattr(state, "session_id", None) except Exception: - return False - - -def _session_id() -> Optional[str]: - try: - from utils.cloud.cloud_utils import get_state_context - state = get_state_context() - return getattr(state, "session_id", None) if state else None - except Exception: - return None + return False, None def gate_command( @@ -141,11 +133,13 @@ def gate_command( return _ALLOWED token = _gate_inflight_command.set(cmd_hash) + approved_token = _guardrails_approved_command.set(None) try: return _gate_impl(user_id=user_id, tool_name=tool_name, command=command, cmd_hash=cmd_hash) finally: _gate_inflight_command.reset(token) + _guardrails_approved_command.reset(approved_token) def gate_action( @@ -167,14 +161,15 @@ def gate_action( if not user_id: # Preserve prior behavior of wait_for_user_confirmation helpers, # which required a user and otherwise denied. - return _block("USER_DECLINED", "Tool call not allowed by user") + return _block("TOOL_NOT_ALLOWED", "Tool call not allowed: no user context") - if not _is_foreground(): - return _block("USER_DECLINED", "Tool call not allowed by user") + foreground, session_id = _get_context() + if not foreground: + return _block("BACKGROUND_DENIED", "Tool call not allowed in background context") decision = _prompt_user( user_id=user_id, - session_id=_session_id(), + session_id=session_id, tool_name=tool_name, command=summary, block_code="ACTION_CONFIRM", @@ -245,8 +240,7 @@ def _gate_impl(*, user_id: str, tool_name: str, command: str, cmd_hash: str) -> from utils.security.command_safety import evaluate_command as safety_evaluate org_id = get_org_id_for_user(user_id) - foreground = _is_foreground() - session_id = _session_id() + foreground, session_id = _get_context() # Evaluate all layers unconditionally so we can report the combined # block state to the user. The previous short-circuit (return on first @@ -281,8 +275,10 @@ def _gate_impl(*, user_id: str, tool_name: str, command: str, cmd_hash: str) -> else None ) layers = [l for l in (safety_layer, policy_layer) if l] - if not layers and tainted: - layers = ["session_taint"] + if tainted: + layers.append("session_taint") + if not layers: + layers = ["unknown"] block_layer = "+".join(layers) code = safety_code or ("POLICY_DENIED" if policy_blocked else "SESSION_TAINTED") @@ -294,7 +290,7 @@ def _gate_impl(*, user_id: str, tool_name: str, command: str, cmd_hash: str) -> "organization policy: " + (policy_verdict.rule_description or "matched organization policy")[:200] ) - if not reasons and tainted: + if tainted: reasons.append("session flagged by input safety check; approval required") block_reason = "Command blocked by " + "; ".join(reasons) diff --git a/server/utils/auth/command_policy.py b/server/utils/auth/command_policy.py index 31d8427f9..5e4d1e5c5 100644 --- a/server/utils/auth/command_policy.py +++ b/server/utils/auth/command_policy.py @@ -18,6 +18,7 @@ Fail-open on DB error: if rules cannot be fetched, commands are allowed. """ +import collections import logging import re import shlex @@ -30,6 +31,7 @@ logger = logging.getLogger(__name__) _CACHE_TTL = 30 # seconds +_CACHE_MAX = 512 # max orgs in cache before LRU eviction _CacheEntry = Tuple[ List["PolicyRule"], # allow rules @@ -37,7 +39,7 @@ "ListStates", float, # monotonic timestamp ] -_cache: Dict[str, _CacheEntry] = {} +_cache: "collections.OrderedDict[str, _CacheEntry]" = collections.OrderedDict() @dataclass(frozen=True) @@ -128,10 +130,14 @@ def _get_cached(org_id: str) -> Tuple[List[PolicyRule], List[PolicyRule], ListSt if entry is not None: allow, deny, states, ts = entry if time.monotonic() - ts < _CACHE_TTL: + _cache.move_to_end(org_id) return allow, deny, states allow, deny, states = _fetch(org_id) _cache[org_id] = (allow, deny, states, time.monotonic()) + _cache.move_to_end(org_id) + if len(_cache) > _CACHE_MAX: + _cache.popitem(last=False) # evict LRU return allow, deny, states @@ -309,8 +315,16 @@ def evaluate_compound_command( return last_verdict +_REDOS_RE = re.compile(r"(\(.*\+.*\)[\+\*]|(\.\*){3,}|\(\?:.*\)[\+\*]{2})") +_PATTERN_MAX_LEN = 500 + + def validate_pattern(pattern: str) -> Optional[str]: - """Return an error string if *pattern* is not valid regex, else None.""" + """Return an error string if *pattern* is not a safe, valid regex, else None.""" + if len(pattern) > _PATTERN_MAX_LEN: + return f"pattern too long (max {_PATTERN_MAX_LEN} chars)" + if _REDOS_RE.search(pattern): + return "pattern contains nested quantifiers that could cause ReDoS" try: re.compile(pattern) return None @@ -319,7 +333,7 @@ def validate_pattern(pattern: str) -> Optional[str]: # Leading `VAR=value` env assignments the user didn't intend as the "command". -_ENV_ASSIGN_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=\S*$") +_ENV_ASSIGN_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=.*$") # Shell keywords / operators that mean "the first token is not a CLI name". # When we see one of these at position 0 (after stripping sudo / env assigns), @@ -348,6 +362,7 @@ def derive_pattern_from_command(command: str) -> str: Examples: "sudo kubectl delete pod foo" -> ^kubectl delete pod\\b + "sudo -u root aws ec2 terminate" -> ^aws ec2 terminate\\b "KUBECONFIG=/x kubectl get pods" -> ^kubectl get pods\\b "aws ec2 terminate-instances --id" -> ^aws ec2 terminate-instances\\b "for i in {1..10}; do echo $i; done" -> ^for i in \\{1\\.\\.10\\}; do echo \\$i; done$ @@ -357,8 +372,21 @@ def derive_pattern_from_command(command: str) -> str: tokens = shlex.split(stripped, posix=True) except ValueError: tokens = stripped.split() - while tokens and (tokens[0] == "sudo" or _ENV_ASSIGN_RE.match(tokens[0])): + # Strip leading env assignments (VAR=value) and sudo with its flags/values. + while tokens and _ENV_ASSIGN_RE.match(tokens[0]): + tokens.pop(0) + if tokens and tokens[0] == "sudo": tokens.pop(0) + # Consume sudo's own flags (e.g. -u root, -E, --user=root) until we + # reach the actual CLI token. + while tokens and tokens[0].startswith("-"): + flag = tokens.pop(0) + # Flags without embedded '=' consume the next token as a value. + if "=" not in flag and tokens and not tokens[0].startswith("-"): + tokens.pop(0) + # Strip any remaining env assignments set after sudo (sudo VAR=val cmd). + while tokens and _ENV_ASSIGN_RE.match(tokens[0]): + tokens.pop(0) if not tokens or tokens[0] in _SHELL_NON_CLI_LEADERS: return r"^" + re.escape(stripped) + r"$" parts = [tokens[0]] @@ -803,7 +831,9 @@ def seed_default_command_policy(org_id: str, created_by: str) -> None: from utils.db.connection_pool import db_pool from utils.auth.stateless_auth import store_org_preference - tpl = get_policy_templates()[0] # observability_only + tpl = next((t for t in get_policy_templates() if t["id"] == "observability_only"), None) + if tpl is None: + raise ValueError("observability_only policy template not found") try: with db_pool.get_admin_connection() as conn: with conn.cursor() as cur: diff --git a/server/utils/cloud/infrastructure_confirmation.py b/server/utils/cloud/infrastructure_confirmation.py index c7de6a9aa..ff960b784 100644 --- a/server/utils/cloud/infrastructure_confirmation.py +++ b/server/utils/cloud/infrastructure_confirmation.py @@ -143,8 +143,8 @@ def cancel_pending_confirmations_for_session(session_id: str) -> int: return 0 cancelled_count = 0 - for confirmation_id, confirmation_data in _pending_confirmations.items(): - if confirmation_data.get('result') is None: + for confirmation_id, confirmation_data in list(_pending_confirmations.items()): + if confirmation_data.get('session_id') == session_id and confirmation_data.get('result') is None: confirmation_data['result'] = 'cancel' cancelled_count += 1 logger.info(f"Cancelled pending confirmation {confirmation_id} for session {session_id}") @@ -217,6 +217,7 @@ def wait_for_user_confirmation_ex( _pending_confirmations[confirmation_id] = { 'result': None, 'user_id': user_id, + 'session_id': session_id, 'timestamp': time.time(), } # Register the waiter before sending so a fast client response can never From 7914dda9c3ddad4a38c42e2efa472e6fcd04fe42 Mon Sep 17 00:00:00 2001 From: Damian Loch Date: Fri, 1 May 2026 09:02:23 -0400 Subject: [PATCH 11/13] Replace REDOS_RE with character-scan to fix CodeQL warning _REDOS_RE ran a complex regex on user-supplied input, which is exactly the polynomial-regex-on-uncontrolled-data pattern CodeQL flags. Replace it with _has_nested_quantifiers(), a simple character walk that detects (X+)+ style nesting without executing any regex against user data. Made-with: Cursor --- server/utils/auth/command_policy.py | 35 +++++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/server/utils/auth/command_policy.py b/server/utils/auth/command_policy.py index 5e4d1e5c5..0e3003b8a 100644 --- a/server/utils/auth/command_policy.py +++ b/server/utils/auth/command_policy.py @@ -315,15 +315,46 @@ def evaluate_compound_command( return last_verdict -_REDOS_RE = re.compile(r"(\(.*\+.*\)[\+\*]|(\.\*){3,}|\(\?:.*\)[\+\*]{2})") _PATTERN_MAX_LEN = 500 +def _has_nested_quantifiers(pattern: str) -> bool: + """Character-scan heuristic for (X+)+ style ReDoS patterns. + + Walks the pattern without running any regex on user data (avoids the + CodeQL "polynomial regex on uncontrolled input" warning). Tracks open + groups and flags when a quantifier follows a closing group that itself + contained a quantifier. + """ + group_has_quant: list[bool] = [] + i = 0 + while i < len(pattern): + ch = pattern[i] + if ch == "\\": + i += 2 # skip escaped char + continue + if ch == "(": + group_has_quant.append(False) + elif ch == ")": + if group_has_quant: + had = group_has_quant.pop() + j = i + 1 + # skip non-greedy modifier + if j < len(pattern) and pattern[j] == "?": + j += 1 + if had and j < len(pattern) and pattern[j] in "+*": + return True + elif ch in "+*{" and group_has_quant: + group_has_quant[-1] = True + i += 1 + return False + + def validate_pattern(pattern: str) -> Optional[str]: """Return an error string if *pattern* is not a safe, valid regex, else None.""" if len(pattern) > _PATTERN_MAX_LEN: return f"pattern too long (max {_PATTERN_MAX_LEN} chars)" - if _REDOS_RE.search(pattern): + if _has_nested_quantifiers(pattern): return "pattern contains nested quantifiers that could cause ReDoS" try: re.compile(pattern) From c97572f50f0353725385634fecaad7fcef516881 Mon Sep 17 00:00:00 2001 From: Damian Loch Date: Fri, 1 May 2026 09:09:39 -0400 Subject: [PATCH 12/13] Skip safety judge when command matches an explicit allowlist rule An org admin explicitly adding a pattern to the allowlist is a trust signal that should not be overridden by the LLM safety judge. Previously, safety_evaluate ran unconditionally so allowlisted commands (e.g. gcloud instances delete) still triggered a Yes/No prompt if the judge flagged them as dangerous. Now: if evaluate_compound_command returns allowed=True with a non-default rule_description (i.e. an explicit allow rule matched, not just "lists disabled"), skip the safety judge entirely. Default-allow (lists off) and no-rule-match paths still run the full judge. Made-with: Cursor --- server/utils/auth/command_gate.py | 27 ++++++++++++++++++--------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/server/utils/auth/command_gate.py b/server/utils/auth/command_gate.py index 0b78d97fa..ed9a055aa 100644 --- a/server/utils/auth/command_gate.py +++ b/server/utils/auth/command_gate.py @@ -242,17 +242,26 @@ def _gate_impl(*, user_id: str, tool_name: str, command: str, cmd_hash: str) -> org_id = get_org_id_for_user(user_id) foreground, session_id = _get_context() - # Evaluate all layers unconditionally so we can report the combined - # block state to the user. The previous short-circuit (return on first - # safety block) prevented Always from showing when the policy layer - # would have also fired. - safety_decision = safety_evaluate( - command, tool=tool_name, user_id=user_id, session_id=session_id, - ) policy_verdict: CommandVerdict = evaluate_compound_command(org_id, command) - - safety_blocked = safety_decision.blocked policy_blocked = not policy_verdict.allowed + + # An explicit allowlist match is a trust signal from the org admin — skip + # the safety judge so approved patterns don't generate spurious prompts. + # "Policy lists are disabled" / None means default-allow, not an explicit + # rule hit, so we still run the safety judge in those cases. + explicitly_allowed = ( + policy_verdict.allowed + and policy_verdict.rule_description not in (None, "Policy lists are disabled") + ) + if not explicitly_allowed: + safety_decision = safety_evaluate( + command, tool=tool_name, user_id=user_id, session_id=session_id, + ) + safety_blocked = safety_decision.blocked + else: + safety_decision = None + safety_blocked = False + tainted = foreground and is_session_tainted(session_id, user_id) if not (safety_blocked or policy_blocked or tainted): From cd9dcdbcc7fe997f7f1c1390ecba0f08044898fa Mon Sep 17 00:00:00 2001 From: Damian Loch Date: Fri, 1 May 2026 09:12:48 -0400 Subject: [PATCH 13/13] Revert "Skip safety judge when command matches an explicit allowlist rule" This reverts commit c97572f50f0353725385634fecaad7fcef516881. --- server/utils/auth/command_gate.py | 27 +++++++++------------------ 1 file changed, 9 insertions(+), 18 deletions(-) diff --git a/server/utils/auth/command_gate.py b/server/utils/auth/command_gate.py index ed9a055aa..0b78d97fa 100644 --- a/server/utils/auth/command_gate.py +++ b/server/utils/auth/command_gate.py @@ -242,26 +242,17 @@ def _gate_impl(*, user_id: str, tool_name: str, command: str, cmd_hash: str) -> org_id = get_org_id_for_user(user_id) foreground, session_id = _get_context() - policy_verdict: CommandVerdict = evaluate_compound_command(org_id, command) - policy_blocked = not policy_verdict.allowed - - # An explicit allowlist match is a trust signal from the org admin — skip - # the safety judge so approved patterns don't generate spurious prompts. - # "Policy lists are disabled" / None means default-allow, not an explicit - # rule hit, so we still run the safety judge in those cases. - explicitly_allowed = ( - policy_verdict.allowed - and policy_verdict.rule_description not in (None, "Policy lists are disabled") + # Evaluate all layers unconditionally so we can report the combined + # block state to the user. The previous short-circuit (return on first + # safety block) prevented Always from showing when the policy layer + # would have also fired. + safety_decision = safety_evaluate( + command, tool=tool_name, user_id=user_id, session_id=session_id, ) - if not explicitly_allowed: - safety_decision = safety_evaluate( - command, tool=tool_name, user_id=user_id, session_id=session_id, - ) - safety_blocked = safety_decision.blocked - else: - safety_decision = None - safety_blocked = False + policy_verdict: CommandVerdict = evaluate_compound_command(org_id, command) + safety_blocked = safety_decision.blocked + policy_blocked = not policy_verdict.allowed tainted = foreground and is_session_tainted(session_id, user_id) if not (safety_blocked or policy_blocked or tainted):