📘

Trim a Message History to a Token Budget

Problem Statement

An agent's context window is finite. Before every model call the runtime must decide which of the accumulated messages still fit. Naively slicing the last N messages corrupts the conversation: you can end up starting mid-turn, with an assistant message answering a user message that is no longer present.

Write the trimmer that a production agent loop would use.

Background

A chat transcript is a list of messages, each with a role (system, user or assistant) and a pre-computed tokens count. Two invariants must survive trimming:

System messages are pinned. They carry the agent's instructions and tool definitions, so they are never dropped and their tokens are spent first.
Turns are atomic. A turn starts at a user message and includes every message after it up to (but not including) the next user message. You drop whole turns from the oldest end, never half a turn. This guarantees the retained history always begins on a user message.

Any non-system messages appearing before the first user message belong to no turn and are discarded.

Your Task

Implement:

def trim_history(messages, max_tokens):

messages: list of dicts {"id": str, "role": str, "tokens": int}, in chronological order.
max_tokens: int, the budget for the whole kept set.

Algorithm:

Sum the tokens of all system messages. This is the base cost.
Group the remaining messages into turns as described.
Starting from the full history, while the total (base + kept turns) exceeds max_tokens, drop the oldest remaining turn.
If, after dropping every turn, the system messages alone still exceed the budget, return only the system messages.

Input/Output Format

Return the kept messages as a list of the original dicts, in the original chronological order (system messages keep their original positions).

Sample

msgs = [
    {"id": "s0", "role": "system", "tokens": 20},
    {"id": "u1", "role": "user", "tokens": 30},
    {"id": "a1", "role": "assistant", "tokens": 30},
    {"id": "u2", "role": "user", "tokens": 25},
    {"id": "a2", "role": "assistant", "tokens": 25},
]
print([m["id"] for m in trim_history(msgs, 100)])
# ['s0', 'u2', 'a2']

Full history costs 20 + 60 + 50 = 130 > 100. Drop the oldest turn (u1,a1, 60 tokens) leaving 70, which fits.

Example:

Input:

msgs = [{'id':'s0','role':'system','tokens':20},{'id':'u1','role':'user','tokens':30},{'id':'a1','role':'assistant','tokens':30},{'id':'u2','role':'user','tokens':25},{'id':'a2','role':'assistant','tokens':25}]
print([m['id'] for m in trim_history(msgs, 100)])

Output:

['s0', 'u2', 'a2']

Reasoning:

Base cost is the system message: 20 tokens. Turn 1 = u1+a1 = 60, turn 2 = u2+a2 = 50. Total 130 > 100, so the oldest turn is dropped, leaving 20 + 50 = 70 <= 100. The kept set starts on a user message, as required.

Constraints:

0 <= len(messages) <= 500
1 <= max_tokens <= 10**6; every tokens value is a non-negative int
role is one of "system", "user", "assistant"
System messages are always kept, even when they alone blow the budget
The returned list must preserve the original chronological order
Never split a turn: a user message and its following assistant messages are dropped together

More from AI Agents 1

Cheapest Cache TTL Meeting a Freshness SLAHard Bootstrap Confidence Interval Width for Eval AccuracyMedium Route to the Cheapest Model Meeting QualityMedium pass@k EstimatorMedium Savings from a Semantic CacheMedium

1. Background Knowledge

In Large Language Model (LLM) applications, the context window represents the maximum number of tokens the model can process in a single request. As a conversation grows, the accumulated history often exceeds this limit. Simply truncating the message list by index is dangerous because it can break the semantic structure of the dialogue. For instance, removing a user prompt while keeping the subsequent assistant response leaves the model with an orphaned answer, leading to incoherent or hallucinated outputs.

To maintain coherence, we treat conversation history as a sequence of atomic turns. A turn typically begins with a user input and includes all subsequent messages (e.g., assistant responses, tool calls) until the next user input. By treating these groups as indivisible units, we ensure that any retained history is logically complete. Additionally, system messages contain critical instructions and persona definitions; they are usually pinned to the beginning of the context and are never dropped, regardless of the token budget.

