How to add an AI agent to an Axum app

What an agent in an Axum app is made of

• A server endpoint: an async handler on your Router, such as POST /api/agent, that calls the model. The agent lives in router state, and the model API key is read from the server environment, never rendered into a page or a static script.
• Tools: Rust functions the model may call, each wrapping an operation your app already has (an sqlx query, a service call, an internal API).
• The agent loop: the model calls a tool, receives the result and decides the next step until it can answer. Cap the number of turns.
• A chat UI: your server-rendered page plus a small script that posts the message and reads the reply as server-sent events, the stream Axum's Sse response produces.
• Authorization and approval: an extractor turns your session into the signed-in user, tools act only as that user, and changes to data wait for the user to confirm them.

Ways to add an agent to Axum

Build an agent with Rig

Rig is a Rust library for LLM applications: one API over many model providers, typed tools and an agent loop that runs tool calls and streams every step. The examples below use Rig 0.46 (which needs Rust 1.97.1 or newer) with Axum 0.8 on Tokio. Add the crates with:

cargo add [email protected] [email protected] futures serde serde_json tokio --features serde/derive,tokio/full

The examples call Claude through Rig's Anthropic provider, which reads the API key from the ANTHROPIC_API_KEY variable. Other providers, such as OpenAI or OpenRouter, only change the client and model lines.

1. Tools scoped to the user

The #[rig::rig_tool] macro turns an async function into a tool: the doc comment becomes the description, and the arguments become the JSON Schema the model fills in. The signed-in user's ID is not an argument. It travels in Rig's tool context, which the handler fills for each run and the model cannot see or set, so the model can choose an order ID but never whose orders it reads. Your data layer still checks that the order belongs to that user.

cancel_order does not cancel anything. It checks that the order is open and returns a proposal that the page asks the user to confirm. Rig can also gate tool calls in an on_dispatch hook or pause a durable run until someone approves, but a hook waits inside the running request, and a durable run needs storage for the paused state. A proposal keeps the write on the endpoint your app already has.

use rig::prelude::*;
use rig::providers::anthropic::{self, Anthropic};
use rig::tool::ToolExecutionError;
use serde::{Deserialize, Serialize};

use crate::orders; // your existing data layer

/// The signed-in user's ID. The handler sets it for each run; the model never does.
#[derive(Serialize, Deserialize, rig::ContextValue)]
pub struct UserId(pub String);

/// List the signed-in user's recent orders.
#[rig::rig_tool(params(status = "Only orders in this state: open, shipped or cancelled"))]
async fn list_orders(
    ctx: &mut rig::tool::ToolContext,
    status: Option<String>,
) -> Result<Vec<orders::Order>, ToolExecutionError> {
    let user = ctx.require::<UserId>()?;
    orders::list(&user.0, status.as_deref())
        .await
        .map_err(|_| ToolExecutionError::other("could not load orders"))
}

#[derive(Serialize)]
pub struct PendingCancel {
    order_id: String,
    needs_confirmation: bool,
}

/// Ask to cancel one of the signed-in user's open orders. The user confirms it in the page.
#[rig::rig_tool(params(order_id = "The order to cancel"))]
async fn cancel_order(
    ctx: &mut rig::tool::ToolContext,
    order_id: String,
) -> Result<PendingCancel, ToolExecutionError> {
    let user = ctx.require::<UserId>()?;
    match orders::find_open(&user.0, &order_id).await {
        Ok(Some(_)) => Ok(PendingCancel { order_id, needs_confirmation: true }),
        _ => Err(ToolExecutionError::not_found("no open order with that id")),
    }
}

pub fn build_agent() -> Result<Agent, Box<dyn std::error::Error>> {
    // Reads ANTHROPIC_API_KEY from the environment.
    let model = Anthropic::from_env()?.completion(anthropic::CLAUDE_SONNET_5_5);
    Ok(AgentBuilder::new(model)
        .preamble("You help the signed-in user manage their orders. Use the tools; never guess order data.")
        .tool(ListOrders)
        .tool(CancelOrder)
        .build())
}

2. The SSE handler

CurrentUser implements Axum's FromRequestParts on top of your session, so any handler that takes it returns 401 to signed-out requests before the model is called. The handler puts the user ID into the tool context, caps the run with max_turns(5) and calls stream(), which yields the model's text fragments, tool calls and tool results as they happen. It forwards text as text events and a successful cancel_order result as a confirm event through Axum's Sse response. Each event's data is JSON, so newlines in the model's text cannot break the event format.

use std::convert::Infallible;

