+1 (726) 227-3241
NeuralArmada

Build a Tool-Using Agent on a SageMaker AI Endpoint with MCP

Every ML roadmap we have reviewed in the last two quarters has "agents" on it, and most of them stall at the same place: a demo notebook that calls a hosted API, with no answer for how the model gets deployed privately, how tools are governed, or who pays for the idle GPU. This tutorial builds the boring, production-shaped version on Amazon SageMaker AI: your own model on a SageMaker AI endpoint, tools exposed over the Model Context Protocol (MCP), and a small agent loop you can run in AWS Lambda or a container.

Nothing here requires a managed agent service. If you later move the orchestration to Bedrock AgentCore or a framework like Strands or LangGraph, the endpoint and MCP server you build below stay exactly the same — that is the point of the split.

The architecture in one paragraph

A caller (API Gateway, a Slack bot, a batch job) invokes an agent runtime. The runtime holds the conversation, decides when to call a tool, and calls the model on a SageMaker AI real-time endpoint. Tools live behind an MCP server — one process that publishes a typed catalogue of callable functions. The runtime never hard-codes tool logic; it lists tools from MCP, hands the schemas to the model, and executes whatever the model asks for, subject to an allowlist.

caller -> agent runtime (Lambda) -> SageMaker AI endpoint (model)
                     |
                     +--> MCP server (tools: SQL, search, internal APIs)

Step 1: deploy a tool-calling model with an OpenAI-compatible API

Agent loops are far simpler when the endpoint speaks the OpenAI chat-completions schema, because tool calls come back as structured JSON instead of text you have to parse. The Large Model Inference (LMI) container built on vLLM does this out of the box. Pick a model that was actually trained for tool use — Llama 3.x Instruct, Qwen2.5-Instruct, and Mistral-Small-Instruct all are.

import sagemaker
from sagemaker.djl_inference import DJLModel

role = sagemaker.get_execution_role()

model = DJLModel(
    model_id="meta-llama/Llama-3.1-8B-Instruct",
    role=role,
    env={
        "OPTION_ROLLING_BATCH": "vllm",
        "OPTION_MAX_MODEL_LEN": "16384",
        "OPTION_TENSOR_PARALLEL_DEGREE": "1",
        "OPTION_ENABLE_AUTO_TOOL_CHOICE": "true",
        "OPTION_TOOL_CALL_PARSER": "llama3_json",
        "HF_TOKEN": "<from Secrets Manager, not a literal>",
    },
)

predictor = model.deploy(
    initial_instance_count=1,
    instance_type="ml.g6.2xlarge",
    endpoint_name="agent-llm",
    container_startup_health_check_timeout=900,
)

Two settings do the heavy lifting. OPTION_ENABLE_AUTO_TOOL_CHOICE lets the server decide when a tool call is warranted, and OPTION_TOOL_CALL_PARSER tells it how that model family emits calls (llama3_json for Llama 3.x, hermes for most Qwen builds, mistral for Mistral). Mismatch the parser and you will get tool calls as plain text in content — the single most common failure we get called in to debug.

Sanity-check the schema before you build anything on top of it:

import boto3, json

rt = boto3.client("sagemaker-runtime")
body = {
    "messages": [{"role": "user", "content": "What is the order status for A-119?"}],
    "tools": [{
        "type": "function",
        "function": {
            "name": "get_order_status",
            "description": "Look up the fulfilment status of an order.",
            "parameters": {
                "type": "object",
                "properties": {"order_id": {"type": "string"}},
                "required": ["order_id"],
            },
        },
    }],
    "tool_choice": "auto",
    "temperature": 0.1,
}
resp = rt.invoke_endpoint(
    EndpointName="agent-llm",
    ContentType="application/json",
    Body=json.dumps(body),
)
print(json.loads(resp["Body"].read())["choices"][0]["message"])

You want to see a tool_calls array with a JSON arguments string. If you do, the rest is plumbing.

Step 2: publish your tools over MCP

MCP matters here for an unglamorous reason: it decouples the tool catalogue from the agent. One MCP server can back a Slack agent, an internal web app, and a nightly batch job, and a platform team can audit and version it in one place.

# mcp_server.py
from mcp.server.fastmcp import FastMCP
import boto3

mcp = FastMCP("ops-tools")
ddb = boto3.resource("dynamodb").Table("orders")

@mcp.tool()
def get_order_status(order_id: str) -> dict:
    """Look up the fulfilment status of an order by ID."""
    item = ddb.get_item(Key={"order_id": order_id}).get("Item")
    if not item:
        return {"found": False}
    return {"found": True, "status": item["status"], "eta": item.get("eta")}

