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/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/client/src/components/tool-calls/ToolExecutionWidget.tsx b/client/src/components/tool-calls/ToolExecutionWidget.tsx index 2d26b0219..29940207e 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("{")) { @@ -466,56 +473,14 @@ 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 +551,152 @@ const ToolExecutionWidget = ({ tool, className, sendMessage, sendRaw, onToolUpda } export default ToolExecutionWidget + +interface ConfirmationPanelProps { + tool: ToolCall + command: string + userId?: string + sessionId?: string + sendRaw?: (data: string) => boolean + onToolUpdate?: (updatedTool: Partial) => void +} + +// 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) + + 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 [alwaysOpen, setAlwaysOpen] = React.useState(false) + + 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 + const sent = sendRaw(JSON.stringify(payload)) + if (!sent) return + if (decision === 'cancel') { + onToolUpdate?.({ status: 'completed', output: 'Operation cancelled by user' }) + } else { + onToolUpdate?.({ status: 'running' }) + } + } + + return ( +
+
+ + + 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 && ( + + {c.pattern} + + )} + + ) : ( + <> + + setEdited(prev => ({ ...prev, [String(i)]: e.target.value }))} + spellCheck={false} + aria-label="Allow-rule regex pattern" + /> + + )} +
  • + ))} +
+
+ +
+
+
+ )} +
+
+ ) +} diff --git a/client/src/hooks/useChatHistory.ts b/client/src/hooks/useChatHistory.ts index dc6ded63e..819a07154 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.is_own ? data.pending_turn : null; + 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/client/src/hooks/useMessageHandler.ts b/client/src/hooks/useMessageHandler.ts index 62534a333..c332b4ccd 100644 --- a/client/src/hooks/useMessageHandler.ts +++ b/client/src/hooks/useMessageHandler.ts @@ -115,7 +115,10 @@ export const useMessageHandler = ({ ...tc, status: 'awaiting_confirmation' as const, confirmation_id: confirmationId, - confirmation_message: message.data.message + 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, }; } return tc; 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 61ef35ff4..4a0ec3999 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,26 +1156,6 @@ 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", - }) - def _run_on_account(conn: dict) -> dict: account_id = conn.get("account_id", "unknown") region = conn.get("region") or "us-east-1" @@ -1409,25 +1317,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 + "scaleway": "scw", "ovh": "ovhcloud", "tailscale": "tailscale"} 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(), }) @@ -1834,42 +1739,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/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/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/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/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/routes/auth_routes.py b/server/routes/auth_routes.py index 8f6b1576b..b5276a1b5 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("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) @@ -251,6 +257,12 @@ def setup_org(user_id): except Exception as casbin_err: logging.warning(f"Failed to assign Casbin role for {user_id}: {casbin_err}") + 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("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, 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/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_gate.py b/server/utils/auth/command_gate.py new file mode 100644 index 000000000..0b78d97fa --- /dev/null +++ b/server/utils/auth/command_gate.py @@ -0,0 +1,444 @@ +"""Command-execution gate — unified policy + safety enforcement with HITL. + +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) + 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 / BACKGROUND_DENIED / + # TOOL_NOT_ALLOWED + block_reason: str = "" + + +_ALLOWED = GateDecision(allowed=True) + + +def _block(code: str, reason: str) -> GateDecision: + return GateDecision(allowed=False, code=code, block_reason=reason) + + +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, None + return not bool(getattr(state, "is_background", False)), getattr(state, "session_id", None) + except Exception: + return False, 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) + 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( + *, + 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("TOOL_NOT_ALLOWED", "Tool call not allowed: no user context") + + 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, + 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). + + Tainted sessions force every command through user confirmation, even when + 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 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 AND user_id = %s", + (session_id, user_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, + ) + 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, 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 + 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 + # signature+judge for this command on the same invocation. + _guardrails_approved_command.set(cmd_hash) + return _ALLOWED + + # 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 + ) + layers = [l for l in (safety_layer, policy_layer) if l] + 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") + 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 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) + + # 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=code, + block_reason=block_reason, + block_layer=block_layer or "unknown", + allow_yes_always=bool(changes), + yes_always_changes=changes, + org_id=org_id if policy_blocked else None, + cmd_hash=cmd_hash, + ) + if decision.allowed: + _guardrails_approved_command.set(cmd_hash) + 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, + 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, + } + # 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:", + "changes": [ + { + "action": ch.action, + "rule_id": ch.rule_id, + "pattern": ch.pattern, + "description": ch.description, + "editable": ch.editable, + } + for ch in yes_always_changes + ], + } + + # 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, + 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 offer_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..0e3003b8a 100644 --- a/server/utils/auth/command_policy.py +++ b/server/utils/auth/command_policy.py @@ -11,12 +11,17 @@ 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. """ +import collections import logging import re +import shlex import time from dataclasses import dataclass, field from typing import Dict, List, Optional, Tuple @@ -26,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 @@ -33,13 +39,17 @@ "ListStates", float, # monotonic timestamp ] -_cache: Dict[str, _CacheEntry] = {} +_cache: "collections.OrderedDict[str, _CacheEntry]" = collections.OrderedDict() @dataclass(frozen=True) 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) @@ -120,15 +130,26 @@ 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 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) @@ -138,18 +159,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_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) - return CommandVerdict(allowed=False, rule_description="No matching allow rule") + 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+|<\(|>\(") @@ -277,8 +315,47 @@ def evaluate_compound_command( return last_verdict +_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 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 _has_nested_quantifiers(pattern): + return "pattern contains nested quantifiers that could cause ReDoS" try: re.compile(pattern) return None @@ -286,6 +363,150 @@ 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_]*=.*$") + +# 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. + + 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 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$ + """ + stripped = command.strip() + try: + tokens = shlex.split(stripped, posix=True) + except ValueError: + tokens = stripped.split() + # 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]] + for tok in tokens[1:]: + if tok.startswith("-"): + break + parts.append(tok) + # 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" + + +@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) @@ -627,3 +848,54 @@ 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 = 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: + # 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 + ) diff --git a/server/utils/cloud/infrastructure_confirmation.py b/server/utils/cloud/infrastructure_confirmation.py index 901125b8a..ff960b784 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 -# 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,83 +101,89 @@ 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 + 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}") + # 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}") - + 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: + 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}") - + return cancelled_count -def wait_for_user_confirmation( +def wait_for_user_confirmation_ex( + *, user_id: str, message: str, - tool_name: str = "action", - session_id: str = None, - workflow_instance = None, -) -> bool: - """ - 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. + 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 used exclusively by the command gate. + + 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 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() if state and getattr(state, 'is_background', False): logger.warning(f"[BackgroundChat] Denying confirmation for {tool_name} -- no interactive user") - return False + return {"decision": "cancel", "edited_patterns": {}} 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}" - payload = { + payload: Dict[str, Any] = { "type": "execution_confirmation", "data": { "message": message, @@ -120,146 +191,59 @@ def wait_for_user_confirmation( "user_id": user_id, "confirmation_id": confirmation_id, "tool_name": tool_name, - "options": [ - {"text": "Execute", "value": "execute"}, - {"text": "Cancel", "value": "cancel"}, - ], + "options": options, }, } - - # Add session and user information at the top level for filtering + if extra: + payload["data"].update(extra) 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) + # 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(), + }) - # Store the confirmation request (without asyncio.Event for simplicity) _pending_confirmations[confirmation_id] = { 'result': None, 'user_id': user_id, - 'timestamp': time.time() + 'session_id': session_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}") - logger.debug(f"WEBSOCKET: Waiting for user confirmation via WebSocket with ID: {confirmation_id}") - + decision: Optional[str] = None + edited: Dict[str, Any] = {} 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) - 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}") + 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: # Timeout occurred - 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] - - logger.debug(f"WEBSOCKET: Decision for {confirmation_id}: {decision}") - return decision == "execute" - - -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.""" - 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: - # Update the tool call to show it's awaiting confirmation - tool_call['status'] = 'awaiting_confirmation' - tool_call['confirmation_id'] = confirmation_id - 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())}") - - # 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) - + logger.warning(f"WEBSOCKET: Timeout waiting for confirmation_ex {confirmation_id}") except Exception as e: - logger.error(f"Error saving UI messages with confirmation: {e}") - return False \ No newline at end of file + 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} diff --git a/server/utils/db/db_utils.py b/server/utils/db/db_utils.py index 3a07dff4e..34fe5c682 100644 --- a/server/utils/db/db_utils.py +++ b/server/utils/db/db_utils.py @@ -1461,6 +1461,40 @@ 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 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(""" 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()