use axum::{
    Json,
    extract::{FromRequestParts, State},
    http::{StatusCode, request::Parts},
    response::sse::{Event, KeepAlive, Sse},
};
use futures::{Stream, StreamExt};
use rig::agent::{Agent, MultiTurnStreamItem};
use rig::streaming::{Item, StreamEvent};
use rig::tool::ToolContext;
use serde::Deserialize;

use crate::auth; // your existing session lookup
use crate::tools::UserId;

/// The signed-in user. A handler that takes it rejects signed-out requests with 401.
pub struct CurrentUser {
    pub id: String,
}

impl<S: Send + Sync> FromRequestParts<S> for CurrentUser {
    type Rejection = StatusCode;

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, StatusCode> {
        let id = auth::session_user_id(parts).await.ok_or(StatusCode::UNAUTHORIZED)?;
        Ok(CurrentUser { id })
    }
}

#[derive(Deserialize)]
pub struct ChatRequest {
    message: String,
}

pub async fn chat(
    user: CurrentUser,
    State(agent): State<Agent>,
    Json(req): Json<ChatRequest>,
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, StatusCode> {
    let mut context = ToolContext::new();
    context.insert(UserId(user.id)).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    let stream = agent
        .prompt(req.message)
        .tool_context(context)
        .max_turns(5)
        .stream()
        .filter_map(|item| async move {
            let event = match item {
                Ok(MultiTurnStreamItem::StreamAssistantItem(Item::Event(StreamEvent::Text { text, .. }))) => {
                    Event::default().event("text").json_data(text)
                }
                Ok(MultiTurnStreamItem::ToolResult { tool_result })
                    if tool_result.name == "cancel_order" && !tool_result.is_error =>
                {
                    Event::default().event("confirm").json_data(tool_result.content.first()?.as_json()?)
                }
                Err(_) => Event::default().event("error").json_data("The agent stopped. Try again."),
                _ => return None,
            };
            event.ok().map(Ok)
        });

    Ok(Sse::new(stream).keep_alive(KeepAlive::default()))
}

Rig's Agent is Clone, so it can be the router state directly:

mod agent;
mod auth;
mod orders;
mod tools;

use axum::{Router, routing::post};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let app = Router::new()
        .route("/api/agent", post(agent::chat))
        // ...your existing routes and session layer
        .with_state(tools::build_agent()?);

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await?;
    axum::serve(listener, app).await?;
    Ok(())
}

This handler answers one message at a time. To keep a conversation, store the messages from the final response per user on the server and pass them to the next run with .history(..); see Rig's streaming guide.

3. The browser client

The browser's EventSource only sends GET requests, so the client posts the message with fetch and reads the events from the response body. On a confirm event it shows a button; the click calls your existing cancel endpoint with the same session cookie and CSRF token as any other form on the page.

const log = document.querySelector("#agent-log");
const form = document.querySelector("#agent-form");
const csrfToken = document.querySelector('meta[name="csrf-token"]').content;

form.addEventListener("submit", async (event) => {
  event.preventDefault();
  const message = form.elements.message.value;
  form.reset();
  const reply = log.appendChild(document.createElement("p"));

  const res = await fetch("/api/agent", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ message }),
  });
  if (!res.ok) {
    reply.textContent = "Error " + res.status;
    return;
  }

  // Read the server-sent events from the response body.
  const reader = res.body.pipeThrough(new TextDecoderStream()).getReader();
  let buffer = "";
  for (;;) {
    const { value, done } = await reader.read();
    if (done) break;
    buffer += value;
    const frames = buffer.split("\n\n");
    buffer = frames.pop();
    for (const frame of frames) {
      const name = /^event: (.*)$/m.exec(frame)?.[1];
      const data = /^data: (.*)$/m.exec(frame)?.[1];
      if (!data) continue;
      if (name === "text") reply.textContent += JSON.parse(data);
      if (name === "error") reply.textContent = JSON.parse(data);
      if (name === "confirm") confirmCancel(JSON.parse(data).order_id);
    }
  }
});

// The agent only proposes the cancellation; the user's click calls your normal endpoint.
function confirmCancel(orderId) {
  const button = log.appendChild(document.createElement("button"));
  button.textContent = "Cancel order " + orderId;
  button.addEventListener("click", async () => {
    button.disabled = true;
    const res = await fetch("/orders/" + encodeURIComponent(orderId) + "/cancel", {
      method: "POST",
      headers: { "X-CSRF-Token": csrfToken },
    });
    button.textContent = res.ok ? "Order " + orderId + " cancelled" : "Could not cancel (" + res.status + ")";
  });
}

