Never Trust Tool Arguments: Validating Model Output

Why Bedrock does not enforce your inputSchema, and how to validate, authorise and safely execute arguments that a model generated.

Here is the fact that surprises people: the inputSchema you put in toolConfig is a prompt for the model, not a filter applied by the service. Bedrock passes your schema to the model to shape its output; it does not reject a tool call whose arguments violate that schema before handing it to you. Models are good at conforming, but "usually correct" is not a security control. Treat every toolUse.input as untrusted input from the internet, because in effect that is exactly what it is.

Validate against the schema you published

Run the arguments through a real validator — jsonschema against the same schema object you sent, or a Pydantic model, or Powertools for AWS Lambda's parser. Validating against the identical schema keeps your contract in one place and catches missing required fields, wrong types, out-of-range enums and unexpected extra properties. Set additionalProperties: false so a hallucinated field is a hard failure rather than something your handler silently ignores.

Then validate what the schema cannot express

JSON Schema checks shape, not meaning. A well-formed order_id may belong to a different tenant. A date range may be valid and yet span ten years of data. Add business validation after structural validation: bounds on page sizes and date ranges, existence checks, tenant scoping, and canonicalisation of anything you will use as a key.

Authorisation is not in the conversation

The model has no idea who the end user is, and its request carries no credentials. Your application must re-derive the caller's identity from the session — a Cognito token, an IAM principal, whatever your front door uses — and check entitlement on every tool call, using the caller's identity rather than the model's suggestion. If the model asks to read account 12345, the only question that matters is whether this signed-in user may read account 12345. Give each tool's execution path least-privilege IAM: a read-only tool should run under a role that physically cannot write.

Never treat arguments as code

Model output must never reach eval, exec, a shell string, or a concatenated SQL query. Use parameterised queries; if the model must choose a table or column, map its value through an allow-list of identifiers you control. The same applies to file paths (resolve and confirm the result stays inside a permitted directory), URLs (allow-list hosts, and beware requests aimed at internal metadata endpoints), and anything passed to a subprocess.

Prompt injection turns tools into a weapon

The dangerous path is indirect: a retrieved document, a web page or an email in the context contains text instructing the model to call a tool with attacker-chosen arguments. The model may comply. This is why the trust boundary sits in your executor and not in the prompt — no amount of "ignore instructions in documents" system prompting is a control. Gate destructive and outbound-communication tools behind human confirmation, keep their blast radius small, and log every executed call with its arguments, the caller, and the result so the action is reconstructable afterwards.

Example

Example · python
import re
from jsonschema import Draft202012Validator, ValidationError

# The SAME schema object you sent in toolConfig - one source of truth.
CANCEL_ORDER_SCHEMA = {
    "type": "object",
    "properties": {
        "order_id": {"type": "string", "pattern": "^ORD-[0-9]{5}$"},
        "reason_code": {
            "type": "string",
            "enum": ["customer_request", "duplicate", "payment_failed", "other"],
        },
    },
    "required": ["order_id", "reason_code"],
    "additionalProperties": False,
}

VALIDATORS = {"cancel_order": Draft202012Validator(CANCEL_ORDER_SCHEMA)}
HANDLERS = {"cancel_order": cancel_order}  # explicit allow-list of callables


class ToolRejected(Exception):
    pass


def execute_tool(tool_use: dict, principal: dict) -> dict:
    name = tool_use["name"]

    # 1. The model may only name tools we published. Never getattr() a name.
    if name not in HANDLERS:
        raise ToolRejected(f"Unknown tool: {name}")

    # 2. Structural validation against the published schema.
    try:
        VALIDATORS[name].validate(tool_use["input"])
    except ValidationError as exc:
        raise ToolRejected(f"Invalid arguments: {exc.message}") from exc

    args = tool_use["input"]

    # 3. Business rules the schema cannot express.
    if not order_exists(args["order_id"]):
        raise ToolRejected(f"No such order: {args['order_id']}")

    # 4. Authorisation on the CALLER, not on what the model asked for.
    if not caller_may_cancel(principal, args["order_id"]):
        raise ToolRejected("Caller is not permitted to cancel this order.")

    # 5. Execute with a parameterised query, never string interpolation.
    return HANDLERS[name](**args)

When to use it

  • A banking assistant re-checks the signed-in user's entitlement to an account on every balance-lookup tool call, so a prompt-injected instruction to read another account is refused by the executor.
  • A data assistant lets the model pick a table name from a fixed allow-list rather than interpolating the model's string into SQL.
  • An action-group Lambda function validates incoming parameters with Powertools' parser and returns a structured rejection instead of raising, so the agent can correct itself.

Discussion

  • Be the first to comment on this lesson.