DEV Community

Cover image for Building a Production AI Agent, End to End
Vipul Dhaigude
Vipul Dhaigude

Posted on

Building a Production AI Agent, End to End

Every AI agent demo calls a tool and prints the result. None of them show what the agent is actually reasoning over, what happens when a tool call needs a human to say yes first, when it needs to actually compute something instead of guessing, or when the token bill for describing a few dozen tools on every single turn quietly eats your margin. ai-agent-template is a LangGraph agent, built to handle exactly that: Bedrock Knowledge Base retrieval, called on demand rather than auto-injected on every turn, that gives it something real to ground an answer in — arguably the actual brain of the thing — tools sourced entirely from MCP instead of hand-registered one by one, human-in-the-loop approval that survives a server restart mid-decision, a real code sandbox for the math and charts a tool call can't do, an input guardrail that never takes the chat down even when it fails, and prompt caching that actually keeps the token bill in check instead of just claiming to. It runs on OpenAI by default — clone it, drop in an API key, no AWS account needed — with Bedrock as the recommended path once you actually need the production-grade pieces (real guardrails, persistent memory, CloudWatch observability) that only AWS can back.

What follows: the middleware chain that makes each of those real, a deliberate caching decision that looks wasteful until you check the numbers, three bugs that shipped and got fixed, and what a live run against Bedrock actually looks like.

Architecture at a glance

Architecture at a glance

The core idea: middleware, not a hand-rolled graph

Instead of baking every concern — memory, prompt caching, guardrails, context limits — directly into graph node logic, each one is a separate middleware, composed in order, testable on its own. Roughly, in the order they run: lifecycle (read memory before the call, write a turn summary after), dynamic context injection (memory and documents appended per turn, never written back into message history), prompt-cache settings, a context-overflow safety net that trims and retries once if a call comes back too large, an input guardrail that runs once per turn rather than once per model call, and an HITL policy paired with a risk classifier that decides which tool calls need a human.

Tools are bound fresh on every model call rather than wired into a static graph, and the system prompt itself is fetched from the companion MCP server at startup rather than hardcoded — the agent doesn't define its own tools, it sources them.

Execution Flow

What you get out of the box

  • HITL with stateless resume. Approve, reject, respond, or edit a proposed tool call. The interrupt happens at graph execution, not inside middleware — the turn ends, and a separate new turn resumes it from a checkpoint once the decision comes back. That's what makes approval survive a server restart instead of depending on a connection staying open.
  • Tools sourced entirely from MCP. No hand-registered tool list — it connects to a companion MCP server over HTTP, lists tools once at startup, and converts them to LangChain tools automatically.
  • Knowledge-base retrieval via Bedrock KB. kb_retrieve is a real tool, bound alongside the ones sourced from MCP, backed by Bedrock Knowledge Base's read-only Retrieve API — the model calls it on demand when it decides it needs outside knowledge, rather than every turn getting auto-injected context whether it needs it or not. Results come back deduped and cited by source. No KB ID configured, the tool simply never registers — the agent runs normally on its other tools either way. This is the part that lets the agent ground an answer in your own documents instead of guessing — arguably doing more of the actual reasoning work than any single tool. Bedrock-only regardless of which model provider is active.
  • Guardrails. A real ApplyGuardrail check runs before a turn completes — actual content/topic filtering, not a keyword list — and it's fail-open by design, so a guardrail outage degrades gracefully instead of taking the chat down.
  • Long-term memory across conversations. AgentCore-backed session memory that persists facts across turns and sessions, not just within one, with automatic fallback to in-memory state if the backing store is unreachable.
  • Observability. Structured logs with a turn correlation ID, plus normalized token and cache accounting wired into CloudWatch traces and metrics — not print statements you have to go dig through.
  • Prompt caching that's actually measured. Normalized token and cache accounting across every call, because a cache metric nobody can see is worse than no metric at all.
  • A real code sandbox for math and charts (run_python). A remote AgentCore sandbox the agent hands a task to when the answer needs actual computation — statistics, aggregation, a Plotly chart — instead of a tool call alone. Server owns the whole lifecycle (start, execute, register any produced file, stop) in one self-contained call; the model never sees a sandbox id. It has no access to this agent's own tools at all — on purpose, see below.
  • An orchestration tool for dependent tool calls (run_orchestration). A separate, local, sub-second sandbox for the one shape a single tool call can't handle: calls that depend on each other — chain, branch, reshape, loop — without spending a round trip per step. Read-only only; more on why below.
  • A WebSocket API, plus a health endpoint.

Everything past HITL and MCP tool sourcing is off by default — each one is a single env var away from turning on, and the agent completes a full turn cleanly with all of them unset.

Binding every tool, every turn — on purpose

The obvious optimization is to retrieve only the tools relevant to what the user just asked and bind those. It's also the wrong call once prompt caching is in the picture: a cached prefix is cheap to read but costs a premium to write, and that premium only pays off if the prefix stays stable turn to turn. Narrowing the tool set per turn changes the prefix every time, which busts the cache you were trying to use. So this template does the opposite on purpose — binds every tool, unconditionally, every turn. The block is bigger, but it's byte-identical call to call, so it sits in the cache: pay the write premium once, read cheap after that.

