DEV Community

kirandeepjassal-crypto
kirandeepjassal-crypto

Posted on Originally published at prepstack.co.in

How to Build an AI Agent That Doesn't Go Off the Rails — The Agent Loop, Planning, Tool Use, Memory & Guardrails

A demo agent takes ten lines and a weekend. A production agent is a different animal: it loops forever on an ambiguous task, calls a tool that doesn't exist, spends $40 in tokens on one question, silently returns a confidently wrong answer, and — the first time you let it near a write action — deletes something it shouldn't.

The gap between "I wired up the OpenAI API with tools" and "this agent runs unattended in production" is exactly where senior AI engineering lives. This is the whole ladder: what an agent actually is, the loop that drives it, planning, tool use done right, memory, reflection, and the guardrails that turn a toy into a system you can trust.

Every example is from Mattrx — a real multi-tenant analytics platform whose Insights agent answers questions like "why did tenant X's conversions drop last week?" by reasoning over tools, orchestrated with LangGraph and traced with Langfuse on a FastAPI service. We hardened it from a demo that looped and overspent into an agent with p95 4.2s → 1.8s and bounded cost.

The demo-to-production gap

The first agent everyone builds works in the demo and fails in production, because a demo hides five things:

  1. It loops forever — on a vague task, the model keeps "thinking" with no exit.
  2. It hallucinates tools — calls send_email when you only gave it read_kpis.
  3. It has no budget — one hard question fans out into 30 LLM calls and a surprise bill.
  4. It fails silently — a tool throws, the model shrugs, you get a confident wrong answer.
  5. It has no brakes — the moment a tool can write, an unbounded agent is a liability.

What is an agent, really?

Strip the hype: an agent is an LLM in a loop with tools and memory, deciding its own next step until a goal is met. Separate three things people conflate:

  • A single LLM call — one prompt, one answer. Control flow you write.
  • A workflow / chain — a fixed sequence of steps. You decide the path; predictable, cheaper, testable.
  • An agent — the LLM decides the path at runtime. Flexible but non-deterministic and harder to bound.

The senior instinct: don't reach for an agent when a workflow will do. Agents earn their non-determinism only when the path genuinely can't be known in advance.

The naive agent, and why it fails

# DO NOT ship this. No exit, no tool validation, no budget, no error handling.
def naive_agent(task, tools):
    messages = [{"role": "user", "content": task}]
    while True:                                   # (1) infinite loop risk
        resp = llm(messages, tools=tools)
        if resp.tool_calls:
            for call in resp.tool_calls:
                result = tools[call.name](**call.args)   # (2) KeyError on hallucinated tool
                messages.append(tool_result(call, result))  # (4) crashes if the tool throws
        else:
            return resp.content                  # (3) no token/step budget
Enter fullscreen mode Exit fullscreen mode

Four bugs in ten lines. The fix isn't a bigger model — it's structure: a bounded loop, validated tools, budgets, error handling, and guardrails.

The agent loop (ReAct): reason, act, observe

The core pattern is ReAct — Reason, Act, Observe, repeat. The single most important addition over the naive version: a hard step cap so it can never loop forever.

def run_agent(task: str, tools: dict, max_steps: int = 8) -> str:
    messages = [system_prompt(tools), {"role": "user", "content": task}]
    for step in range(max_steps):                      # HARD cap -> no infinite loop
        resp = llm(messages, tools=tool_schemas(tools))
        if not resp.tool_calls:
            return resp.content                        # the model decided it's done
        for call in resp.tool_calls:
            if call.name not in tools:                 # reject hallucinated tools
                result = f"error: unknown tool '{call.name}'"
            else:
                result = safe_invoke(tools[call.name], call.args)  # try/except inside
            messages.append(tool_result(call, result)) # observe -> feeds the next reason
    return "stopped: step budget exhausted"            # fail SAFE, never forever
Enter fullscreen mode Exit fullscreen mode

That for step in range(max_steps) line is the difference between an agent that gives up gracefully and one that racks up a bill at 3am. Bound everything.

Planning: ReAct vs plan-and-execute

ReAct decides one step at a time — great for exploratory tasks, but it can wander. For multi-step tasks with a knowable shape, plan-and-execute is often better: the model first writes a plan, then executes, re-planning if reality diverges.

When to use which: ReAct for open-ended "figure it out" tasks; plan-and-execute for multi-step tasks where up-front structure cuts cost and wandering.

Tools done right

An agent is only as good as its tools, and most agent failures are tool failures. A production tool is a typed function + a clear schema + strict validation + error handling that teaches the model what went wrong.

@tool
def get_campaign_kpis(campaign_id: str, date_from: str, date_to: str) -> dict:
    """Return KPIs (impressions, clicks, conversions, spend) for one campaign in a date
    range. Dates are ISO-8601 (YYYY-MM-DD). Use this to answer performance questions."""
    validate_iso_date(date_from); validate_iso_date(date_to)   # never trust model args
    require_tenant_scope(campaign_id)                          # AUTHZ: agent can't cross tenants
    return kpi_service.query(campaign_id, date_from, date_to)
