Pruning Old Results

The fix is three lines.

That's the part that's going to feel anticlimactic. You measured the problem in the last lesson, watched the input tokens climb, sketched out the disaster scenario at step thirty. Now we add three lines and the curve flattens.

The lines themselves are easy. Where they go and why is the lesson.

Outcome

prepareStep runs pruneMessages before every step of the loop, removing tool call/result pairs older than the last three messages. The token-growth curve from the last lesson plateaus instead of climbing forever.

Fast Track

1. Import pruneMessages from ai
2. Add a prepareStep to your ToolLoopAgent config
3. Inside it, return { messages: pruneMessages({ messages, toolCalls: "before-last-3-messages" }) }

Hands-on Exercise 5.2

Wire pruning into the agent and re-run the same multi-step task from lesson 5.1.

Requirements:

1. Add prepareStep: ({ messages }) => ({...}) to the agent config
2. Return a pruned messages array from it
3. Use toolCalls: "before-last-3-messages" for now (the simplest reasonable strategy)
4. Confirm input tokens plateau across steps instead of climbing

Implementation hints:

• prepareStep runs before each step, and it receives the messages the SDK is about to send for that step. You're modifying them on the way in
• Return only what you want to change. Anything you leave out (model, tools, instructions) keeps its current value
• messages is always an array in prepareStep, even on step 0, so there's no undefined case to guard

The fix

import { ToolLoopAgent, isStepCount, tool, pruneMessages } from "ai";
 
const agent = new ToolLoopAgent({
  // ... existing config
  prepareStep: ({ messages }) => ({
    messages: pruneMessages({ messages, toolCalls: "before-last-3-messages" }),
  }),
});

Three lines, one import.

What's actually happening

Before each step, prepareStep runs. It receives the messages the SDK is about to send for that step. We return a pruned version, dropping every tool call and result that's older than the last three messages. The SDK uses that array for this step and builds the next step on top of it.

Before pruning at step 15:
  [user prompt]
  [assistant + tool_call] -> [tool_result]    (old, will be pruned)
  [assistant + tool_call] -> [tool_result]    (old, will be pruned)
  ... 12 more pairs ...
  [assistant + tool_call] -> [tool_result]    (recent, kept)
  [assistant + tool_call] -> [tool_result]    (recent, kept)
  [assistant] -> [user]                       (recent, kept)

After pruning:
  [user prompt]                               (kept, original prompt)
  [assistant + tool_call] -> [tool_result]    (recent)
  [assistant + tool_call] -> [tool_result]    (recent)
  [assistant] -> [user]                       (recent)

The original user prompt always survives. The recent tool interactions survive. The middle of the conversation, where the tool results pile up, gets dropped on each call.

Why three messages

The toolCalls: "before-last-3-messages" setting keeps the last three messages of conversation, not just the last three tool pairs. That's enough context for the model to know where it is in a multi-step task without keeping the whole history.

You can tune this. before-last-1-messages is more aggressive and saves more tokens. before-last-5-messages is gentler and keeps more context. Three is a reasonable default that works well across task shapes. Start there. Tune later if you have a specific task that needs it.

Try It

Run the same multi-step task from lesson 5.1 and compare the token curves:

bun run index.ts . "Read package.json, tsconfig, index.ts, then summarize"

You should see something like:

Step 0: 1,200 input, 450 output
Step 1: 2,800 input, 200 output
Step 2: 3,100 input, 180 output    (old results pruned)
Step 3: 3,400 input, 350 output    (growth plateaus)
Step 4: 3,200 input, 600 output    (stays flat)

The exact numbers depend on your project. The shape is what matters. Input tokens plateau by step 2 or 3 instead of climbing forever.

npx tsc --noEmit

Commit

git add index.ts
git commit -m "feat(context): prune old tool results in prepareStep"

Done-When

• pruneMessages is imported from ai
• prepareStep is wired into the agent config
• prepareStep returns the pruned messages array
• On a 4+ step task, input tokens plateau instead of growing linearly
• npx tsc --noEmit passes

Solution

import { ToolLoopAgent, isStepCount, tool, pruneMessages } from "ai";
 
const agent = new ToolLoopAgent({
  model: "anthropic/claude-haiku-4.5",
  instructions: buildSystemPrompt({ /* ... */ }),
  tools,
  stopWhen: isStepCount(15),
  onStepEnd: ({ usage, stepNumber }) => {
    console.error(
      `Step ${stepNumber}: ${usage.inputTokens} input, ${usage.outputTokens} output`,
    );
  },
  prepareStep: ({ messages }) => ({
    messages: pruneMessages({ messages, toolCalls: "before-last-3-messages" }),
  }),
});