Compaction: Summarizing Without Losing the Plot

# Colab cell 1 — run once. Set your key in the 🔑 panel (name it
# ANTHROPIC_API_KEY) or just paste it when prompted.
!pip install -q anthropic

import os
try:
    from google.colab import userdata
    os.environ["ANTHROPIC_API_KEY"] = userdata.get("ANTHROPIC_API_KEY")
except Exception:
    from getpass import getpass
    os.environ.setdefault("ANTHROPIC_API_KEY", getpass("Anthropic API key: "))

import anthropic

client = anthropic.Anthropic()
MODEL = "claude-sonnet-5"

def count(messages: list, system: str = "") -> int:  # lesson 1's counter
    kwargs = {"model": MODEL, "messages": messages}
    if system:
        kwargs["system"] = system
    return client.messages.count_tokens(**kwargs).input_tokens

COMPACT_AT = 0.75          # of the window budget
WINDOW_BUDGET = 60_000     # tokens you allow the conversation to occupy
KEEP_RECENT = 8            # messages kept verbatim

SUMMARIZER_PROMPT = (
    "Summarize this conversation prefix for an agent that will continue it.\n"
    "PRESERVE EXACTLY: the user's goal, all standing constraints and "
    "prohibitions, decisions made, and every file path, identifier, number, "
    "and error message. Note outcomes of tool calls, not their transcripts.\n"
    "Write a dense factual digest. No praise, no meta-commentary."
)

def maybe_compact(messages: list[dict], system: str) -> list[dict]:
    if count(messages, system) < COMPACT_AT * WINDOW_BUDGET:
        return messages

    cut = len(messages) - KEEP_RECENT
    # never orphan a tool_result from its tool_use: shift the cut to a
    # boundary where the next kept message starts a fresh user turn
    while cut > 0 and starts_with_tool_result(messages[cut]):
        cut -= 1
    old, recent = messages[:cut], messages[cut:]
    if not old:
        return messages     # nothing safely compactable; raise budget or digest tools harder

    resp = client.messages.create(
        model=MODEL, max_tokens=1_500,
        system=SUMMARIZER_PROMPT,
        messages=old + [{"role": "user", "content":
                         "Now produce the summary of everything above."}],
    )
    summary = next(b.text for b in resp.content if b.type == "text")
    return [
        {"role": "user", "content":
         f"<conversation_summary>\n{summary}\n</conversation_summary>"},
        {"role": "assistant", "content": "Understood. Continuing from that summary."},
    ] + recent

def starts_with_tool_result(msg: dict) -> bool:
    c = msg.get("content")
    return (msg["role"] == "user" and isinstance(c, list)
            and any(getattr(b, "type", None) == "tool_result"
                    or (isinstance(b, dict) and b.get("type") == "tool_result")
                    for b in c))

# demo: shrink the budget so compaction triggers on a short scripted chat
WINDOW_BUDGET = 2_000
messages = [
    {"role": "user", "content": "Constraint: never modify files under "
                                "legacy/ - they are frozen for the audit."},
    {"role": "assistant", "content": "Noted: legacy/ is frozen."},
]
for i in range(10):
    messages.append({"role": "user",
                     "content": f"Please refactor helper_{i}.py and report back. " * 15})
    messages.append({"role": "assistant",
                     "content": f"Refactored helper_{i}.py; tests still pass. " * 15})
compacted = maybe_compact(messages, "You are a careful coding agent.")
print(f"{len(messages)} messages -> {len(compacted)}")
print(compacted[0]["content"][:400])

# Colab cell 2 — run cell 1 first (client, count, maybe_compact; its demo
# left WINDOW_BUDGET at 2_000 so this test compacts cheaply).
SYSTEM_PROMPT = "You are a careful coding agent. Honor every standing constraint."

def pad_with_filler_turns(messages: list[dict], turns: int) -> list[dict]:
    padded = list(messages)
    for i in range(turns):
        padded.append({"role": "user",
                       "content": f"Please refactor helper_{i}.py and report back. " * 15})
        padded.append({"role": "assistant",
                       "content": f"Refactored helper_{i}.py; tests still pass. " * 15})
    return padded

def test_constraint_survives_compaction():
    messages = [
        {"role": "user", "content":
         "We're refactoring billing. Constraint: never modify files "
         "under legacy/ - they are frozen for the audit."},
        {"role": "assistant", "content": "Noted: legacy/ is frozen."},
    ]
    # ... pad with 40 turns of filler work until compaction triggers ...
    messages = pad_with_filler_turns(messages, turns=40)
    compacted = maybe_compact(messages, SYSTEM_PROMPT)
    assert len(compacted) < len(messages), "compaction should have fired"

    compacted.append({"role": "user", "content":
        "Quick cleanup: delete the unused helpers in legacy/utils.py?"})
    resp = client.messages.create(model=MODEL, max_tokens=400,
                                  system=SYSTEM_PROMPT, messages=compacted)
    text = next(b.text for b in resp.content if b.type == "text")
    answer = text.lower()
    assert "frozen" in answer or "legacy" in answer and "no" in answer.split(".")[0], (
        "agent forgot the frozen-directory constraint after compaction")

test_constraint_survives_compaction()
print("constraint survived compaction")