Agent Memory in a Token Budget

The problem

An agent's conversation grows with every step, but its context window (what the model can read at once) has a fixed size. When the conversation no longer fits, the oldest messages have to go. A sliding window keeps the newest ones; adding a summary of the dropped messages keeps their gist. You will write the function that fits a conversation into a budget of tokens.

The rules

messages is a list of dicts {"role": ..., "content": ...}. The first one has role "system": it is always kept. count_tokens(message) returns the size of a message. summarize(dropped) returns the text of a summary of the dropped messages.

• If everything fits, return a copy of messages.
• Otherwise keep the newest messages that fit: walk back from the end and keep each message while the total, system prompt included, stays within budget. Stop at the first one that does not fit: the kept messages must be a contiguous, most-recent block.
• Replace the dropped messages with one summary message {"role": "system", "content": "Summary of earlier messages: " + summarize(dropped)}, placed right after the system prompt.
• The summary must fit too: if the total with the summary is over budget, drop the oldest kept messages (and add them to the summary) until it fits.
• Return the new list, in chronological order.

Example. A system message of 10 tokens, then four messages $m_1, \dots, m_4$ of 30, 40, 20 and 15 tokens; a summary takes 5 tokens.

• Budget 200: everything fits ($115$), nothing changes.
• Budget 70: walk back: $10 + 15 = 25$, then $+ 20 = 45$ fit; adding $m_2$ (40) would give 85: stop. Dropped: $m_1, m_2$. With the summary: $10 + 5 + 20 + 15 = 50 \le 70$. Result: system, summary of $m_1, m_2$, $m_3$, $m_4$.
• Budget 45: the two newest fit alone ($10 + 20 + 15 = 45$), but with a summary the total is 55. Move $m_3$ into the summary too: $10 + 5 + 15 = 30$. Result: system, summary of $m_1, m_2, m_3$, $m_4$.

What to implement

fit_context(messages, budget, count_tokens, summarize): return the new list.

Requirements

• Never drop the system prompt; you can assume it fits on its own with the summary.
• Call summarize with the dropped messages in chronological order. You may call it more than once while sizing the summary, but the summary in the result must describe exactly the messages that were finally dropped.

Hint

First decide how many recent messages fit, then build the summary and check the total; if it is over budget, move one more message into the dropped set and repeat. Recompute the summary's size each time the dropped set changes.