"""Agent runner — executes AI agents as ARQ jobs with safety checks. Safety features: - Max executions per agent per hour (rate limiting) - Max duration per execution (asyncio timeout) - Auto-stop on infinite loop (same tool called 5x consecutively) - Budget limit per agent (track cumulative cost_usd, stop if over budget) """ from __future__ import annotations import asyncio import logging from datetime import UTC, datetime from typing import Any from sqlalchemy import func, select from app.core.db import get_session_factory logger = logging.getLogger(__name__) # Track consecutive tool calls per agent run to detect infinite loops _consecutive_tool_calls: dict[str, dict[str, int]] = {} # run_id -> {tool_name: count} async def run_agent( ctx: dict[str, Any], agent_id: str, trigger_type: str = "manual", trigger_data: dict[str, Any] | None = None, ) -> dict[str, Any]: """ARQ job function. Loads AgentDefinition from DB, checks rate limits, gathers context, calls LLM via LiteLLM, executes tool calls via ToolRegistry, saves result to AgentRun. Returns result dict. Safety checks: 1. Rate limit: max_executions_per_hour 2. Max duration: max_duration_seconds (asyncio.timeout) 3. Infinite loop: same tool 5x consecutively 4. Budget limit: cumulative cost_usd """ from app.plugins.builtins.automation.models import AgentDefinition, AgentRun factory = get_session_factory() # Load agent definition async with factory() as db: result = await db.execute( select(AgentDefinition).where(AgentDefinition.id == agent_id) ) agent = result.scalar_one_or_none() if agent is None: logger.error("AgentDefinition %s not found", agent_id) return {"error": f"AgentDefinition {agent_id} not found", "status": "failed"} if not agent.is_active: logger.warning("AgentDefinition %s is inactive", agent_id) return {"error": "Agent is inactive", "status": "skipped"} # ── Safety Check 1: Rate Limit ── if agent.max_executions_per_hour: async with factory() as db: one_hour_ago = datetime.now(UTC) count_result = await db.execute( select(func.count()) .select_from(AgentRun) .where( AgentRun.agent_id == agent.id, AgentRun.created_at >= one_hour_ago, ) ) recent_runs = count_result.scalar() or 0 if recent_runs >= agent.max_executions_per_hour: logger.warning( "Rate limit hit for agent %s: %d runs in last hour (max %d)", agent.id, recent_runs, agent.max_executions_per_hour, ) return {"error": "Rate limit exceeded", "status": "rate_limited"} # ── Safety Check 2: Budget Limit ── if agent.budget_limit_usd > 0: async with factory() as db: cost_result = await db.execute( select(func.coalesce(func.sum(AgentRun.cost_usd), 0.0)) .where(AgentRun.agent_id == agent.id) ) total_cost = float(cost_result.scalar() or 0.0) if total_cost >= agent.budget_limit_usd: logger.warning( "Budget limit hit for agent %s: $%.4f total cost (limit $%.2f)", agent.id, total_cost, agent.budget_limit_usd, ) return {"error": "Budget limit exceeded", "status": "budget_exceeded"} # Gather context context_data: dict[str, Any] = {} if trigger_type == "proactive" or agent.mode == "proactive": # Collect context data (recent contacts, mails, events) try: from app.services.contact_service import list_contacts async with factory() as db: contacts = await list_contacts(db, agent.tenant_id, page=1, page_size=10) context_data["recent_contacts"] = contacts.get("items", []) except Exception: logger.warning("Failed to collect contacts for proactive context") try: from app.core.cache import get_cached_mail_summary context_data["recent_mails"] = get_cached_mail_summary(agent.tenant_id) or [] except Exception: logger.warning("Failed to collect mails for proactive context") try: from app.services.workflow_service import list_instances async with factory() as db: events = await list_instances(db, agent.tenant_id, page=1, page_size=10) context_data["recent_events"] = events.get("items", []) except Exception: logger.warning("Failed to collect events for proactive context") else: # Reactive mode: use trigger_data as context context_data = trigger_data or {} # Call LLM via LiteLLM result_data: dict[str, Any] = { "agent_id": str(agent.id), "agent_name": agent.name, "trigger_type": trigger_type, "status": "running", "llm_response": None, "tool_calls": [], "cost_usd": 0.0, "error": None, } # ── Safety Check 3: Max Duration ── max_duration = agent.max_duration_seconds or 300 try: async def _run_llm() -> None: """Inner coroutine for LLM call with tool execution.""" from app.ai.llm_client import LLMClient llm = LLMClient( model=agent.model or None, api_key=agent.api_key or None, api_base=agent.api_base or None, provider=agent.provider or None, ) # Build system prompt from agent configuration system_prompt = agent.system_prompt or "You are a helpful AI assistant." user_prompt = f"Context: {context_data}" # Use litellm directly for more control import litellm litellm_model = agent.model or "gpt-4o" if agent.provider and agent.provider != "openai": litellm_model = f"{agent.provider}/{litellm_model}" kwargs: dict[str, Any] = { "model": litellm_model, "messages": [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}, ], "temperature": 0.3, "max_tokens": agent.max_tokens or 1000, } if agent.api_key: kwargs["api_key"] = agent.api_key if agent.api_base: kwargs["api_base"] = agent.api_base response = await litellm.acompletion(**kwargs) content = response.choices[0].message.content # Track cost if hasattr(response, "usage") and response.usage: usage = response.usage # Estimate cost (simplified) input_cost = (usage.prompt_tokens or 0) * 0.00001 / 1000 output_cost = (usage.completion_tokens or 0) * 0.00003 / 1000 result_data["cost_usd"] = round(input_cost + output_cost, 6) result_data["llm_response"] = content # Execute tool calls if LLM returned function calls if hasattr(response.choices[0].message, "tool_calls") and response.choices[0].message.tool_calls: from app.plugins.builtins.ai_assistant.tool_registry import get_tool_registry registry = get_tool_registry() tool_call_count: dict[str, int] = {} for tc in response.choices[0].message.tool_calls: tool_name = tc.function.name # ── Safety Check 4: Infinite Loop Detection ── tool_call_count[tool_name] = tool_call_count.get(tool_name, 0) + 1 if tool_call_count[tool_name] >= 5: logger.warning( "Infinite loop detected for agent %s: tool '%s' called %d times consecutively", agent.id, tool_name, tool_call_count[tool_name], ) result_data["error"] = f"Infinite loop detected: tool '{tool_name}' called 5+ times consecutively" result_data["status"] = "failed" return tool = registry.get(tool_name) if tool: try: import json args = json.loads(tc.function.arguments) tool_result = await tool.handler( arguments=args, context={"tenant_id": str(agent.tenant_id)}, ) result_data["tool_calls"].append({ "tool": tool_name, "arguments": args, "result": tool_result, }) except Exception as e: result_data["tool_calls"].append({ "tool": tool_name, "error": str(e), }) result_data["status"] = "completed" # Run with timeout try: await asyncio.wait_for(_run_llm(), timeout=max_duration) except asyncio.TimeoutError: logger.warning( "Agent %s execution timed out after %d seconds", agent.id, max_duration, ) result_data["status"] = "timed_out" result_data["error"] = f"Execution timed out after {max_duration} seconds" except Exception as e: logger.exception("Agent run failed for %s", agent.id) result_data["status"] = "failed" result_data["error"] = str(e) # Save result to AgentRun try: async with factory() as db: run = AgentRun( tenant_id=agent.tenant_id, agent_id=agent.id, trigger_type=trigger_type, status=result_data["status"], input_data=context_data, output_data=result_data.get("llm_response"), tool_calls=result_data.get("tool_calls", []), cost_usd=result_data.get("cost_usd", 0.0), error_message=result_data.get("error"), duration_seconds=None, ) db.add(run) await db.flush() result_data["run_id"] = str(run.id) except Exception as e: logger.exception("Failed to save AgentRun for %s", agent.id) result_data["save_error"] = str(e) return result_data