Add a web search tool to the Vercel AI SDK with Mux

Q: Why not use a provider's SDK directly as the tool?

You can, and at very high volume with a negotiated contract a direct relationship is cheaper. Mux is for the case where you want more than one index behind the same tool: models gives you fallback when a provider is slow or down, the response shape is the same for all three, and billing, rate limits and per-key budgets are in one place.

Q: Does the AI SDK's step limit stop runaway search costs?

Only indirectly. isStepCount(5) caps the number of model calls, and the model may call the tool several times within one step. The per-key credit limit on the Mux side is the hard stop, and it is enforced atomically before the search runs.

Q: Can I use this tool with OpenAI or Google models?

Yes. The tool is model-agnostic. Install the matching provider package, for example @ai-sdk/openai, and pass that provider's model to generateText. Nothing in web-search.ts changes.

Q: How do I keep the model from flooding its context with results?

Keep maxResults small by default (5 is plenty for most questions), return only title, URL, snippet and date as the example does, and reserve content: "full" for a separate tool the model calls on one URL at a time. Get a key and run your first search in the quickstart.

A web search tool for the Vercel AI SDK is a tool() definition whose execute function calls a search API and returns a small list of results the model can read. This guide wires that tool to Mux, so one MUX_API_KEY gives the agent Exa, Parallel and Keenable with automatic fallback, and shows how to run it in a multi-step generateText or streamText loop. The code targets AI SDK 7 (ai 7.0.x) with Zod 4.

What you need#

Node 18 or newer.
A Mux API key from the dashboard with a few dollars of credit. See the quickstart.
A key for whichever model provider you use with the AI SDK. The example below uses the Anthropic provider and reads ANTHROPIC_API_KEY.

npm install ai @ai-sdk/anthropic zod getmux
export MUX_API_KEY="mux-..."
export ANTHROPIC_API_KEY="sk-ant-..."

The versions this guide was written against:

Package	Version
`ai`	7.0.143
`@ai-sdk/anthropic`	4.0.82
`zod`	4.6.7
`getmux`	0.1.0

AI SDK 7 renamed a few things from AI SDK 6. Tools take inputSchema (not parameters), the step limit is stopWhen: isStepCount(n) (stepCountIs still works as a deprecated alias), and the per-step callback is onStepEnd. If you are on AI SDK 5 or 6, swap those names and the rest is identical.

The tool#

Create web-search.ts. The schema exposes query, an optional maxResults and an optional includeDomains. The execute function calls Mux with an ordered fallback list and trims the response to four fields per result, which keeps tool output around a few hundred tokens instead of several thousand.

// web-search.ts
import { tool } from "ai";
import { z } from "zod";
import { Mux, MuxError } from "getmux";

// Reads MUX_API_KEY from the environment.
const mux = new Mux();

export const webSearch = tool({
  description:
    "Search the web for current information. Returns titles, URLs, short snippets and publish dates. Use it for anything after your training data or that needs a source.",
  inputSchema: z.object({
    query: z.string().min(1).max(2000).describe("The search query, phrased like a web search"),
    maxResults: z.number().int().min(1).max(20).optional().describe("How many results to return (default 5)"),
    includeDomains: z
      .array(z.string())
      .optional()
      .describe("Only return results from these domains, e.g. ['arxiv.org']"),
  }),
  execute: async ({ query, maxResults = 5, includeDomains }) => {
    try {
      const res = await mux.search({
        query,
        models: ["parallel/fast", "exa/auto"],
        maxResults,
        includeDomains,
      });
      return {
        model: res.model,
        results: res.results.map((r) => ({
          title: r.title,
          url: r.url,
          snippet: r.snippet,
          published_date: r.publishedDate,
        })),
      };
    } catch (e) {
      if (e instanceof MuxError) {
        // Return the problem as data so the model can tell the user instead of crashing the run.
        return { error: e.code, message: e.message, results: [] };
      }
      throw e;
    }
  },
});

Two choices worth explaining. models: ["parallel/fast", "exa/auto"] tries Parallel first at $0.001 per search and falls back to Exa at $0.007 only if Parallel fails or is at capacity. You pay for the model that answered; failed attempts are refunded. And the result object omits content, usage and attempts, because the model does not need them. Log res.usage.costUsd and res.attempts on your side if you want them.

Multi-step generateText#

The agent loop is one call. stopWhen: isStepCount(5) lets the model search, read the results, search again if it wants, and then answer, up to five model calls.

// agent.ts
import { generateText, isStepCount } from "ai";
import { anthropic } from "@ai-sdk/anthropic";
import { webSearch } from "./web-search";

const result = await generateText({
  model: anthropic("claude-opus-5-5"),
  tools: { webSearch },
  stopWhen: isStepCount(5),
  system:
    "You are a research assistant. Use webSearch for current facts and cite the URL of every source you use.",
  prompt: "What changed in the EU AI Act implementation timeline this month?",
  onStepEnd: ({ toolCalls }) => {
    for (const call of toolCalls) console.log("searched:", call.input);
  },
});

console.log(result.text);
console.log(`${result.steps.length} steps`);

Run it with npx tsx agent.ts. result.steps holds every step with its tool calls and tool results, which is useful for building a sources list from the URLs the model actually looked at.

Streaming with streamText#