@mcp.tool()
def search_kb(query: str, top_k: int = 5) -> list[dict]:
    """Semantic search over the internal knowledge base."""
    ...

if __name__ == "__main__":
    mcp.run(transport="streamable-http")

The docstrings and type hints become the tool schema the model sees, so write them for the model, not for your teammates: say what the tool returns and when not to call it. Run the server as an ECS/Fargate service behind an internal ALB, or in-process next to the agent for a first iteration.

Step 3: the agent loop

The loop is about forty lines. Resist the urge to make it clever.

import json, boto3
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client

rt = boto3.client("sagemaker-runtime")
ENDPOINT = "agent-llm"
MAX_STEPS = 6
ALLOWED = {"get_order_status", "search_kb"}

def call_model(messages, tools):
    resp = rt.invoke_endpoint(
        EndpointName=ENDPOINT,
        ContentType="application/json",
        Body=json.dumps({"messages": messages, "tools": tools,
                         "tool_choice": "auto", "temperature": 0.1}),
    )
    return json.loads(resp["Body"].read())["choices"][0]["message"]

async def run(user_text, mcp_url):
    async with streamablehttp_client(mcp_url) as (r, w, _):
        async with ClientSession(r, w) as session:
            await session.initialize()
            listed = await session.list_tools()
            tools = [{"type": "function", "function": {
                        "name": t.name,
                        "description": t.description,
                        "parameters": t.inputSchema}}
                     for t in listed.tools if t.name in ALLOWED]

            messages = [
                {"role": "system", "content": "Use tools for anything factual about orders or docs. Never guess an order status."},
                {"role": "user", "content": user_text},
            ]

            for _ in range(MAX_STEPS):
                msg = call_model(messages, tools)
                messages.append(msg)
                calls = msg.get("tool_calls") or []
                if not calls:
                    return msg.get("content", "")
                for c in calls:
                    name = c["function"]["name"]
                    args = json.loads(c["function"]["arguments"] or "{}")
                    if name not in ALLOWED:
                        result = {"error": "tool not permitted"}
                    else:
                        result = (await session.call_tool(name, args)).content
                    messages.append({"role": "tool",
                                     "tool_call_id": c["id"],
                                     "content": json.dumps(result, default=str)})
            return "Step budget exhausted; escalating to a human."

Four details that keep this out of trouble:

  • MAX_STEPS is a hard budget. Unbounded loops are how a $40/day agent becomes a $4,000 one overnight.
  • ALLOWED is enforced in the runtime, not in the prompt. The model proposes; your code disposes.
  • Tool credentials belong to the MCP server, never to the model or the prompt. The agent runtime's IAM role should be able to invoke the endpoint and reach the MCP server, and nothing else.
  • Tool results are untrusted input. A record containing "ignore previous instructions" is a prompt injection vector; keep tool output in role: tool messages and never re-inject it as a system prompt.

Step 4: guardrails, tracing, and the bill

Guardrails. Run input and output through Amazon Bedrock Guardrails via the standalone ApplyGuardrail API — it works with a SageMaker-hosted model, no Bedrock inference required. Screen the user turn before the loop, and the final answer before it leaves.

Tracing. Log one structured record per step: request ID, step index, tool name, arguments hash, latency, token counts. Emit them as CloudWatch EMF so you get metrics for free, and enable data capture on the endpoint so you have real traffic to replay when you swap models. Without per-step traces, "the agent did something weird" is unfalsifiable.

Cost. An always-on ml.g6.2xlarge is roughly $700–900/month before autoscaling. Three levers, in order of impact: put the endpoint behind inference components so several models share a GPU and can scale to zero when idle; cap MAX_STEPS and max_tokens; and cache tool results that are stable within a session. For bursty internal traffic, an async or serverless endpoint for the summarisation legs plus one small real-time endpoint for the loop is usually cheaper than a single large endpoint.

Evaluation. Trajectory quality is not the same as answer quality. Keep a fixture set of 30–50 requests with the expected tool sequence, and assert on it in CI — a model upgrade that improves prose while dropping a lookup step is a regression you will otherwise ship.

When to use this versus a managed agent service

Build this stack when you need a specific or fine-tuned model, VPC-only networking, or tools that already live behind private APIs. Reach for a managed runtime when your model can come from a hosted catalogue and orchestration, memory, and identity are the parts you would rather not own. The MCP server is portable either way, which is the strongest argument for starting there.

If you want a second pair of eyes on an agent design — or a senior SageMaker engineer to build the endpoint, MCP layer, and evaluation harness with your team — get in touch. We work as an extension of in-house teams on exactly this kind of build.