This problem introduces the concept of sliding window memory with structural constraints. Unlike a standard sliding window that moves one element at a time, here the "window" must expand or contract in chunks (turns). This requires a two-phase approach: first, identifying the structural boundaries (turns), and second, applying a greedy removal strategy from the oldest end to satisfy the token constraint while preserving the most recent context.

2. Algorithm Approach

The core algorithmic pattern here is Greedy Truncation with Structural Grouping. We cannot simply iterate through messages one by one because of the atomic turn constraint. Instead, we must first restructure the data into logical units.

Separation: Isolate system messages from the rest. System messages are fixed costs.
Grouping: Partition the non-system messages into turns. A turn is defined as a contiguous block of messages starting with a user role and ending just before the next user role.
Cost Calculation: Compute the token cost for each turn.
Iterative Removal: Start with all turns included. Calculate the total token count (System Cost + Sum of all Turn Costs). While this total exceeds max_tokens, remove the oldest turn (the one at the beginning of the list) and subtract its cost.
Reconstruction: Once the budget is satisfied (or no turns remain), flatten the remaining turns back into a list of messages, prepending the system messages.

This approach ensures that we always keep the most recent context, which is generally the most relevant for the model's next prediction, while strictly adhering to the turn-atomicity invariant.

3. Step-by-Step Strategy

Identify System Messages: Iterate through the input messages list. Collect all messages with role == "system" into a separate list system_msgs. Calculate system_cost as the sum of their tokens.
Filter and Group Non-System Messages: Create a list of non-system messages. Iterate through this list to group them into turns.

Initialize an empty list turns.
Use a temporary buffer current_turn.
For each message, if role == "user", save the current_turn (if not empty) to turns and start a new current_turn.
Append the message to current_turn.
After the loop, append the final current_turn to turns.
Note: Any messages appearing before the first user message are discarded as per the problem statement.

Calculate Turn Costs: For each turn in turns, compute its total token cost. Store this alongside the turn data (e.g., as a tuple (cost, turn_messages)).
Trim Oldest Turns:

Calculate current_total = system_cost + sum(turn_cost for all turns).
While current_total > max_tokens and turns is not empty:
Pop the first turn (oldest) from turns.
Subtract its cost from current_total.

Reconstruct Output:

Flatten the remaining turns back into a single list of message dictionaries.
Prepend system_msgs to this list.
Return the combined list.

4. Common Pitfalls

Ignoring Pre-User Messages: The problem states that non-system messages before the first user message belong to no turn and are discarded. Failing to handle this edge case can lead to incorrect grouping or including invalid context.
Modifying Lists While Iterating: When removing turns, avoid iterating over the turns list with a for loop while popping elements. Use a while loop with an index or pop(0) carefully to avoid index errors or skipping elements.
System Message Positioning: The output must maintain the original chronological order. System messages usually appear at the start, but if they were interspersed in the input (though rare in this specific problem structure), you must ensure they are placed correctly in the final output. The problem implies system messages are pinned, so prepending them is the standard approach.
Empty Turn Handling: Ensure that your grouping logic handles cases where there are no user messages at all. In such cases, turns will be empty, and only system messages should be returned (if they fit).
Token Count Mismatch: Double-check that you are summing the tokens field correctly. A common error is summing the length of the text string instead of the pre-computed tokens integer.

5. Time & Space Complexity

Time Complexity: $O(N)$ , where $N$ is the number of messages.
We iterate through the messages once to separate system messages and group turns.
We iterate through the turns to calculate costs and potentially remove them. In the worst case, we remove all turns, which is linear with respect to the number of turns.
Flattening the remaining turns is also linear.
Since the number of turns is less than or equal to $N$ , the overall complexity remains linear.
Space Complexity: $O(N)$ .
We create new lists for system_msgs, turns, and the final output.
In the worst case, we store references to all original message dictionaries.
No additional complex data structures are used, so the space overhead is proportional to the input size.

solution.py

Python 3.13

Test Results

0/0

Run code to see test results.