streamText takes the same options. The text stream only yields text, so the tool calls happen silently between chunks. Iterate result.stream instead if you want to show a "searching..." indicator when a tool-call part arrives.

import { streamText, isStepCount } from "ai";
import { anthropic } from "@ai-sdk/anthropic";
import { webSearch } from "./web-search";

const result = streamText({
  model: anthropic("claude-opus-5-5"),
  tools: { webSearch },
  stopWhen: isStepCount(5),
  prompt: "Summarize this week's news about solid-state batteries with sources.",
});

for await (const part of result.stream) {
  if (part.type === "tool-call") process.stdout.write("\n[searching] ");
  if (part.type === "text-delta") process.stdout.write(part.text);
}

In a Next.js route handler, return result.toUIMessageStreamResponse() and the useChat hook renders tool calls as parts of the message.

Tuning the tool#

Recency. Add publishedAfter to the schema and pass it through. Mux normalizes dates to YYYY-MM-DD across providers. Note that publishedBefore is not applied by Parallel models.
More text per result. content: "full" returns page text on Exa and joined excerpts on Parallel. It is much larger, so use it in a separate readPage-style tool rather than in every search.
Quality over price. Swap the fallback order to ["exa/auto", "parallel/advanced"] for research-heavy agents. The models page lists latency and price for each option.

Model	Per 1,000 requests	Included results	Extra result	Latency
exa/instant Exa instant, lowest-latency neural search	$4.00	10	+$1.00 / 1k	~300 ms
exa/fast Exa fast	$7.00	10	+$1.00 / 1k	~500 ms
exa/auto Exa auto (default, best general quality)	$7.00	10	+$1.00 / 1k	~1 s
exa/deep Exa deep multi-query search	$12.00	10	+$1.00 / 1k	~5 s
parallel/turbo Parallel turbo, ~200ms	$1.00	10	+$1.00 / 1k	~200 ms
parallel/fast Parallel fast, <1s	$1.00	10	+$1.00 / 1k	~800 ms
parallel/basic Parallel basic	$5.00	10	+$1.00 / 1k	~1 s
parallel/advanced Parallel advanced, highest quality excerpts	$5.00	10	+$1.00 / 1k	~3 s
keenable/pro Keenable pro index	$4.00	10	none	~700 ms
keenable/realtime Keenable realtime index	$4.00	10	none	~700 ms

Cost control#

An agent in a loop can issue many searches, and the AI SDK's step limit caps model calls, not spend. Put a hard ceiling on the Mux side by giving the agent its own key with a credit limit. Create it once with a management key (mux-mgmt-...):

import { Mux } from "getmux";

const admin = new Mux({ apiKey: process.env.MUX_MGMT_KEY });
const { key } = await admin.keys.create({
  name: "research-agent",
  limit: 10,            // USD
  limitReset: "daily",  // or "weekly" | "monthly"
});
// Put `key` in MUX_API_KEY for the agent process. It is shown once.

When the key reaches its limit, Mux returns 402 key_limit_exceeded and charges nothing. The dashboard shows spend per key, so one key per agent or per customer gives you a usage breakdown for free. Details on managing keys.

Handling 402 and 429#

The getmux client already retries 429 rate_limited and most 5xx responses twice with backoff, honoring Retry-After. The errors you should handle yourself:

Error	Meaning	What to do in a tool
`402 insufficient_credits`	Account balance is empty	Return an error object so the model says it cannot search right now; alert an operator to add credits.
`402 key_limit_exceeded`	This key hit its daily, weekly or monthly limit	Same as above. Raise the limit or wait for the reset.
`429 rate_limited`	More than 20 requests per second on one key	Already retried by the SDK. If it still surfaces, lower the agent's concurrency.
`429 too_many_concurrent_requests`	More than 50 searches in flight on the account	Also retried. Queue tool executions if you run many agents in parallel.
`502 all_providers_failed`	Every model in `models` failed	Not retried by the SDK. Safe to retry once; nothing was charged.

The catch block in the tool above turns every MuxError into data rather than throwing. That is usually the right call inside an agent, because a thrown error aborts the whole generateText run, while an error object lets the model explain the problem or answer from what it already knows. Switch on e.code if you want different behavior for billing errors versus transient ones. The full list is on the errors page.

FAQ#

Why not use a provider's SDK directly as the tool?#

You can, and at very high volume with a negotiated contract a direct relationship is cheaper. Mux is for the case where you want more than one index behind the same tool: models gives you fallback when a provider is slow or down, the response shape is the same for all three, and billing, rate limits and per-key budgets are in one place.

Does the AI SDK's step limit stop runaway search costs?#

Only indirectly. isStepCount(5) caps the number of model calls, and the model may call the tool several times within one step. The per-key credit limit on the Mux side is the hard stop, and it is enforced atomically before the search runs.

Can I use this tool with OpenAI or Google models?#

Yes. The tool is model-agnostic. Install the matching provider package, for example @ai-sdk/openai, and pass that provider's model to generateText. Nothing in web-search.ts changes.

How do I keep the model from flooding its context with results?#

Keep maxResults small by default (5 is plenty for most questions), return only title, URL, snippet and date as the example does, and reserve content: "full" for a separate tool the model calls on one URL at a time.

Get a key and run your first search in the quickstart.