Production checklist for Axum

• Authorize in the extractor and in every tool. Taking CurrentUser in the handler signature means the route cannot run without a user, and every query a tool makes is scoped to the ID from the tool context, never to an ID the model supplies.
• Protect writes against CSRF. Axum has no built-in CSRF protection. The Json extractor rejects requests without Content-Type: application/json, which a cross-site form cannot send and a cross-site fetch only sends after a CORS preflight, so do not add a CORS layer that allows credentials from other origins. The cancel endpoint keeps whatever CSRF check your app already uses.
• Treat tool input and tool results as untrusted. The model chooses tool arguments, and data a tool returns (user-written text, emails, web pages) can contain instructions. Validate arguments with your own business rules, and keep anything irreversible behind the user's confirmation.
• Limit cost. Cap turns, rate-limit the route per user (in your own middleware or with a tower layer such as tower_governor), and set spending limits with your model provider.
• Keep the runtime responsive. Tools run on Tokio's worker threads: use async database drivers, or move blocking calls into tokio::task::spawn_blocking.
• Allow for long streams. KeepAlive::default() sends a keep-alive message every 15 seconds while the model works, so idle timeouts do not close the connection. Make sure your reverse proxy does not buffer text/event-stream responses and allows reads as long as several model turns take.
• Keep secrets server-side. ANTHROPIC_API_KEY belongs in the server environment or a secret manager, never in templates or static files.

Option: embed an in-app agent with Bunkhouse

If the agent's job is to help signed-in users do work your Axum app already supports, you may not need to build the tools, agent loop and chat UI yourself. Bunkhouse adds an Assistant pane to an existing web app with one script tag. When a user opens it, the assistant works with what the page already offers: buttons, links, form fields and tables, the API requests your app makes, a hints file you publish for operations it cannot see, and WebMCP tools your page declares.

The model writes small programs that run in a sandbox on the embed origin, not in your page. They can only call capabilities the assistant has discovered and your configuration allows. Write actions are shown to the user for review before they run, the model never sees credentials, and your Axum backend keeps enforcing its own authorization on every request.

In Axum, put the identity hook and the loader in the <head> of the page template your handlers render, before your application's own scripts, so they load once per document. For signed-in users, add a GET /api/agent-identity handler that takes your user extractor (an OptionalFromRequestParts implementation over your session) and signs a short-lived token for that user with jsonwebtoken. Merge its router before your session layer is added, so the layer wraps it; the signing secret stays in a server environment variable.

<head>
  <!-- The identity hook first, then the loader, before your application's own scripts -->
  <script src="/agent-identity.js"><\/script>
  <script src="https://embed.bunkhouse.ai/embed.js" data-public-key="pk_YOUR_SITE_PUBLIC_KEY"><\/script>
</head>

The tested Axum integration guide has the complete code for each session mode (signed identity token, server mint and anonymous), CSRF headers and Content Security Policy settings.

When Bunkhouse fits, and when to build it yourself

• Bunkhouse fits when users want tasks done inside your existing pages, such as filtering, editing or submitting through the controls and API you already have, and you would rather not write and maintain a tool definition for each action, a sandbox for model-written code, an approval UI and usage metering.
• Build it yourself when the agent must run without a browser page (background jobs, email, chat platforms, scheduled work), when you need to choose the model provider, use your own keys or run inference on your own infrastructure, or when the conversation must be part of your own UI components.
• Plan for the limits: discovery cannot see server-only routes, workflows the user has not visited, closed shadow roots or cross-origin frames (describe those in a hints file); some controls need a real user gesture; password, payment and file fields are never read, and you can mark other content private.
• Cost: production use is $49 per month, including $20 of usage credit per paid billing period, then $1 per million input tokens and $3 per million output tokens. See billing.

Try it on your Axum app first

A coding agent such as Claude Code, Codex or Cursor can run a free, disposable agent trial on a dummy page of your Axum app. No account, card or model API key is needed: your coding agent acts as the assistant's model. A trial lasts 24 hours and allows 100 turns. Give your agent this instruction, changing the address if your app runs elsewhere:

Read https://api.bunkhouse.ai/agent-trial.md and follow it to run a Bunkhouse agent trial on my dummy Axum page at http://localhost:3000.

The trial evaluation checklist shows how to test a read, a rejected write and an approved change, and what a trial does not prove about production use.

Related guides

• How to add AI to your website: chat assistants, in-app agents and custom builds compared
• MCP vs WebMCP: choosing where outside agents act on your site
• AI agents in other Rust frameworks: Actix Web