Pack an agent's context window under a token budget

Agent buildHard40 min

Your task

Before every model call, an agent's runtime packs the conversation into the model's context window. Long tool results quickly fill it, so older history has to go, but the request must stay valid: model APIs reject a tool result without the call that asked for it, and a tool call without its result. Write pack_context(messages, budget, note_cost).

messages is the conversation, oldest first, as [role, tokens, ref] with tokens a positive integer. The roles are:

"system": instructions.
"user" and "assistant": ordinary turns (ref is "").
"call": the assistant asking for a tool call. ref is the call id; call ids are unique.
"result": a tool's result. ref is the id of the call it answers.

Rules

Pinned. Every "system" message and the last "user" message are always kept.
Pairs. A "call" and every "result" with the same ref form one group, kept or dropped together. A "result" whose call is not in the list, and a "call" with no result, are malformed: they are always dropped, cost nothing, and their dropping alone never needs a note.
Singles. Every other message (an "assistant" turn, or a "user" turn that is not the pinned one) is a group of its own.
Everything fits. If the pinned messages plus all groups fit in budget tokens, keep all of them, with no note.
Otherwise, a note saying "earlier history was dropped" is added, costing note_cost tokens. Order the groups from newest to oldest, where a group's position is the index of its earliest message. Keep each group while it fits in what is left of budget - pinned - note_cost, and stop at the first group that does not fit: no older group is kept, even a small one, so the kept history has no holes in the middle.
Cannot fit. If the pinned messages alone exceed budget, or a note is needed and the pinned messages plus the note exceed it, return None (null in JavaScript).

Return {"keep": [...], "note": bool}, with keep the indices of the kept messages in increasing order.

Example, with these messages:

0  ["system", 50, ""]
1  ["user", 30, ""]
2  ["assistant", 20, ""]
3  ["call", 10, "c1"]
4  ["result", 200, "c1"]
5  ["assistant", 40, ""]
6  ["user", 25, ""]          the last user message: pinned
7  ["call", 10, "c2"]
8  ["result", 60, "c2"]

With budget 500 and note_cost 15, everything (445 tokens) fits: {"keep": [0, 1, 2, 3, 4, 5, 6, 7, 8], "note": False}.
With budget 300 and note_cost 15, the pinned messages cost 75, so 210 tokens are left for groups. Newest first: messages 7 and 8 (70) fit, message 5 (40) fits, messages 3 and 4 (210) do not, and the walk stops there, so messages 1 and 2 are dropped even though they would fit. The result is {"keep": [0, 5, 6, 7, 8], "note": True}.
With budget 70, the pinned messages alone (75) do not fit: None.

Examples

pack_context(messages: list[list], budget: int, note_cost: int) -> dict | None

Example 1

Input

pack_context([["system", 50, ""], ["user", 30, ""], ["assistant", 20, ""], ["call", 10, "c1"], ["result", 200, "c1"], ["assistant", 40, ""], ["user", 25, ""], ["call", 10, "c2"], ["result", 60, "c2"]], 500, 15)

Output{"keep": [0, 1, 2, 3, 4, 5, 6, 7, 8], "note": False}

Example 2

Input

pack_context([["system", 50, ""], ["user", 30, ""], ["assistant", 20, ""], ["call", 10, "c1"], ["result", 200, "c1"], ["assistant", 40, ""], ["user", 25, ""], ["call", 10, "c2"], ["result", 60, "c2"]], 300, 15)

Output{"keep": [0, 5, 6, 7, 8], "note": True}

Example 3

Input

pack_context([["system", 50, ""], ["user", 30, ""], ["assistant", 20, ""], ["call", 10, "c1"], ["result", 200, "c1"], ["assistant", 40, ""], ["user", 25, ""], ["call", 10, "c2"], ["result", 60, "c2"]], 70, 15)

OutputNone