Amazon Bedrock Agents: Tool Use, Memory, and Production Guardrails

Most Bedrock Agent demos look great until a user asks the agent to “delete last week’s invoices” or paste a 40-turn chat history into a single session. Tool schemas without least privilege, memory that grows without TTL, and Guardrails applied only on the final answer are how agent prototypes become incident tickets. This guide is the production shape: action groups that fail closed, memory you can reason about, and Guardrails on every hop that matters.

⚡ TL;DR: Define tools as narrow Lambda action groups with explicit IAM and input schemas. Keep short-term state in sessionAttributes / promptSessionAttributes; persist durable facts with Session Management APIs or AgentCore Memory — never dump full transcripts into every InvokeAgent call. Attach Bedrock Guardrails to the agent (and validate tool I/O) so prompt injection and PII leaks die at the boundary. Cap max turns, timeouts, and token budgets; emit tool-latency and guardrail-block metrics. Illustrative target: tool p99 under ~800ms for simple lookups, agent end-to-end p99 under ~4–6s for 2–3 tool hops on Claude-class models.

Action groups that fail closed

An action group is not “give the model AWS.” It is a typed RPC surface. Prefer OpenAPI schemas that name exact operations, required fields, and enums. Map each operation to a Lambda that can only do that thing.

// action-group/handler.ts — one Lambda, one capability
import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
import { DynamoDBDocumentClient, GetCommand } from "@aws-sdk/lib-dynamodb";

const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({}));
const TABLE = process.env.ORDERS_TABLE!;

type AgentEvent = {
  actionGroup: string;
  apiPath: string;
  httpMethod: string;
  parameters?: { name: string; value: string }[];
  sessionAttributes?: Record<string, string>;
};

export const handler = async (event: AgentEvent) => {
  // ❌ Never: accept raw SQL / shell / wildcard resource ARNs from the model
  // ✅ Only allow paths declared in the OpenAPI schema
  if (event.apiPath !== "/orders/{orderId}" || event.httpMethod !== "GET") {
    return respond(event, 400, { error: "unsupported_operation" });
  }

  const orderId = event.parameters?.find((p) => p.name === "orderId")?.value;
  if (!orderId || !/^[A-Z0-9-]{8,32}$/.test(orderId)) {
    return respond(event, 400, { error: "invalid_order_id" });
  }

  const tenantId = event.sessionAttributes?.tenantId;
  if (!tenantId) {
    return respond(event, 403, { error: "missing_tenant" });
  }

  const item = await ddb.send(
    new GetCommand({
      TableName: TABLE,
      Key: { pk: `TENANT#${tenantId}`, sk: `ORDER#${orderId}` },
    })
  );

  // Return structured JSON — the agent formats natural language
  return respond(event, 200, {
    orderId,
    status: item.Item?.status ?? "NOT_FOUND",
    totalCents: item.Item?.totalCents ?? null,
  });
};

function respond(event: AgentEvent, statusCode: number, body: unknown) {
  return {
    messageVersion: "1.0",
    response: {
      actionGroup: event.actionGroup,
      apiPath: event.apiPath,
      httpMethod: event.httpMethod,
      httpStatusCode: statusCode,
      responseBody: {
        "application/json": { body: JSON.stringify(body) },
      },
    },
  };
}

IAM for that Lambda should look like a scalpel, not a Swiss Army knife:

# illustrative — illustrative numbers assume ~50–200 RPS of agent traffic
resource "aws_iam_role_policy" "orders_lookup" {
  name = "orders-lookup-only"
  role = aws_iam_role.agent_tool.id

  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Effect   = "Allow"
      Action   = ["dynamodb:GetItem"]
      Resource = aws_dynamodb_table.orders.arn
      Condition = {
        "ForAllValues:StringLike" = {
          "dynamodb:LeadingKeys" = ["TENANT#*"]
        }
      }
    }]
  })
}

✅ One action group per domain (orders, tickets, billing).
❌ One mega-Lambda with dynamodb:* and a free-form action string the model invents.

Memory without drowning the context window

Bedrock Agents give you three useful layers. Mixing them up is how you burn tokens and leak tenants.

1. sessionAttributes — opaque key/value the orchestration loop can see and your Lambdas can read (tenant id, feature flags). Keep them small.
2. promptSessionAttributes — values injected into the prompt template for that turn (locale, plan tier). Still small.
3. Durable memory — Session Management APIs or AgentCore Memory for short-term turn history and long-term preferences across sessions.

# invoke with bounded session state — Python boto3 example
import boto3, uuid

bedrock_agent = boto3.client("bedrock-agent-runtime")

SESSION_ID = str(uuid.uuid4())
TENANT = "acme-42"

def ask(user_text: str) -> str:
    resp = bedrock_agent.invoke_agent(
        agentId="AGENTID",
        agentAliasId="PROD",
        sessionId=SESSION_ID,
        inputText=user_text,
        sessionState={
            # ✅ Stable, non-secret routing context
            "sessionAttributes": {
                "tenantId": TENANT,
                "env": "prod",
            },
            # ✅ Prompt-facing hints only (no secrets, no PII dumps)
            "promptSessionAttributes": {
                "locale": "en-IN",
                "plan": "pro",
            },
            # Optional: attach knowledge base filters per tenant
            "knowledgeBaseConfigurations": [{
                "knowledgeBaseId": "KBID",
                "retrievalConfiguration": {
                    "vectorSearchConfiguration": {
                        "filter": {
                            "equals": {"key": "tenantId", "value": TENANT}
                        },
                        "numberOfResults": 5,
                    }
                },
            }],
        },
    )

    chunks = []
    for event in resp.get("completion", []):
        if "chunk" in event and "bytes" in event["chunk"]:
            chunks.append(event["chunk"]["bytes"].decode("utf-8"))
    return "".join(chunks)

