# Plan: Graceful Context Window Management

## Problem

When the conversation history exceeds the model's context window, `call_llm()` raises a LiteLLM exception (typically `ContextWindowExceededError`, status 400) and the agent crashes with exit code 1. There is no recovery, no pruning, and no visibility into how much context is being consumed.

## Goals

1. Catch context overflow errors and recover by compacting/trimming messages
2. Notify the user (on stderr) when compaction happens
3. In verbose mode, log estimated context size before and after each turn

## Current State

- Messages list grows unboundedly (mix of dicts and LiteLLM response objects)
- Per-tool-result output is capped at 50 KB, but many results accumulate
- System prompt is ~3.3 KB (fixed)
- No token counting anywhere
- `call_llm()` catches all exceptions with `sys.exit(1)`

## Design

### 1. Token counting with tiktoken

Add `tiktoken` as a dependency (`uv add tiktoken`). Use it to get accurate token counts for the message history.

Since the agent talks to local models via LM Studio (not OpenAI directly), there's no "correct" encoding to pick — the local model's tokenizer will differ. Use `cl100k_base` as the default: it's the standard GPT-4/3.5 encoding, widely representative of modern BPE tokenizers, and close enough for context budget decisions. The actual overflow is still detected by the API error; tiktoken just gives us good estimates for logging and could later enable proactive trimming.

Add an `estimate_tokens(messages, tools)` function in `agent.py`:

```python
import tiktoken

_encoder = tiktoken.get_encoding("cl100k_base")

def estimate_tokens(messages: list, tools: list | None = None) -> int:
    """Count tokens across all messages using tiktoken."""
    total = 0
    for m in messages:
        if isinstance(m, dict):
            content = m.get("content", "") or ""
        else:
            content = getattr(m, "content", "") or ""
            tool_calls = getattr(m, "tool_calls", None)
            if tool_calls:
                for tc in tool_calls:
                    content += tc.function.name + (tc.function.arguments or "")
        total += len(_encoder.encode(content))
    if tools:
        total += len(_encoder.encode(json.dumps(tools)))
    # Per-message overhead (role, separators) — ~4 tokens each
    total += 4 * len(messages)
    return total
```

The encoder is initialized once at module level (fast, no network calls). The per-message overhead of ~4 tokens accounts for role markers and message separators that the API adds internally.

### 2. Verbose context logging

Before calling `call_llm()` each turn, log the estimated token count:

```
--- Turn 3/50 (~12,400 tokens) ---
```

After getting the response and appending tool results, log the new size:

```
Context after turn 3: ~18,200 tokens
```

This goes through the existing `log()` function (stderr, verbose-only).

### 3. Detect context overflow errors

Define a custom exception and a classifier that catches both the typed LiteLLM exception and generic 400s whose message mentions context length:

```python
class ContextOverflowError(Exception):
    """Raised when the LLM call fails due to context window overflow."""
    pass
```

Change `call_llm()` to raise it:

```python
def call_llm(...):
    import litellm

    ...

    try:
        response = litellm.completion(...)
    except litellm.ContextWindowExceededError:
        raise ContextOverflowError("context window exceeded (typed)")
    except litellm.BadRequestError as e:
        # Some providers return a generic 400 with context-length keywords.
        # Use regex to avoid false positives on unrelated 400s.
        import re
        msg_text = str(e)
        _CONTEXT_OVERFLOW_RE = re.compile(
            r"context.{0,10}(length|window|limit)"
            r"|maximum.{0,10}(context|token)"
            r"|token.{0,10}limit"
            r"|exceed.{0,10}(context|token|max)",
            re.IGNORECASE,
        )
        if _CONTEXT_OVERFLOW_RE.search(msg_text):
            raise ContextOverflowError(f"context window exceeded (inferred): {e}")
        print(f"Error: LLM call failed: {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Error: LLM call failed: {e}", file=sys.stderr)
        sys.exit(1)

    ...
```

The agent loop catches `ContextOverflowError` only. This handles both LiteLLM's typed exception and the common case where providers surface overflow as a generic `BadRequestError` with context-length text in the message.

### 4. Turn atomicity invariant

The OpenAI message format requires that every tool-result message has a matching `tool_call_id` in a preceding assistant message's `tool_calls` array. Breaking this pairing produces a 400 error — which would defeat the purpose of compaction.

Both compaction functions must preserve **complete turns**: an assistant message with `tool_calls` and all of its corresponding tool-result messages are treated as an atomic unit. They are either all kept or all removed/truncated together.

Add a helper to group messages into turns:

```python
def group_into_turns(messages: list) -> list[list]:
    """Group messages into atomic turns.

    A turn is one of:
    - A single message (system, user, or assistant without tool_calls)
    - An assistant message with tool_calls + all its matching tool results
    """
    turns = []
    i = 0
    while i < len(messages):
        msg = messages[i]
        role = msg.get("role") if isinstance(msg, dict) else getattr(msg, "role", None)
        tool_calls = msg.get("tool_calls", None) if isinstance(msg, dict) else getattr(msg, "tool_calls", None)

        if role == "assistant" and tool_calls:
            # Collect this assistant msg + all following tool results
            turn = [msg]
            tc_ids = {tc.id if hasattr(tc, "id") else tc["id"] for tc in tool_calls}
            j = i + 1
            while j < len(messages):
                next_msg = messages[j]
                next_role = next_msg.get("role") if isinstance(next_msg, dict) else getattr(next_msg, "role", None)
                tc_id = next_msg.get("tool_call_id") if isinstance(next_msg, dict) else getattr(next_msg, "tool_call_id", None)
                if next_role == "tool" and tc_id in tc_ids:
                    turn.append(next_msg)
                    j += 1
                else:
                    break
            turns.append(turn)
            i = j
        else:
            turns.append([msg])
            i += 1
    return turns
```

All compaction operates on turns, then flattens back to a message list.

### 5. Compaction strategy

When `ContextOverflowError` is caught in the agent loop, compact the message history using this approach:

**Step 1 — Truncate large tool results.** Walk the turns and replace any tool-result `content` longer than 1000 characters with a summary:

```
[compacted — originally N chars]
```

