Sliding-Window Context Truncation

An agent’s context window is finite, but a conversation is not. Every extra turn you carry forward costs tokens, and once the transcript would exceed the model’s budget you have to drop something before the next call. The standard move is a sliding window: pin the system prompt (it holds the instructions the agent must never forget), then keep only the most recent turns that still fit. Old small talk falls off the front; the latest exchange — the part the model actually needs to answer — is protected.

The subtlety is that you decide what to keep by scanning from the newest message backwards, but you must return the survivors in their original order so the model reads the conversation top to bottom.

Task

Complete truncate_context(messages, max_tokens):

1. If the first message has role == "system", keep it and subtract its token_count from the budget. (If even that one message exceeds max_tokens, keep nothing.)
2. Walk the remaining non-system messages from newest to oldest. Keep each one whose token_count still fits in the remaining budget, subtracting as you go. Stop as soon as a message would overflow — do not skip it to squeeze in a smaller older one.
3. Return the kept messages (the system message, if kept, plus the recent ones) in their original chronological order.

Example

messages = [
    {"role": "system",    "token_count": 10, "content": "You are a helpful assistant."},
    {"role": "user",      "token_count": 20, "content": "Hi"},
    {"role": "assistant", "token_count": 30, "content": "Hello!"},
    {"role": "user",      "token_count": 25, "content": "What's 2+2?"},
    {"role": "assistant", "token_count": 15, "content": "4"},
]

kept = truncate_context(messages, max_tokens=60)
# system(10) is pinned -> 50 left.
# newest 15 fits -> 35 left; 25 fits -> 10 left; 30 would overflow, stop.
# kept == [system(10), user(25), assistant(15)]
[m["token_count"] for m in kept]
# [10, 25, 15]

Constraints

• max_tokens >= 0. The sum of token_count over the returned list must be <= max_tokens.
• Only the first message can be a system message; treat any later message as ordinary conversation.
• If there is no system message, spend the whole budget on the most recent non-system messages.
• Stop at the first message (scanning newest to oldest) that does not fit; do not skip ahead to fit a smaller older message.
• Every token_count is a non-negative integer. Do not mutate the input list or its dicts; return a new list.

Approach

Pay the fixed cost first: if the transcript opens with a system message and it fits, pin it and subtract its tokens from the budget. Then greedily fill the leftover budget by walking the non-system messages from newest to oldest, keeping each until one would overflow — recency is what a sliding window protects. Because we collected survivors while scanning backwards, we reverse that recent-slice at the end so the returned list is in chronological order, with the pinned system message in front.

Solution

def truncate_context(messages: list[dict], max_tokens: int) -> list[dict]:
    kept: list[dict] = []
    budget = max_tokens

    # 1. Pin the leading system message if it fits.
    body = messages
    if messages and messages[0].get("role") == "system":
        system = messages[0]
        body = messages[1:]
        if system["token_count"] <= budget:
            kept.append(system)
            budget -= system["token_count"]
        else:
            return []  # can't even fit the pinned instruction

    # 2. Keep the most recent non-system messages that fit, newest first.
    recent: list[dict] = []
    for msg in reversed(body):
        if msg["token_count"] <= budget:
            recent.append(msg)
            budget -= msg["token_count"]
        else:
            break  # first overflow stops the window

    # 3. Restore chronological order: system, then oldest-to-newest recents.
    recent.reverse()
    kept.extend(recent)
    return kept

Walkthrough

Tracing the example with max_tokens=60:

1. The first message is a system costing 10, and 10 <= 60, so we pin it. kept = [system], budget = 50. body is the four non-system messages.
2. We scan body in reverse. The newest is assistant(15): 15 <= 50, keep it, budget = 35. Next newest is user(25): 25 <= 35, keep it, budget = 10. Next is assistant(30): 30 > 10, so we break — we do not skip it to reach the older user(20).
3. recent is [assistant(15), user(25)] in newest-first order. Reversing gives [user(25), assistant(15)]. Appending to kept yields [system(10), user(25), assistant(15)], whose token counts are [10, 25, 15] summing to 50 — within budget and in chronological order.

If the transcript had no system message, step 1 leaves body = messages and the whole 60-token budget goes to the most recent turns. If the lone system message cost 70, step 1 returns [] immediately.

Complexity & notes

• Time is O(n) over the messages (a single reverse scan plus a final reverse of the kept slice); space is O(k) for the k messages returned. No sorting is needed because the input is already chronological.
• The stop-at-first-overflow rule makes this a true sliding window, not a knapsack: we never skip a big recent message to fit a smaller older one, because dropping a recent turn to keep older history is exactly the wrong trade for an agent.
• Pinning the system message before filling recents is what guarantees the instruction survives; a naive “keep newest until full” would let a long tail of chatter evict the very prompt that defines the agent’s behavior.
• Returning [] when the system message alone overflows is a deliberate signal that the budget is misconfigured — silently dropping the instruction would be worse than returning nothing.
• We build a new list and never mutate the inputs, so the caller can keep the full transcript for logging while sending only the truncated view to the model.
• In production the token_count would come from the model’s real tokenizer and you might reserve headroom for the response; the windowing logic here stays the same regardless of how the counts are computed.