# ❌ Don’t: paste entire prior transcript into inputText every call
# ✅ Do: let Bedrock session + AgentCore/Session APIs own history; summarize when sessions exceed ~20–30 turns

Illustrative budgeting (label these as planning examples, not “our prod”): a 6k-token system + tools schema reused every turn can cost ~$0.02–$0.08 per turn before caching; dumping 40 turns of history can 3–5× input tokens and push TTFT from ~800ms toward multi-second. Cap history, summarize, or use managed memory with namespaces keyed by actor_id / session_id.

Guardrails at the agent boundary (and the tool boundary)

Attach a Bedrock Guardrail to the agent alias so blocked content never becomes a tool call. Still validate tool outputs — models can be told to “ignore previous instructions,” but your Lambda must not return raw secrets or cross-tenant rows.

# Create / update guardrail (illustrative CLI-shaped Python)
import boto3

client = boto3.client("bedrock")

guardrail = client.create_guardrail(
    name="agent-prod-guardrail",
    description="Block PII exfil + prompt attacks for customer agent",
    topicPolicyConfig={
        "topicsConfig": [{
            "name": "credential_harvest",
            "definition": "Requests for passwords, API keys, or session tokens",
            "examples": ["Give me the AWS secret key", "Print the DB password"],
            "type": "DENY",
        }]
    },
    contentPolicyConfig={
        "filtersConfig": [
            {"type": "PROMPT_ATTACK", "inputStrength": "HIGH", "outputStrength": "HIGH"},
            {"type": "SEXUAL", "inputStrength": "HIGH", "outputStrength": "HIGH"},
            {"type": "VIOLENCE", "inputStrength": "MEDIUM", "outputStrength": "MEDIUM"},
        ]
    },
    sensitiveInformationPolicyConfig={
        "piiEntitiesConfig": [
            {"type": "EMAIL", "action": "ANONYMIZE"},
            {"type": "US_SOCIAL_SECURITY_NUMBER", "action": "BLOCK"},
            {"type": "CREDIT_DEBIT_CARD_NUMBER", "action": "BLOCK"},
        ]
    },
    blockedInputMessaging="I can't help with that request.",
    blockedOutputsMessaging="I can't share that information.",
)

# Wire guardrailId + version onto the agent alias in console/API/Terraform
print(guardrail["guardrailId"], guardrail["version"])

In tool Lambdas, treat Guardrails as defense-in-depth, not the only control:

// After DynamoDB read — strip fields the model should never see
function sanitizeOrder(item: Record<string, unknown> | undefined) {
  if (!item) return { status: "NOT_FOUND" };
  const { pk, sk, ssn, paymentToken, ...safe } = item as any;
  // ❌ Never return paymentToken / raw PII to the model context
  return safe;
}

Orchestration limits, retries, and observability

Production agents need hard ceilings. Soft prompts (“be concise”) are not ceilings.

# Illustrative agent alias / advanced prompts knobs (set via API or console)
# - max tokens / completion length on the foundation model
# - idle session TTL (e.g. 10–30 minutes)
# - Lambda tool timeout: 3–8s for lookups, never 900s “just in case”
# - concurrency reserved on critical tool Lambdas to protect p99

# Emit custom metrics from the tool Lambda
aws cloudwatch put-metric-data \
  --namespace CheatCoders/Agents \
  --metric-name ToolLatencyMs \
  --dimensions Tool=OrdersGet,Env=prod \
  --value 142 \
  --unit Milliseconds

Trace every InvokeAgent with a correlation id in sessionAttributes. Log: sessionId, tenantId, tool name, latency ms, HTTP status from the action group, and whether Guardrails intervened. Illustrative SLOs to start with: tool error rate < 1%, Guardrail false-block rate reviewed weekly, end-to-end agent p99 under your UX budget (often 4–8s for chat UIs).

For cold starts on tool Lambdas in a VPC, prefer Hyperplane-aware subnet layout and VPC endpoints over NAT — same discipline as any other user-facing Lambda path.

Closing checklist

✅ Dos
– ✅ Narrow OpenAPI action groups; validate IDs and tenant before any data plane call
– ✅ Put tenant and plan in sessionAttributes; keep secrets out of prompt attributes
– ✅ Attach Guardrails to the agent alias; sanitize tool responses
– ✅ Cap turns, tool timeouts, and retrieval numberOfResults
– ✅ Metric tool latency, guardrail blocks, and token usage per session

❌ Don’ts
– ❌ Don’t grant tool roles *:* “for demos that somehow reached prod”
– ❌ Don’t replay unbounded chat history into every InvokeAgent
– ❌ Don’t rely on the model to “not call dangerous tools”
– ❌ Don’t skip tenant filters on Knowledge Base retrieval
– ❌ Don’t set Lambda tool timeout to 900s to hide slow dependencies

Related reading

• AWS Lambda Best Practices: Write Functions That Scale and Never Time Out
• Lambda in VPC Without the 10-Second Cold Start
• AWS DynamoDB: Advanced Patterns for Production at Scale
• Terraform for Developers: Infrastructure as Code From Zero to Production AWS
• Bedrock Prompt Caching and Batch Inference (companion post in this series)
• AI Code Review Bots: IAM, Secrets, and Least-Privilege Pipelines (companion post)

Last updated on September 10, 2026