feat: context compaction strategies for the react loop

[yelkurdi]

Component PR

Use this template when adding or modifying components in mellea/stdlib/components/.

Description

• Link to Issue: Adds feature Compaction #1099.

Implementation Checklist

Protocol Compliance

• parts() returns list of constituent parts (Components or CBlocks)
• format_for_llm() returns TemplateRepresentation or string
• _parse(computed: ModelOutputThunk) parses model output correctly into the specified Component return type

Content Blocks

• CBlock used appropriately for text content
• ImageBlock used for image content (if applicable)

Integration

• Component exported in mellea/stdlib/components/__init__.py or, if you are adding a library of components, from your sub-module

Testing

• Tests added to tests/components/
• New code has 100% coverage
• Ensure existing tests and github automation passes (a maintainer will kick off the github automation when the rest of the PR is populated)

Attribution

• AI coding assistants used

Summary

Adds a generic, sync Compactor protocol to mellea.stdlib.context for shrinking a running Context, with three implementations — WindowCompactor, ThresholdCompactor, LLMSummarizeCompactor — plus composable pin predicates that decide which leading components survive. ReACT gains an optional compactor= parameter; ChatContext gains compactor= / window_size= for per-append compaction.

Empirically on the BCP benchmark with Granite 4.1-8b, LLM-summary compaction cuts inference cost by 23.7% and raises accuracy by 3.5 pp — a dual win, not a cost/quality trade-off. Backwards compatible: the defaults preserve existing behavior exactly.

Motivation

Long agentic loops pile up tool responses. Each react iteration re-sends the full history, so prompt-token cost grows quadratically and the loop can exhaust the context window before finishing. Compaction trims that history, cutting both dollar cost and the chance of hitting the context / timeout wall.

On BrowseCompPlus (BCP) with Granite 4.1-8b (131K context, 830 questions, loop_budget=400, per_question_timeout_s=1800, mean of 3 runs/strategy):

Compaction	Est. cost @ 80% cache hit	Accuracy	Mean wall-clock per Q
none	$1801.0M	15.6%	724 s
clear_all	$1499.2M (−16.8%)	18.5% (+2.9 pp)	980 s
llm_summarize	$1373.9M (−23.7%)	19.1% (+3.5 pp)	709 s

llm_summarize is the clear winner — cost −23.7% and accuracy +3.5 pp (+23% relative) at similar wall-clock, since the smaller prompt makes each iteration faster. clear_all also helps (−16.8% cost, +2.9 pp) but its +35% wall-clock is less attractive when the budget allows the extra summarization call.

Benchmark labels are the external eval-harness strategy names: clear_all → WindowCompactor(size=0, pin_predicate=pin_react_initiator), llm_summarize → LLMSummarizeCompactor(pin_predicate=pin_react_initiator), both gated via ThresholdCompactor.

Hardware / infrastructure

• Model: ibm-granite/granite-4.1-8b (bf16, native 131K context)
• Node: single 8× NVIDIA H100 80 GiB host (IBM Bluevela LSF cluster, exclusive GPU allocation)
• Inference: vLLM 0.19.1, 7 instances × TP=1 (1 GPU each), --enable-prefix-caching on. GPU 0 reserved for the BCP local search service.
• Agent: Mellea react via OpenAIBackend → local vLLM, concurrency=56, loop_budget=400, per_question_timeout_s=1800, --compaction_threshold 50000 (~38% of 131K), --compaction_keep_n 5.
• Judge: openai/gpt-oss-120b on the same node after agent teardown.
• All 830 BCP questions; 3 independently seeded runs per strategy (means reported).

Measured with a separate BCP eval harness running on Mellea's react framework; the feature itself has no external runtime dependency.

Design

class Compactor(Protocol):
    def compact(self, ctx: T, *, backend: Backend | None = None) -> T: ...

A Compactor returns a fresh, compacted copy of a context (never mutates the input). Sync and generic in T (a Context subtype), so a chat compactor narrows to ChatContext -> ChatContext. Backend-calling compactors hide the async work behind the sync method. Two usage patterns: per-append (ChatContext(compactor=...), runs on every add()) and per-turn / manual (react(compactor=...) or a direct compact(ctx, backend=...) between turns).

