Engenharia de contexto para agentes de longa duração: orçamento, limpeza, compactação e memória

python -m venv .venv && . .venv/bin/activate
pip install tiktoken pytest

"""Token counting for context budgets.
 
Uses tiktoken's o200k_base encoding when it is installed. That is an OpenAI
tokenizer, not Claude's or Gemini's, so the counts are estimates: good for
budgets and trends, not for reconciling an invoice (use the provider's
token counting endpoint for that). Without tiktoken it falls back to
len(text) / 4, the usual rule of thumb for English prose and code.
"""
 
from __future__ import annotations
 
from functools import lru_cache
 
# Role markers and framing the chat format adds around every message.
# Approximate: providers don't document it and it varies by model.
MESSAGE_OVERHEAD = 4
 
try:
    import tiktoken
 
    _encoding = tiktoken.get_encoding("o200k_base")
    BACKEND = "tiktoken o200k_base"
 
    def _count(text: str) -> int:
        return len(_encoding.encode(text, disallowed_special=()))
 
except ImportError:
    BACKEND = "chars / 4"
 
    def _count(text: str) -> int:
        return max(1, len(text) // 4)
 
 
@lru_cache(maxsize=8192)
def count_tokens(text: str) -> int:
    return _count(text)

from __future__ import annotations
 
import hashlib
import re
from dataclasses import dataclass, field, replace
from typing import Callable, Literal, Protocol
 
from ctxkit.tokens import MESSAGE_OVERHEAD, count_tokens
 
Role = Literal["system", "user", "assistant", "tool", "summary"]
 
ERROR_LINE = re.compile(r"\b([A-Z]\w*(?:Error|Exception)): ([^\n\]]+)")
 
 
@dataclass(frozen=True)
class Message:
    role: Role
    content: str
    turn: int = 0
    tool: str | None = None  # tool results: which tool produced them
    args: str | None = None  # and with which arguments, kept after clearing
    cleared: bool = False
 
    @property
    def tokens(self) -> int:
        return count_tokens(self.content) + MESSAGE_OVERHEAD
 
    @property
    def key(self) -> str:
        return hashlib.sha256(f"{self.role}\0{self.content}".encode()).hexdigest()
 
 
class Summarizer(Protocol):
    def summarize(self, messages: list[Message], notes: str) -> str: ...
 
 
class ContextOverflow(RuntimeError):
    pass
 
 
@dataclass(frozen=True)
class Budget:
    window: int = 200_000  # the model's hard limit
    trigger: int = 30_000  # above this, the policies run
    clear_at_least: int = 10_000  # never break the cache for less than this
    keep_tool_results: int = 3  # newest tool results always stay verbatim
    compact_above: int = 20_000  # still above this after clearing: summarize
    keep_recent: int = 10  # messages a compaction keeps word for word
    pinned_tools: frozenset[str] = frozenset({"notes_read"})
 
 
def stub(message: Message) -> str:
    """What replaces a stale tool result: the call survives, the payload doesn't."""
    text = f"[cleared: {message.tool}({message.args}) returned {message.tokens:,} tokens. Call it again if you need it."
    if error := ERROR_LINE.search(message.content):
        text += f" First error: {error.group(0)}"
    return text + "]"
 
 
@dataclass
class ContextManager:
    system: str
    summarizer: Summarizer
    budget: Budget = field(default_factory=Budget)
    read_notes: Callable[[], str] = lambda: ""
    history: list[Message] = field(default_factory=list)
    events: list[str] = field(default_factory=list)
    requests: int = 0
 
    def append(self, message: Message) -> None:
        self.history.append(message)
 
    def total(self) -> int:
        return count_tokens(self.system) + MESSAGE_OVERHEAD + sum(m.tokens for m in self.history)
 
    def prepare(self) -> list[Message]:
        """Apply the policies, cheapest first, and return the next request."""
        self.requests += 1
        if self.total() > self.budget.trigger:
            self._clear_tool_results()
            if self.total() > self.budget.compact_above:
                self._compact()
        if self.total() > self.budget.window:
            raise ContextOverflow(f"{self.total():,} tokens after every policy ran")
        return [Message("system", self.system), *self.history]
 
    def _clear_tool_results(self) -> None:
        results = [
            i
            for i, m in enumerate(self.history)
            if m.role == "tool" and not m.cleared and m.tool not in self.budget.pinned_tools
        ]
        stale = results[: -self.budget.keep_tool_results or None]
        stubs = {i: replace(self.history[i], content=stub(self.history[i]), cleared=True) for i in stale}
        freed = sum(self.history[i].tokens - s.tokens for i, s in stubs.items())
        # Every edit invalidates the cached prefix from that point on, so
        # small edits cost more in cache writes than they save.
        if freed < self.budget.clear_at_least:
            return
        for i, cleared in stubs.items():
            self.history[i] = cleared
        self.events.append(f"request {self.requests}: cleared {len(stale)} tool results (-{freed:,})")
 
    def _compact(self) -> None:
        cut = len(self.history) - self.budget.keep_recent
        # A tool result must never lose the call that produced it.
        while cut > 0 and self.history[cut].role == "tool":
            cut -= 1
        old, recent = self.history[:cut], self.history[cut:]
        # Same rule as clearing: a summary rewrites the prefix, so it must pay for itself.
        if sum(m.tokens for m in old) < self.budget.clear_at_least:
            return
        before = self.total()
        summary = self.summarizer.summarize(old, self.read_notes())
        self.history = [Message("summary", summary, turn=old[-1].turn), *recent]
        self.events.append(f"request {self.requests}: compacted {len(old)} messages (-{before - self.total():,})")

from __future__ import annotations
 
import re
 
from ctxkit.context import ERROR_LINE, Message
 
MARKER = re.compile(r"^(TASK|USER|DECISION|TODO|DONE|ERROR|FILE): (.+)$", re.M)
PATH = re.compile(r"\b(?:src|tests)/[\w/.-]+\.py\b")
SEEN = re.compile(r"^(.*) \(last seen turn (\d+)\)$")
 
 
def _touch(d: dict, key, value) -> None:
    d.pop(key, None)  # re-insert so the dict stays ordered by recency
    d[key] = value
 
 
class ExtractiveSummarizer:
    """Deterministic and offline: keeps what later turns tend to need.
 
    It relies on a convention the system prompt asks for: the agent writes
    DECISION:, TODO: and DONE: lines in its messages. Its own output uses
    the same markers, so summarizing a summary loses nothing.
    """
 
    def __init__(self, max_errors: int = 5, max_files: int = 12) -> None:
        self.max_errors = max_errors
        self.max_files = max_files
 
    def summarize(self, messages: list[Message], notes: str) -> str:
        task, users, decisions = "", [], []
        todos: dict[str, None] = {}
        errors: dict[str, int] = {}
        files: dict[str, str] = {}
        for m in messages:
            if m.role == "user":
                if task:
                    users.append(f"(turn {m.turn}) {m.content}")
                else:
                    task = m.content
            elif m.role in ("assistant", "summary"):
                turn = "" if m.role == "summary" else f"(turn {m.turn}) "
                for kind, text in MARKER.findall(m.content):
                    if kind == "TASK":
                        task = text
                    elif kind == "USER":
                        users.append(text)
                    elif kind == "DECISION":
                        decisions.append(turn + text)
                    elif kind == "TODO":
                        todos[text] = None
                    elif kind == "DONE":
                        todos.pop(text, None)
                    elif kind == "ERROR" and (seen := SEEN.match(text)):
                        _touch(errors, seen.group(1), int(seen.group(2)))
                    elif kind == "FILE":
                        _touch(files, PATH.findall(text)[0], text)
            elif m.role == "tool":
                for error in ERROR_LINE.finditer(m.content):
                    _touch(errors, error.group(0), m.turn)
                for path in PATH.findall(m.args or ""):
                    _touch(files, path, f"{path}, last {m.tool} at turn {m.turn}")
 
        lines = [f"TASK: {task}"]
        lines += [f"USER: {u}" for u in users]
        lines += [f"DECISION: {d}" for d in decisions]
        lines += [f"TODO: {t}" for t in todos]
        lines += [f"ERROR: {e} (last seen turn {t})" for e, t in list(errors.items())[-self.max_errors :]]
        lines += [f"FILE: {f}" for f in list(files.values())[-self.max_files :]]
        header = f"Summary of turns 1 to {messages[-1].turn}. Older tool output was dropped; re-read what you need."
        return "\n".join([header, *lines, "", "Current notes:", notes or "(empty)"])

from ctxkit.context import Message
 
PROMPT = """You are compacting the history of a coding agent so it can continue its task.
Keep, verbatim where possible:
- the original task and every instruction the user added later
- every decision and its reason, as DECISION: lines
- open work, as TODO: lines (drop items marked DONE:)
- exact error messages and approaches that failed, as ERROR: lines
- every file path touched and what was last done to it, as FILE: lines
Drop tool output the agent can fetch again. Never invent facts."""
 
 
class LLMSummarizer:
    def __init__(self, client, model: str, max_tokens: int = 2_000) -> None:
        self.client, self.model, self.max_tokens = client, model, max_tokens
 
    def summarize(self, messages: list[Message], notes: str) -> str:
        transcript = "\n\n".join(f"[{m.role}, turn {m.turn}] {m.content}" for m in messages)
        response = self.client.messages.create(
            model=self.model,
            max_tokens=self.max_tokens,
            system=PROMPT,
            messages=[{"role": "user", "content": transcript}],
        )
        return f"{response.content[0].text}\n\nCurrent notes:\n{notes or '(empty)'}"

class NotesFile:
    """Memory the agent writes and reads itself; it survives every compaction."""
 
    def __init__(self, path: Path, max_tokens: int = 1_500) -> None:
        self.path = path
        self.max_tokens = max_tokens
 
    def read(self) -> str:
        return self.path.read_text(encoding="utf-8") if self.path.exists() else ""
 
    def write(self, section: str, entry: str) -> str:
        sections = self._sections()
        sections.setdefault(section, []).append(entry)
        return self._save(sections, f"ok: {section} has {len(sections[section])} entries")
 
    def rewrite(self, section: str, entries: list[str]) -> str:
        sections = self._sections()
        sections[section] = entries
        if not entries:
            del sections[section]
        return self._save(sections, f"ok: {section} rewritten")
 
    def _sections(self) -> dict[str, list[str]]:
        sections: dict[str, list[str]] = {}
        for line in self.read().splitlines():
            if line.startswith("## "):
                sections[line[3:]] = []
            elif line.startswith("- ") and sections:
                sections[list(sections)[-1]].append(line[2:])
        return sections
 
    def _save(self, sections: dict[str, list[str]], ok: str) -> str:
        text = "\n\n".join(f"## {name}\n" + "\n".join(f"- {e}" for e in items) for name, items in sections.items())
        # Notes that grow without bound are just a second history.
        if count_tokens(text) > self.max_tokens:
            return f"error: notes would exceed {self.max_tokens} tokens; merge entries with notes_rewrite first"
        self.path.write_text(text + "\n", encoding="utf-8")
        return ok

@dataclass
class Meter:
    """Per-request tokens and cost, with an append-only prompt cache model."""
 
    previous: list[str] = field(default_factory=list)
    context: list[int] = field(default_factory=list)
    fresh: int = 0
    cost_plain: float = 0.0
    cost_cached: float = 0.0
    cost_cold: float = 0.0
 
    def request(self, messages: list[Message], output: int, turn: int) -> None:
        keys, tokens = [m.key for m in messages], [m.tokens for m in messages]
        shared = 0
        while shared < min(len(keys), len(self.previous)) and keys[shared] == self.previous[shared]:
            shared += 1
        read = sum(tokens[:shared]) if sum(tokens[:shared]) >= MIN_CACHEABLE else 0
        total = sum(tokens)
        self.context.append(total)
        self.fresh += total - read
        self.cost_plain += (total * PRICE_IN + output * PRICE_OUT) / 1e6
        self.cost_cached += self._cached(total, read, output)
        self.cost_cold += self._cached(total, 0 if turn in COLD else read, output)
        self.previous = keys
 
    @staticmethod
    def _cached(total: int, read: int, output: int) -> float:
        return (read * PRICE_IN * CACHE_READ + (total - read) * PRICE_IN * CACHE_WRITE + output * PRICE_OUT) / 1e6
 
 
def replay(turns, manager: ContextManager | None, notes: NotesFile) -> tuple[Meter, list[Message]]:
    meter, history = Meter(), []
    add = manager.append if manager else history.append
    request = manager.prepare if manager else lambda: [Message("system", SYSTEM), *history]
    for t, user, call, result in turns:
        if user:
            add(Message("user", user, turn=t))
        meter.request(request(), output=call.tokens, turn=t)
        if result.tool in ("notes_write", "notes_rewrite"):
            section, entry = result.args.split(" | ")
            ok = notes.write(section, entry) if result.tool == "notes_write" else notes.rewrite(section, [entry])
            result = Message("tool", ok, turn=t, tool=result.tool, args=result.args)
        add(call)
        add(result)
    return meter, request()

$ python simulate.py
tokenizer: tiktoken o200k_base, system prompt + tools: 612 tokens
turn     naive   managed
   1       649       649
   5     8,826     8,826
  10    24,950    24,950
  15    43,753    21,850
  20    63,403    22,804
  25    79,622    20,491
  30    93,418    15,497
  35   108,138    17,992
  40   118,983    28,837
  45   146,066    28,298
  50   176,486    33,041
  55   187,866    29,429
  60   212,592    20,965
 
manager events:
  request 12: cleared 8 tool results (-21,903)
  request 18: cleared 6 tool results (-18,696)
  request 24: cleared 6 tool results (-18,532)
  request 30: cleared 6 tool results (-18,790)
  request 35: cleared 5 tool results (-12,225)
  request 41: cleared 6 tool results (-17,048)
  request 44: cleared 3 tool results (-10,574)
  request 46: cleared 2 tool results (-15,562)
  request 49: cleared 3 tool results (-10,115)
  request 52: cleared 3 tool results (-14,992)
  request 56: cleared 4 tool results (-10,063)
  request 59: cleared 3 tool results (-13,393)
  request 59: compacted 108 messages (-9,734)
 
                                 naive     managed
input tokens, all turns      5,856,377   1,286,153
peak context                   212,592      36,936
tokens prefilled fresh         213,241     310,929
cost without caching            $17.71       $4.00
cost with caching                $2.63       $1.60
... and 2 cold caches            $3.33       $1.75
naive passes 128,000 tokens at turn 42
naive passes 200,000 tokens at turn 58
 
facts still in the managed context at the end:
  yes  task: keep USD behavior unchanged
  yes  decision t8: ROUND_HALF_EVEN
  yes  decision t16: never float
  yes  decision t35: format_money()
  yes  user t20: 3 decimals for BRL
  yes  error t13: TypeError: unsupported operand
  yes  file: src/billing/tax.py
  yes  open task: TODO: run the full suite

Summary of turns 1 to 53. Older tool output was dropped; re-read what you need.
TASK: Invoice totals in BRL are sometimes off by one cent. Find the cause, fix it, add regression tests, and keep USD behavior unchanged.
USER: (turn 20) Also check the CSV export, finance says it shows 3 decimals for BRL.
DECISION: (turn 8) round once, at the invoice total, with ROUND_HALF_EVEN in src/billing/money.py
DECISION: (turn 16) tax rates load as Decimal from strings, never float
DECISION: (turn 35) export formats money with format_money(), never str(Decimal)
TODO: run the full suite and summarize the change for the PR
ERROR: AssertionError: expected Decimal('10.01'), got Decimal('10.00') (last seen turn 11)
ERROR: TypeError: unsupported operand type(s) for *: 'float' and 'decimal.Decimal' (last seen turn 15)
ERROR: AssertionError: '10.010' != '10.01' (last seen turn 32)
FILE: src/billing/invoice.py, last read_file at turn 17
FILE: tests/test_export.py, last run_tests at turn 31
FILE: src/billing/export.py, last edit_file at turn 37
FILE: src/billing/tax.py, last read_file at turn 42
FILE: tests/test_money.py, last run_tests at turn 46
FILE: src/billing/currency.py, last read_file at turn 47
FILE: src/billing/money.py, last read_file at turn 48
FILE: tests/test_invoice.py, last read_file at turn 49
 
Current notes:
## Decisions
- round once at the total, ROUND_HALF_EVEN, money.py
- export uses format_money(), never str(Decimal)
 
## Open
- run the full suite and summarize the change for the PR

[cleared: run_tests(tests/test_invoice.py) returned 3,654 tokens. Call it again if you need it. First error: AssertionError: expected Decimal('10.01'), got Decimal('10.00')]

$ python -m pytest -v
collecting ... collected 5 items
 
test_ctxkit.py::test_clearing_keeps_calls_and_newest_results PASSED      [ 20%]
test_ctxkit.py::test_small_gains_do_not_break_the_cache PASSED           [ 40%]
test_ctxkit.py::test_compaction_never_orphans_a_tool_result PASSED       [ 60%]
test_ctxkit.py::test_summarizing_a_summary_loses_nothing PASSED          [ 80%]
test_ctxkit.py::test_notes_refuse_to_grow_without_bound PASSED           [100%]
 
============================== 5 passed in 0.32s ==============================

{
  "context_management": {
    "edits": [
      {
        "type": "clear_tool_uses_20250919",
        "trigger": { "type": "input_tokens", "value": 30000 },
        "keep": { "type": "tool_uses", "value": 3 },
        "clear_at_least": { "type": "input_tokens", "value": 10000 },
        "exclude_tools": ["notes_read"]
      }
    ]
  }
}