LLM agent token limit exceeded in TypeScript: detect before the provider throws, recover with partial output

import OpenAI from "openai";
import { guard, ContextOverflowError, LoopDetectedError, BudgetExceededError } from "@runguard/sdk";
import { encode } from "gpt-tokenizer"; // or tiktoken, your choice

const client = new OpenAI();

type Message = OpenAI.Chat.ChatCompletionMessageParam;

async function callModel(messages: Message[]): Promise<OpenAI.Chat.ChatCompletion> {
  return client.chat.completions.create({
    model: "gpt-4o",
    messages,
  });
}

const guardedCall = guard(callModel, {
  // Project token count before each call — fires before the HTTP request
  context: {
    maxContextTokens: 128_000,
    headroom: 8_000,  // reserve 8k for the completion
  },
  tokens: (messages: Message[]) => {
    // gpt-tokenizer gives token-accurate counts for GPT-4o
    return messages.reduce((total, msg) => {
      const text = typeof msg.content === "string"
        ? msg.content
        : (msg.content ?? [])
            .map((c) => (c.type === "text" ? c.text : ""))
            .join("");
      return total + encode(text).length + 4; // 4 overhead per message
    }, 3); // 3 for reply priming
  },
  loop: { repeats: 3, maxCycleLen: 4 },
  budget: { maxUsd: 3.0 },
});

async function runAgent(initialMessages: Message[]): Promise<string> {
  const history: Message[] = [...initialMessages];

  while (true) {
    try {
      const completion = await guardedCall(history);
      const reply = completion.choices[0].message;
      history.push(reply);

      if (!reply.tool_calls?.length) {
        return reply.content ?? "";
      }

      // Execute tools, append results...
      for (const tc of reply.tool_calls) {
        const result = await executeTool(tc);
        history.push({
          role: "tool",
          tool_call_id: tc.id,
          content: JSON.stringify(result),
        });
      }

    } catch (err) {
      if (err instanceof ContextOverflowError) {
        // Context is still in memory — compact and retry
        console.warn(
          `Context guard fired: projected ${err.projectedTokens} tokens ` +
          `vs ${err.maxContextTokens - err.headroom} limit. Compacting.`
        );
        return await compactAndResume(history);
      }
      if (err instanceof LoopDetectedError) {
        return `Agent loop detected (pattern: ${err.pattern}). Run aborted.`;
      }
      if (err instanceof BudgetExceededError) {
        return `Budget cap reached ($${err.spent.toFixed(3)}). Run aborted.`;
      }
      throw err;
    }
  }
}

async function compactAndResume(history: Message[]): Promise<string> {
  // Strategy: summarize the middle of the conversation, keep system + last 5 turns
  const system = history.filter((m) => m.role === "system");
  const tail = history.filter((m) => m.role !== "system").slice(-10);
  const middle = history.filter((m) => m.role !== "system").slice(0, -10);

  const summary = await client.chat.completions.create({
    model: "gpt-4o-mini",  // cheap model for summarization
    messages: [
      { role: "system", content: "Summarize the following conversation history concisely, preserving key findings and decisions." },
      { role: "user", content: middle.map((m) => `${m.role}: ${typeof m.content === "string" ? m.content : "[multi-part]"}`).join("\n") },
    ],
  });

  const compacted: Message[] = [
    ...system,
    { role: "assistant", content: `[Conversation summary: ${summary.choices[0].message.content}]` },
    ...tail,
  ];

  // Re-run with compacted history
  return runAgent(compacted);
}

declare function executeTool(tc: OpenAI.Chat.ChatCompletionMessageToolCall): Promise<unknown>;