Three bugs I hit building this — already fixed here

  • A telemetry counter that read zero forever. Wired to a code path that got refactored out from under it, so it kept reporting a healthy zero instead of erroring — fix: assert the counter in a test against the current path, not just that it exists.
  • HITL state that didn't survive a restart. An approval loop that holds state in the running process loses everything if the process restarts mid-approval — fix: interrupt at graph execution, end the turn, resume as a fresh turn from a checkpoint.
  • A safety rule living only in the prompt. "Only auto-approve read-only tools" as a system-prompt instruction is a rule you're asking a probabilistic model to follow every turn — fix: a naming convention (get_* is read-only, everything else needs approval) enforced by the risk classifier in code, not requested in text.

What I actually verified

Structurally: each middleware and the gateway are tested in isolation with fake models, and integration tests run a full turn against a mock MCP server and a fake Bedrock client — no live model calls needed for day-to-day development. The two sandboxes get the same treatment: run_orchestration's tests confirm the worker-thread/main-loop handoff actually happens (not just that it's called), that a mutation gets refused with the approval message rather than silently skipped, and that a call budget cuts off a runaway loop; run_python's tests run against a fake CodeExecutionService and confirm the sandbox stops in every failure branch, not just the happy path — none of it touches real AWS or MCP. And end to end, for real: a live run against actual Bedrock completed a full turn cleanly — guardrail check applied, response streamed back token by token, turn finished with no manual intervention.

One prerequisite worth calling out if you're trying this yourself: on the default OpenAI provider, all you need is an API key — no AWS account required. Switch to Bedrock and real AWS credentials are required to reach it, and doubly so once AgentCore memory is configured — there's no offline or mocked path for the model call itself. A local AWS_PROFILE, a BEDROCK_API_KEY bearer token, or explicit AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY/AWS_SESSION_TOKEN all work; the agent fails to start without one of them once Bedrock is selected.

Why this shape, specifically

OpenAI by default, Bedrock for production. OpenAI is the default provider, zero-config beyond an API key, and create_model() is a plain if/else dispatch rather than a provider interface — same anti-overengineering call this template makes everywhere else. Bedrock is the recommended provider for production, because six things genuinely don't have an OpenAI equivalent: a real ApplyGuardrail call instead of a keyword filter, Bedrock Knowledge Base retrieval, AgentCore's persistent resumable memory, a real code-interpreter sandbox, cachePoint-tuned prompt caching, cross-region model access, and IAM-based credentials instead of an API key in an env file. The engineering effort went into making every one of those degrade gracefully instead of erroring when you're not on Bedrock: guardrails fall back to a real PII-redaction check built on LangChain's own detector functions (anonymize and continue, not block, matching the framework's own documented guidance), memory falls back to in-memory state automatically even if an AgentCore memory ID is still set in config, and knowledge-base retrieval simply doesn't register as a tool if it isn't configured. 87 tests cover those fallback paths specifically, not just the Bedrock happy path, with zero live calls to either provider anywhere in the suite.

Two sandboxes, kept deliberately separate. run_python and run_orchestration look similar from the outside — both let the model write code instead of reasoning in plain text — but they exist for opposite reasons and neither can do the other's job. run_python is a remote AgentCore sandbox with pandas/numpy/plotly and zero access to this agent's own tools, for math and charts. run_orchestration is a local, in-process sandbox with no math library at all, whose entire purpose is calling already-bound tools in a dependent chain without a round trip per step. Collapsing them into one "code tool" would mean either giving the math sandbox access to live tool calls it has no business making, or bolting pandas onto a tool-calling loop that doesn't need it — kept apart, each one stays small enough to reason about.

Read-only inside the tool-calling sandbox, on purpose. The obvious version of run_orchestration lets a script call any bound tool, mutations included — useful, since a chain/reshape/loop pattern often ends in a write. It's also a real human-in-the-loop bypass: a call made from inside a sandboxed script never reaches the same approval path a normal tool call does, so a mutation buried in a loop would execute without anyone reviewing it, even though the identical call made directly by the model would have stopped for a human first. The fix here is a hard rule, not a prompt instruction: the sandbox's tool bridge checks the same risk classifier the rest of the agent uses, and only read-only (get_*) tools are reachable from inside it. A plan that needs a mutation has to gather what it needs with run_orchestration first, then call the mutation directly, where a human can actually see it. Less capable than the unrestricted version — and the honest tradeoff for keeping the approval guarantee this template already makes elsewhere intact.

Middleware over graph nodes. Concerns like caching and guardrails are easy to add, remove, or reorder when they're independent middleware instead of logic threaded through the graph itself.

Ephemeral context, never persisted to history. Memory and documents get rendered into the prompt at call time and never written back into the message list — keeps the cached prefix stable and keeps message history from bloating with data that was only relevant for one turn.

Resilient fallback over hard failure. If AgentCore memory or persistence is unreachable, the agent degrades to in-memory state and logs a warning once, rather than taking the chat down.

What's next

Last of the three: the chat UI that talks to this agent over WebSocket — a React 19 app that renders streaming tokens, tool calls, and HITL approval prompts as they arrive.

The template is MIT-licensed and public: ai-agent-template. The companion MCP server it sources tools from is mcp-server-template.

Top comments (0)