Build an AI Agent in Typescript / lesson 6 of 8

Context Compaction

Stop blowing through your context window with this automated AI memory management system.

Stop blowing through your context window with this automated AI memory management system. We build a local coding agent that dynamically compresses message history, preventing context overflow while preserving critical tool calls and reasoning. This approach enables infinite-context workflows for large codebases without hitting token limits.

Measuring Token Usage

To estimate token usage, we can use a simple heuristic that approximates the number of tokens in a message by calculating the length of the JSON string representation of the message. This is a rough estimate, but it’s good enough for our purposes.

import type { ModelMessage } from "ai";

export const estimateTokens = (messages: ModelMessage[]): number => {
  return JSON.stringify(messages).length / 4;
};

Using the Token Estimator

const agentLoop = async (messages: ModelMessage[]): Promise<string> => {
 console. log('token count: ${estimateTokens (messages) }');
 let stepsSinceTodo = 0;

Try it out

input >> give me 50 writing prompts
token count: 14
...
...
input >> what is 2+2?
token count: 989.75 # ~4000 characters

Let’s set our agent loop to automatically compact this history

export const compact = (messages: ModelMessage[]): ModelMessage[] => {
  // TODO: implement
};

import type { ModelMessage } from "ai";
import { generateText } from "ai";
import { MODEL, api, CONTEXT_LIMIT } from "./config";

const KEEP_RECENT = 3;
const PRESERVE_TOOLS = new Set(["read_file"]);

type Part = {
  type?: string;
  toolCallId?: string;
  toolName?: string;
  content?: unknown;
};

export const estimateTokens = (messages: ModelMessage[]): number => {
  return JSON.stringify(messages).length / 4;
};

const collectToolResults = (messages: ModelMessage[]):
    { part: { content: unknown; }; toolName: string; }[] => {
  const names = new Map();

  return messages
    .flatMap(m => (Array.isArray(m.content) ? (m.content as Part[]) : []))
    .reduce((acc, p) => {
      if (p.type === 'tool_use' && p.toolCallId) {
        names.set(p.toolCallId, p.toolName);
      }
      if (p.type === 'tool_result' && p.toolCallId) {
        acc.push({
          part: p as { content: unknown; },
          toolName: names.get(p.toolCallId) ?? 'unknown'
        });
      }
      return acc;

    }, [] as { part: { content: unknown; }; toolName: string; }[]
    );

};

export const microCompact = (messages: ModelMessage[]): void => {
  collectToolResults(messages)
    .slice(0, -KEEP_RECENT)
    .forEach(({ part, toolName }) => {
      if (
        !PRESERVE_TOOLS.has(toolName) &&
        typeof part.content === 'string' &&
        part.content.length > 100
      ) {
        part.content = [`Previous: used ${toolName}`];
      }
    });
};

export const autoCompact = async (messages: ModelMessage[]):
    Promise<ModelMessage[]> => {
  const { text } = await generateText({
    model: api.chat(MODEL),
    messages: [
      {
        role: "user",
        content: `
        Summarize this conversation for continuity.
        Include:
        1) What was accomplished,
        2) Current state,
        3) Key decisions made.
        Be concise.

        ${JSON.stringify(messages).slice(-CONTEXT_LIMIT)}
        `,
      }
    ]
  });

  return [
    {
      role: "user",
      content: `[Conversation compressed]\n\n${text || "No summary."}`,
    }
  ];
};

Adding compaction to our Agent Loop

import { generateText, isLoopFinished } from "ai";
import type { ModelMessage } from "ai";
import * as readline from "readline";
import { PARENT_TOOLS, } from "./tools";
import { api, MODEL, WORKDIR } from "./config";
import { api, CONTEXT_LIMIT, MODEL, WORKDIR } from "./config";
import { autoCompact, estimateTokens, microCompact } from "./compact";

/** AGENT LOOP */
const agentLoop = async (messages: ModelMessage[]): Promise<string> => {
  let stepsSinceTodo = 0;
  microCompact(messages);
  const estTokens = estimateTokens(messages);
  console.log (' tokens, ${estTokens}');
  if (estTokens > CONTEXT_LIMIT) {
    console. log ('compacting at $festTokens}');
    const compacted = await autoCompact(messages);
    messages.length = 0;
    messages.push(...compacted);
    console.log ('new tokens, ${estimateTokens (messages) }');
  }

You can try this out by setting a CONTEXT_LIMIT of ~3000.

input >> give me 300 writing prompts
tokens, 14.25
...
...
input >> what is 2+2?
tokens, ???

Share this post on:

Task Persistence