Enter fullscreen mode Exit fullscreen mode

Five rules:

  • The description IS the prompt. The model chooses tools by their docstrings — write them like instructions.
  • Constrain the schema — enums, formats, required fields.
  • Validate every argument. Tool code is a trust boundary, exactly like an API endpoint.
  • Errors that teach. Return "error: date_from is after date_to", not a stack trace — the model can recover.
  • Fewer, task-shaped tools. ~10 intent-level tools beat 40 CRUD endpoints.

Security note: a tool result is untrusted input. A retrieved document can contain injected instructions. Treat tool outputs as data, not commands.

Memory: scratchpad and long-term

  • Short-term (working memory) — the running message list / scratchpad for this task. Watch the context window; summarize or trim old steps on long runs.
  • Long-term — facts that persist across sessions (preferences, prior conclusions, a vector store), retrieved when relevant.
# LangGraph gives durable working memory + resume via a checkpointer.
agent = graph.compile(checkpointer=PostgresSaver(conn))   # state persists per thread_id
context = memory_store.search(query=task, tenant_id=tenant, k=5)  # long-term = retrieval
Enter fullscreen mode Exit fullscreen mode

Reflection and self-correction

Let the agent check its own work before returning — a reflection step asks "is this correct/complete?" and can trigger another loop. It catches hallucinations at the cost of extra calls, so gate it to high-stakes outputs.

answer = run_agent(task, tools)
critique = llm(f"Critique this answer for correctness and completeness:
{answer}")
if critique.has_issues:
    answer = run_agent(f"{task}
Fix these issues: {critique.issues}", tools, max_steps=4)
Enter fullscreen mode Exit fullscreen mode

Guardrails: the part that makes it production

Everything above makes the agent capable; guardrails make it safe to run unattended.

def guarded_step(state, call, tool):
    if state.steps > MAX_STEPS:      raise StopAgent("step budget")     # bound the loop
    if state.tokens > TOKEN_BUDGET:  raise StopAgent("token budget")    # bound the cost
    if tool.name not in ALLOWLIST:   return "error: tool not permitted" # least privilege
    if tool.is_destructive and not human_approved(call):               # confirm write actions
        return "paused: awaiting human approval"
    result = safe_invoke(tool, call.args)                              # never crash the loop
    return validate_output(result)                                     # schema/PII/guard checks
Enter fullscreen mode Exit fullscreen mode

The production checklist:

  • Step & token budgets — hard caps on loop length and spend, per run. Non-negotiable.
  • Tool allowlist / least privilege — the agent gets exactly the tools this task needs (an analytics agent has no delete tool).
  • Human-in-the-loop for destructive/irreversible actions — the agent proposes; a human commits.
  • Input & output validation — validate args going in; validate/scan the final answer going out (schema, PII, moderation).
  • Tool-result-as-untrusted — fence returned content, defending against indirect prompt injection.
  • Timeouts & fallbacks — every tool call has a timeout; a failed run degrades to a helpful message.
  • Full tracing — every run is observable.

Senior signal: you treat "the agent can take an irreversible action unattended" as a design red flag requiring a human gate.

Orchestration frameworks: which and when

  • LangGraph — models the agent as an explicit graph. Best for control, inspectability, durable state, and human-in-the-loop checkpoints.
  • CrewAI — high-level, role-based multi-agent with minimal code.
  • AutoGen — conversation-driven multi-agent.
  • Semantic Kernel / Microsoft.Extensions.AI — the .NET path for teams on the Microsoft stack.
# LangGraph: the agent IS a graph. Explicit, inspectable, resumable.
graph = StateGraph(AgentState)
graph.add_node("reason", call_model)
graph.add_node("act", run_tools)
graph.add_conditional_edges("reason", lambda s: "act" if s.tool_calls else END)
graph.add_edge("act", "reason")                    # observe -> reason again (the loop)
agent = graph.compile(checkpointer=memory)         # durable state + resume + HITL pauses
Enter fullscreen mode Exit fullscreen mode

In production: naive vs hardened

Metric Naive (v1) Hardened (v2)
Agentic p95 latency 4.2s 1.8s
Cost per question ~$0.90 (unbounded) ~$0.12 (budgeted + cached)
Runaway loops frequent zero (hard caps)
Debuggability none full trace per run

The wins came entirely from structure — bounded loop, typed read-only tools, budgets, output validation, and full tracing — not a smarter model.

What's next

Building the agent is step one. The three skills that make it trustworthy in production: evals (test something non-deterministic), observability (trace runs: cost, latency, why it did that), and cost engineering (prompt/semantic caching, model routing). Build it, measure it, watch it, make it cheap.


This post was originally published on PrepStack. The full original has all the code and diagrams.

Top comments (0)