InlineCompactor marker — ChatContext.add() calls compact() without a backend on every append, so ChatContext(compactor=...) accepts only InlineCompactor instances (structural Compactor satisfaction isn't enough), steering backend-calling compactors to react(compactor=...), ThresholdCompactor, or manual compact().

Pin predicates — a PinPredicate returns the index after the protected prefix (components[:idx] preserved, components[idx:] is the body): pin_nothing, pin_system, pin_system_and_initial_user, and pin_react_initiator (keeps the react goal + tool registration).

Compactors:

• WindowCompactor(size, pin_predicate=pin_system) (inline) — keep prefix + last size body components; size=0 drops the body ("clear all").
• ThresholdCompactor(inner, threshold) (inline) — forwards to inner.compact only when measured token size exceeds threshold; the composable token-gate.
• LLMSummarizeCompactor(default_backend, keep_n=5, pin_predicate=pin_nothing, ...) — summarize older body into one [CONTEXT SUMMARY] message, keep last keep_n verbatim. Not inline (calls the backend each time); use via react/ThresholdCompactor/manual.

Token size (_last_usage_tokens) is read from the most recent ModelOutputThunk.generation.usage total_tokens. Helper react_summary_prompt(goal, max_tokens_hint) builds a research-flavoured template for LLMSummarizeCompactor.

Integration — in react, compaction runs once per turn after the is_final check so terminal turns skip it (saves a wasted summary call). Custom compactors used with react must return a ChatContext.

Non-goals

No mellea/core/ changes. The single-module mellea/stdlib/context.py is refactored into a mellea/stdlib/context/ package (chat.py, simple.py, compactor.py); public imports from mellea.stdlib.context are unchanged.

Test plan

• uv run pytest test/stdlib/test_compactor.py — 62 tests, ~12 s, no external services
• uv run pytest test/stdlib/frameworks/test_react_framework.py — compactor= integration

Covers each compactor's output/prefix preservation, all pin predicates, threshold gating (below/at/above, <=0 disables, no-thunk-yet), window sizing incl. size=0 and the window_size= sugar, the InlineCompactor marker check, and LLMSummarizeCompactor rendering + best-effort error handling (backend failure → unchanged; programming errors propagate). Runnable examples in docs/examples/context/.

Pitfalls

1. No generation.usage → token-gated compaction silently no-ops. All mainline backends populate it (AGENTS.md §5); a new backend without it disables ThresholdCompactor.
2. The token count lags by one turn. The threshold is checked against the token usage reported by the last completed model call, so anything added to the context after that call — most notably the current turn's tool responses — isn't counted until the next model call runs. In practice this is negligible, unless a single tool response is large enough to blow past the threshold on its own.
3. Sync→async bridge in LLMSummarizeCompactor (#1258). compact() is sync but must call an async backend, so it bridges through _run_coro_blocking (a worker-thread asyncio.run). Two known failure modes: (a) inside a running loop (e.g. react()) the worker blocks the calling thread, stalling the outer loop's other tasks — callbacks, telemetry flushers, cancellation — until the summary finishes; (b) a backend whose async client (e.g. httpx.AsyncClient) is bound to mellea's shared loop can't be used from the fresh worker loop, surfacing as RuntimeError: This event loop is already running or a hung request. Both are flagged in the docstrings but can't be fixed without a protocol change. The resolution tracked in Async Compactor protocol to replace sync↔async bridge in LLMSummarizeCompactor #1258 is an async acompact() on the Compactor protocol that react() awaits directly — no bridge, no loop blocking, no cross-loop client hazard.
4. Absolute token threshold, not a % of max context — a fraction-of-context threshold is a possible follow-up once backends report the limit reliably.
5. Compaction fires after the is_final check — easy to swap in a refactor; guarded by a comment, ideally a placement test later.

Files

 mellea/stdlib/context/compactor.py              613 +  Compactor protocol, pin predicates, 3 compactors
 mellea/stdlib/context/{__init__,chat,simple}.py 199 +  package split + compactor=/window_size= on ChatContext
 mellea/stdlib/context.py                          82 -  (replaced by the package above)
 mellea/stdlib/components/react.py                 96 +  pin_react_initiator, react_summary_prompt
 mellea/stdlib/frameworks/react.py                 22 ±  optional compactor= integration
 mellea/stdlib/{session,plugins/hooks/session}.py  35 ±  interaction_count tracked out-of-band
 test/stdlib/test_compactor.py                    824 +  (62 tests)
 test/stdlib/frameworks/test_react_framework.py   224 +
 docs/examples/context/*.py                              4 runnable examples + README

[yelkurdi]

please add @ramon-astudillo as an observer

[psschwei]

(original PR body)

Summary

Adds an optional CompactionStrategy to mellea.stdlib.frameworks.react with three concrete implementations (ClearAll, KeepLastN, LLMSummarize) under a new module mellea.stdlib.compaction. Strategies fire when the running context's token count crosses a configurable threshold, measured from the provider-reported usage on the last ModelOutputThunk.

Empirically on the BCP benchmark with Granite 4.1-8b, llm_summarize cuts inference cost by 23.7% and raises accuracy by 3.5 pp — compaction is a dual win, not a quality/cost trade-off.

Backwards compatible: compaction=None (default) preserves existing react() behavior exactly.

Motivation

Long agentic loops — especially retrieval-heavy ones — pile up tool responses. Each react iteration re-sends the full history to the model, so prompt-token cost grows quadratically and the loop can exhaust the model's context window before reaching a final answer. Compaction trims that history, lowering both the dollar cost of inference and the likelihood of hitting the context / timeout wall.

On the BrowseCompPlus (BCP) benchmark with Granite 4.1-8b (131K context, 830 questions, loop_budget=400, per_question_timeout_s=1800, averaged across 3 runs per strategy):

Compaction	Est. cost @ 80% cache hit	Accuracy	Mean wall-clock per Q
none	$1801.0M	15.6%	724 s
clear_all	$1499.2M (−16.8%)	18.5% (+2.9 pp)	980 s
llm_summarize	$1373.9M (−23.7%)	19.1% (+3.5 pp)	709 s

• Without compaction, inference cost on the 830-Q set is ~$1801M under an 80%-cache-hit pricing model, because each react iteration re-sends an ever-larger prompt (history + tool outputs) — the prompt-token total balloons to 5.4B before the 1800 s wall / context limit stops further progress.
• llm_summarize is the clear winner: it cuts inference cost by −23.7% ($1801M → $1374M) and raises accuracy from 15.6% to 19.1% (+3.5 pp, +23% relative), at similar mean wall-clock. The summary shrinks the prompt enough that each iteration is actually faster, so both cost and quality improve.
• clear_all reduces cost by −16.8% and lifts accuracy +2.9 pp, but its +35% wall-clock makes it less attractive than llm_summarize when the judge budget allows the extra summarization call.

Hardware / infrastructure for the table above:

• Model: ibm-granite/granite-4.1-8b (bf16, native 131K context)
• Node: single 8× NVIDIA H100 80 GiB host (IBM LSF cluster, exclusive GPU allocation)
• Inference: vLLM 0.19.1, 7 instances × TP=1 (1 GPU each), with --enable-prefix-caching on (Granite-family default). GPU 0 reserved for the BCP local search service (tevatron + BM25 over the BCP corpus).
• Agent: Mellea react loop via OpenAIBackend → local vLLM, concurrency=56 (8 × num_vllm), loop_budget=400, per_question_timeout_s=1800. Threshold: --compaction_threshold 50000 tokens (~38% of 131K), --compaction_keep_n 5 where applicable.
• Questions: all 830 from the BCP parquet, decrypted at load.
• Judge: openai/gpt-oss-120b on the same node after the agent phase teardown. "Correct" = judge verdict matches the reference answer.
• 3 runs per strategy, numbers in the table are means.

Measurement was done with a separate BCP eval harness (forthcoming PR); the harness data is included here as motivating evidence — the compaction feature itself has no external runtime dependency.

Design

Protocol — CompactionStrategy

class CompactionStrategy(abc.ABC):
    def __init__(self, *, threshold: int = 0) -> None: ...
    def should_compact(self, context: ChatContext) -> bool: ...
    async def maybe_compact(
        self, context, *, backend=None, goal=None
    ) -> ChatContext: ...
    @abc.abstractmethod
    async def compact(
        self, context, *, backend=None, goal=None
    ) -> ChatContext: ...

Threshold is compared against total_tokens from the most recent ModelOutputThunk.usage — i.e. the prompt+completion of the last LLM call.

Concrete strategies

• ClearAll(threshold) — discard everything after the first ReactInitiator, keeping only the system prefix. Cheapest, most aggressive; model must rebuild context each cycle.
• KeepLastN(keep_n, threshold) — retain prefix + the last keep_n body components. Middle-ground; preserves recent tool outputs.
• LLMSummarize(keep_n, threshold) — summarize the older body components via an additional LLM call, keep the last keep_n verbatim. Highest-fidelity, most expensive per-fire; empirically the best for Granite 4.1 on BCP (see table).

Integration site (react.py, +13 lines)

async def react(
    ..., compaction: CompactionStrategy | None = None,
) -> tuple[ComputedModelOutputThunk[str], ChatContext]:
    ...
    while turn_num < loop_budget or loop_budget == -1:
        step, next_context = await mfuncs.aact(...)
        ...
        if is_final:
            return step, context

        # Compact AFTER the final-answer check so terminal turns skip it.
        if compaction is not None:
            context = await compaction.maybe_compact(
                context, backend=backend, goal=goal
            )

Design decision — compact after the is_final check: a terminal turn has no next iteration to benefit from compaction, and for LLMSummarize this saves a full LLM call per question that would otherwise be wasted. Flagged inline so future refactors don't regress it.

Non-goals

• No changes to mellea/core/ — the feature is purely additive under mellea/stdlib/, reusing the existing ChatContext + ModelOutputThunk surfaces.

Test plan

• uv run pytest test/stdlib/test_compaction.py — 26 tests, 11 s locally, no external services required

The test file uses DummyThunk with synthetic usage dicts — no real backend needed. Coverage includes:

• Each strategy's compact() output shape
• Token-count comparison (below, at, above threshold)
• threshold=0 disables compaction
• Empty-context / no-thunk-with-usage returns False
• Prefix preservation (first ReactInitiator never dropped)
• LLMSummarize error handling when backend/goal omitted

Pitfalls (flagged here so reviewers know what to watch for)

1. Backends that don't populate mot.usage silently disable compaction. All mainline Mellea backends set it (OpenAI, HF, Ollama, LiteLLM), and AGENTS.md §5 codifies this as a requirement. If a new backend lands without usage population, its users will see compaction become a no-op with no loud error.

2. One-turn lag in the token-count measurement. The count reflects the prompt+completion of the LLM call that just completed — tool responses appended since are not yet counted. In practice negligible (a typical tool response is <5K tokens relative to a 50K+ threshold). Becomes visible only if a single tool response is very large (e.g. a raw document dump). Documented in the compaction.py docstring.

3. _last_usage_tokens returns None before the first model call. should_compact then returns False — no-op, not an error. Matters for strategies with very low thresholds where the first LLM call itself crosses the bar.

4. LLMSummarize needs backend + goal at compact() time. These are forwarded from the react call site. If a user constructs LLMSummarize and calls compact() directly (outside react), they must pass both. The docstring says so; reviewers may prefer a runtime check.

5. Strategies that drop all body components can leave a thunk-less context. On the next check, _last_usage_tokens returns None and compaction correctly doesn't re-fire immediately — behavior we want, but worth verifying didn't regress.

6. Token-count threshold is absolute, not a percentage of max context. A possible enhancement is to express the threshold as a percentage of the model's max context length rather than an absolute token count. This would require backends to reliably report that limit (e.g. 131K for Granite 4.1), after which --compaction_threshold 0.5 would read as "fire at 50% of context".

7. Compaction firing point is after the is_final check, not before. Easy to accidentally swap in a future refactor. Code comment calls this out; ideally a regression test guards it, but the current tests focus on the strategies themselves rather than react-loop placement. Worth adding if the reviewer flags it.

Files

 mellea/stdlib/frameworks/react_compaction.py
 mellea/stdlib/frameworks/react.py
 test/stdlib/frameworks/test_react_compaction.py
 1 file modified, 2 added

[github-actions]

The PR description has been updated. Please fill out the template for your PR to be reviewed.

[psschwei]

@yelkurdi I updated the PR body so that the update-pr-body check would pass (and copied your original body into a comment). Now that the check is passing, feel free to re-edit to include your original comments in the appropriate section.

[ramon-astudillo]

relevant related discussion #1099

[yelkurdi]

Will update according to #1099

[github-actions]

This comment is managed by a bot. Editing it is fine — checking off boxes, adding notes — but please leave the HTML comment marker on the first line alone, otherwise checklist updates will break.

Component PR Checklist

Use this checklist when adding or modifying components in mellea/stdlib/components/.

Protocol Compliance

• parts() returns list of constituent parts (Components or CBlocks)
• format_for_llm() returns TemplateRepresentation or string
• _parse(computed: ModelOutputThunk) parses model output correctly into the specified Component return type

Content Blocks

• CBlock used appropriately for text content
• ImageBlock used for image content (if applicable)

Integration

• Component exported in mellea/stdlib/components/__init__.py or, if you are adding a library of components, from your sub-module

[planetf1]

Thanks for this — the Compactor architecture is really well thought through, and the test coverage in test_compactor.py is thorough. A few things to sort out before this lands, mostly around LLMSummarizeCompactor edge cases and a couple of quick doc fixes.

[yelkurdi]

@planetf1 I have incorporated your feedback, I've also ran the tests and the full BCP deep search task tests (3 runs of each compaction method), all checks out.

[planetf1]

Second pass after the latest fixes — most of the earlier thread was addressed cleanly, thanks. A few new things from this pass, plus one CI-blocker in the docs example.

Heads up — PR is currently in DIRTY merge state (conflicts against main). Worth a rebase before final review.

Findings inline; one separate top-level note on _run_coro_blocking as a follow-up suggestion (not reopening the resolved thread — your docstring update there is solid).

[planetf1]

Following up on your docstring update on _run_coro_blocking from earlier — wanted to flag this as a separate suggestion rather than reopening the old thread, since the docstring itself is solid.

Have you considered routing the coroutine through mellea.helpers.event_loop_helper._run_async_in_thread? It's mellea's shared sync↔async bridge — used by session.py, functional.py, mcp.py, tools.py, manager.py, ~15 call sites — and it dispatches onto the same event loop where backends create their httpx.AsyncClient.

The benefit is mostly for the sync-from-outside-loop case (someone calling compactor.compact(ctx, backend=…) from a plain script): routing through the shared loop sidesteps the "client bound to a different loop" failure your warning paragraph predicts. The inside-react case isn't fixed by either helper — both fall back to a fresh-loop child when called from the running loop — so your "async Compactor protocol long-term" call-out is still the real answer.

Totally fine as a follow-up issue if you'd rather not expand scope here.

[planetf1]

LGTM. Given the recent discussions in #1099 I'll defer to @nrfulton for final say, but I'm happy with the changes.

[planetf1]

LGTM

[markstur]

Thanks for addressing or replying to everything. I'll accept your explanation for the event loop blocking concern, but I think we need a follow-up issue and need Nathan's approval.

[yelkurdi]

Thanks @planetf1 and @markstur , the follow issue is created to handle the loop blocking concern: #1258

[planetf1]

Fix looks good, but can we rebase to current code and ensure no conflict issues.

[psschwei]

Comparing against the summary of our previous discussion in #1099, in my opinion the PR conforms to that design. There's a couple of things we need to do before we can merge:

• rebase / resolve the merge conflicts
• update the PR description to match the new design

[psschwei]

DCO check is failing on this commit: 1e28704

@ramon-astudillo I think this is yours.

[ramon-astudillo]

This commit will get squashed anyway, but let me check with @yelkurdi

[yelkurdi]

@psschwei I resolved the conflicts and now I'm in the process of updating the PR description.
Regarding the DCO error, we manually resolved that in the past. I can fix it by rebasing, or squashing, will handle it.

[yelkurdi]

PR description updated.

[planetf1]

Third pass — the design work here is solid and the earlier findings are all cleaned up. One thing that needs fixing before merge: six fake Backend subclasses (five in test_compactor.py, one in react_compaction.py) override generate_from_raw instead of the abstract _generate_from_raw that #1264 introduced. The whole TestLLMSummarizeCompactor class errors at fixture setup as a result — details in the inline comment. A rename in each fake unblocks it.

CI is now running and the failing quality check is the active gate — no need for a separate review block on top of it.

[psschwei]

looks like we have another merge conflict (and the CI seems to have gotten stuck, but fixing the conflicts should give it a kick)

All findings addressed.