Skip the most recent 2 turns (they're likely still relevant). The assistant message in each turn is left intact; only the tool-result content is truncated. This preserves turn structure.

**Step 2 — Drop oldest turns.** If step 1 isn't enough (i.e., the retry also hits context overflow), drop turns from the middle of the conversation. Retention is role-based, not positional:

- **Keep all leading system/user messages** (handles both `--no-system-prompt` and normal mode — scan forward while role is `system` or `user`)
- **Keep the last 3 turns** (recent context the model needs to continue)
- **Drop everything in between**, replacing with a synthetic user message:

```
[context compacted — older tool calls and results were removed to fit context window]
```

This avoids hardcoding index assumptions. The leading block is typically 1-2 messages (system + user question) but adapts to whatever bootstrap messages exist.

**Step 3 — Give up.** If after both compaction passes the LLM still raises context overflow, print an error and exit with code 1. This handles pathological cases (e.g., a single enormous system prompt that alone exceeds the window).

### 6. Output-budget clamping

`discover_model()` already returns `context_length`. Pass it into the agent loop. Before each `call_llm()`, if we know the context limit, clamp `max_output_tokens` so that `prompt_tokens + max_output_tokens <= context_length`:

```python
def clamp_output_tokens(messages, tools, context_length, requested_max_output):
    """Reduce max_output_tokens if prompt + output would exceed context."""
    if context_length is None:
        return requested_max_output
    prompt_tokens = estimate_tokens(messages, tools)
    available = context_length - prompt_tokens
    if available < 1:
        return requested_max_output  # Let the overflow error handle it
    return min(requested_max_output, available)
```

This is called before every `call_llm()` invocation, including retries after compaction. Without this, compaction can free prompt space but the original `max_output_tokens` can still push the total over the limit, causing all three retries to fail identically.

When clamping reduces the output budget, log it in verbose mode:

```
Output tokens clamped: 16384 -> 8192 (context_length=32768, prompt=~24576)
```

### 7. Retry logic in the agent loop

Wrap the `call_llm()` invocation in the main loop:

```python
effective_max_output = clamp_output_tokens(messages, tools, context_length, args.max_output_tokens)

try:
    msg, finish_reason = call_llm(..., max_output_tokens=effective_max_output, ...)
except ContextOverflowError:
    print("Warning: context window exceeded, compacting history...", file=sys.stderr)
    messages = compact_messages(messages)
    effective_max_output = clamp_output_tokens(messages, tools, context_length, args.max_output_tokens)
    log(f"Context after compaction: ~{estimate_tokens(messages, tools)} tokens", verbose)
    try:
        msg, finish_reason = call_llm(..., max_output_tokens=effective_max_output, ...)
    except ContextOverflowError:
        print("Warning: still too large, dropping older turns...", file=sys.stderr)
        messages = drop_middle_turns(messages)
        effective_max_output = clamp_output_tokens(messages, tools, context_length, args.max_output_tokens)
        log(f"Context after drop: ~{estimate_tokens(messages, tools)} tokens", verbose)
        try:
            msg, finish_reason = call_llm(..., max_output_tokens=effective_max_output, ...)
        except ContextOverflowError:
            print("Error: context window exceeded even after compaction.", file=sys.stderr)
            sys.exit(1)
```

Each retry re-clamps the output budget after compaction. Three tries with increasingly aggressive trimming, then give up.

### 8. User notification

All compaction warnings go to stderr (visible regardless of verbose mode):

- `"Warning: context window exceeded, compacting history..."` — first compaction
- `"Warning: still too large, dropping older turns..."` — aggressive trim
- `"Error: context window exceeded even after compaction."` — unrecoverable

In verbose mode, the token estimates before/after give full visibility.

## Dependencies

Add `tiktoken` via `uv add tiktoken`. It's OpenAI's official BPE tokenizer library — fast (3-6x faster than alternatives), no network calls at runtime, and the standard choice for token counting in Python LLM tooling.

## Files Changed

**agent.py** is the only source file modified. **pyproject.toml** gains the tiktoken dependency.

### New functions

| Function | Purpose |
|---|---|
| `estimate_tokens(messages, tools)` | Token count via tiktoken (`cl100k_base`) |
| `group_into_turns(messages)` | Group messages into atomic turns (assistant+tool_calls + matching tool results) |
| `compact_messages(messages)` | Truncate large tool results in older turns, preserving turn atomicity |
| `drop_middle_turns(messages)` | Drop oldest turns from middle; keep leading system/user block + last 3 turns |
| `clamp_output_tokens(...)` | Reduce `max_output_tokens` so prompt + output fits context window |

### New exception

| Class | Purpose |
|---|---|
| `ContextOverflowError` | Raised by `call_llm()` on context overflow (typed or inferred from 400 message) |

### Modified functions

| Function | Change |
|---|---|
| `call_llm()` | Catch `ContextWindowExceededError` and context-related `BadRequestError`, raise `ContextOverflowError` |
| `main()` agent loop | Add try/except with compaction fallback; output-budget clamping; token logging |

## Testing

Add `tmp/test_compact.py` with:

### Unit tests

- `test_estimate_tokens` — verify tiktoken-based counting on known inputs; check that tool_calls content is included
- `test_group_into_turns_basic` — single messages become single-element turns
- `test_group_into_turns_tool_calls` — assistant with tool_calls + matching tool results form one turn
- `test_group_into_turns_partial` — tool results without a preceding assistant are kept as standalone (defensive)
- `test_compact_messages_truncates_large_results` — tool results over 1000 chars get replaced; recent 2 turns are untouched
- `test_compact_preserves_turn_atomicity` — after compaction, every tool result still has its assistant message and vice versa
- `test_drop_middle_turns_keeps_boundaries` — leading system/user messages and last 3 turns survive; middle is replaced with splice marker
- `test_drop_middle_turns_no_system` — works correctly when there's no system message (`--no-system-prompt`)
- `test_drop_preserves_turn_atomicity` — after dropping, no orphaned tool results or assistant tool_calls exist
- `test_drop_middle_turns_small_history` — when the history is too small for a middle to exist (leading block + last 3 turns overlap or cover everything), `drop_middle_turns` returns the messages unchanged with no splice marker inserted
- `test_clamp_output_tokens` — verify clamping math: prompt_tokens + result <= context_length; returns original when context_length is None

### Integration-level tests

- `test_compact_then_valid_for_api` — build a realistic message sequence (system, user, multiple assistant+tool turns), compact it, then validate the result would pass OpenAI message format validation: every tool result has a matching tool_call_id in a preceding assistant message
- `test_drop_then_valid_for_api` — same validation after `drop_middle_turns`
- `test_context_overflow_classifier` — verify that `call_llm()` raises `ContextOverflowError` for both `ContextWindowExceededError` and `BadRequestError` with context-length keywords (mock LiteLLM exceptions)

## What This Does NOT Do

- **Proactive trimming** — This is reactive (triggers on error), not preemptive. The output-budget clamping prevents one class of avoidable overflow, but we don't preemptively compact when approaching the limit. With tiktoken and `context_length` both available, a threshold-based pre-trim could be added later.
- **Summarization** — No LLM-based summarization of old messages. That would require an extra LLM call mid-loop and could itself fail. The simple truncate/drop approach is predictable.
- **Exact tokenizer match** — tiktoken's `cl100k_base` won't exactly match whatever tokenizer the local model uses, but it's close enough for budget estimation and output clamping. The actual overflow is still detected by the API error.
