feat(context-engine): add select_context() + on_turn_complete() hooks (RFC #36765; consolidates #41918/#24949/#47109/#50053)

[chaos-xxl]

What does this PR do?

I'm the author of RFC #36765, and this is its implementation. The RFC didn't come from reading the ABC — it came from dogfooding a real context engine (Doctor Chaos, which routes each message to a topic "space" and uses that space's history as the turn's context). Running it on Hermes is what surfaced the gap this PR closes. So this isn't a tidy-up of other people's PRs — it's the concrete landing of a direction that came out of practice, and that a cluster of independent PRs then converged on.

It adds two optional, no-op-default hooks to the ContextEngine ABC — the orthogonal pair the ABC has no vocabulary for today:

• select_context() — pre-request selection. Called every turn after the request messages are assembled and before provider dispatch, independent of should_compress(). Lets an engine select/replace which context enters this one request (retrieval, topic routing, role/branch switching) without mutating persisted history.
• on_turn_complete() — post-turn observation. Called once after the assistant/tool loop finishes, with the finalized transcript snapshot and the turn's canonical token usage (prompt/completion + cache-read/write/reasoning buckets, or None when no provider response was reached). Lets an engine ingest/index/summarize what actually happened — and weigh how large/expensive the selected context was — so the next select_context() can act on it.

select_context()   : "what context should enter THIS request?"  (before)
on_turn_complete() : "what actually happened in the finished turn?" (after)
compress()         : "context is too long -> make it shorter"     (unchanged)

Both are additive, no-op by default, and fail-open. should_compress() / compress() semantics are unchanged; the built-in compressor and every existing engine are unaffected. This builds on the #33750 host contract, adding the two verbs it didn't cover.

Why this matters (from practice, not theory)

Without these hooks, an engine that needs per-turn message access has to force should_compress() to return True so compress() runs every turn purely as a callback. In my own dogfooding that backdoor caused a real correctness bug: when the engine's backend went down, Hermes' built-in compressor still engaged at its default threshold — aggressively compacting a session nowhere near its budget. Installing the engine and having it fail was worse than not installing it at all (the #29370 shape). A no-op-default, fail-open selection hook removes that class of bug at the source — which is why I'd argue the current P3 undersells it.

Related Issue

Implements RFC #36765 (mine).

Consolidates the per-turn request-selection / observation surface that four independent PRs each propose under different names, into one canonical pair:

PR	Author	Proposed surface	Covered here
#41918	@johnnykor82	prepare_request_messages() + on_turn_complete()	✅ both — selection → select_context(), observation → on_turn_complete()
#24949	@100yenadmin	preassemble() (per-turn message rewrite)	✅ selection
#47109	@huangxun375-stack	request-assembly + turn-observation lifecycle hooks	✅ both verbs
#50053	@arminanton	pre_send() + enforce_response()	⚠️ partial — pre_send() (pre-request selection) → select_context(); enforce_response() is NOT subsumed

Credit to all four — this unifies their shared idea. The on_turn_complete() design follows #41918; @johnnykor82 is added via Co-authored-by: (happy to correct the trailer to your preferred address). I'm not attached to the exact names; if maintainers prefer after_turn() / observe_turn() for the observation verb, that's fine.

Scope (corrected, per @arminanton's note): this PR consolidates the per-turn request-selection + post-turn observation surface only. #50053's enforce_response() is a different lifecycle point — active post-generation enforcement that can gate/regenerate/refuse the model's reply, which an observation hook cannot do. It is orthogonal and explicitly out of scope here; the clean split is select_context() (pre-request selection) + on_turn_complete() (post-turn observation), with enforce_response() riding on top as an optional third verb in a follow-up (@arminanton to re-scope #50053 to just that). My earlier "consolidates #50053" was over-broad on the enforcement half — corrected.

Type of Change

• ✨ New feature (non-breaking change that adds functionality)
• ✅ Tests (adding or improving test coverage)

Changes Made

• agent/context_engine.py: add optional select_context(...) and on_turn_complete(...) to the ABC. No-op defaults.
• agent/conversation_loop.py: add _apply_context_engine_selection() (called after request api_messages are assembled, before Anthropic cache-control) and _notify_context_engine_turn_complete() helpers. Both fail-open; the observation helper short-circuits the base no-op so non-implementing engines pay nothing.
• agent/turn_finalizer.py: call the observation helper from finalize_turn() after the existing post_llm_call plugin hook, with the finalized transcript + turn metadata, forwarding the turn's canonical usage when available.
• agent/conversation_loop.py (usage wiring): stash the most recent provider response's usage dict (the same canonical buckets fed to update_from_response) on the agent as _last_turn_usage, reset to None at turn start so turns that never reach a provider response forward None rather than stale usage.
• tests/agent/test_context_engine_select_context.py: cover both hooks' no-op defaults and every host-call-site branch (None/replace/missing-hook/no-engine/exception/non-list, base-no-op skip, metadata forwarding, shallow-copy semantics, persisted history not mutated).
• tests/agent/test_context_engine_on_turn_complete_usage.py: integration test through the real finalize_turn path — a completed turn forwards the full canonical bucket set intact; a no-provider-response turn forwards None.

How to Test

pytest tests/agent/test_context_engine_select_context.py tests/agent/test_context_engine_on_turn_complete_usage.py -q

The logic of both hooks, all host-call-site branches, and the usage-forwarding wiring was verified standalone on Python 3.13. I still couldn't run the full pytest tests/ -q suite locally — my macOS Homebrew Python 3.13 has a pyexpat/system-libexpat ABI mismatch that breaks pip itself (reinstalling python@3.13 + expat didn't relink it), unrelated to this change. Flagging so a maintainer can confirm against CI.

Checklist

Code

• I've read the Contributing Guide
• My commit messages follow Conventional Commits
• I searched for existing PRs (this consolidates them — see table)
• My PR contains only changes related to this feature
• I've run pytest tests/ -q and all tests pass — partial: focused logic verified standalone on 3.13; full suite pending CI
• I've added tests for my changes
• I've tested on my platform: macOS

Documentation & Housekeeping

• Docstrings on both new hooks — or N/A
• cli-config.yaml.example — N/A
• CONTRIBUTING.md / AGENTS.md — N/A
• Cross-platform impact — N/A (pure Python, stdlib only)
• Tool descriptions/schemas — N/A

[johnnykor82]

Thanks for putting this together. I’m supportive of consolidating the request-selection/request-assembly surface into one canonical hook, and select_context() seems like a reasonable name for that half of the problem.

One thing I want to clarify from #41918: that PR has two separate pieces, and this PR only subsumes one of them.

select_context() covers the request-only replacement part of #41918 (prepare_request_messages()): before provider dispatch, let the engine decide which messages should enter this request, without mutating persisted history. That consolidation makes sense to me.

But #41918 also adds on_turn_complete(), which is a different lifecycle verb and is not covered here. That hook is important because external context engines need a reliable post-turn observation point: after the assistant/tool loop has finished, the engine receives the finalized transcript snapshot and can ingest/index/summarize it for future selection. Without that, engines still need some workaround to learn what happened in the completed turn.

Why select_context() alone is not enough:

• select_context() is pre-request selection: “what context should go into this provider call?”
• on_turn_complete() is post-turn ingestion/observation: “what actually happened in the completed turn?”
• The latter lets engines update embeddings, summaries, topic/session state, routing metadata, or external memory from the finalized transcript.
• It avoids abusing should_compress() == True / compress() as a generic per-turn callback just to observe history.
• It also covers the case where a turn completes and there may not be a next request where the engine can infer the previous turn.

So I’d prefer that this PR either:

1. include a minimal on_turn_complete() / after_turn() / equivalent post-turn observation hook as the companion to select_context(), or
2. explicitly leave feat(context-engine): add request preparation hooks #41918 open as a follow-up for the observation/ingestion half, rather than treating feat(context-engine): add request preparation hooks #41918 as fully subsumed.

I’m not attached to the exact names from #41918. If maintainers prefer select_context() for request selection and after_turn() / observe_turn() / on_turn_complete() for observation, that’s fine with me. The key thing is preserving both verbs: selection before the request, and observation after the turn.

Also, since this PR is explicitly consolidating the request-selection surface from #41918, I’d like to be included with a Co-authored-by: trailer if this lands as the replacement path. Happy to provide the exact trailer details.

[chaos-xxl]

Good call — you're right that I was over-claiming. select_context() only covered the request-selection half of #41918; the post-turn observation verb is genuinely distinct and shouldn't be silently dropped. I've added it.

Pushed in this PR:

• on_turn_complete(messages, usage=None, **kwargs) — optional, no-op default, called once from turn_finalizer.finalize_turn() after the assistant/tool loop finishes, with a shallow-copied finalized transcript plus turn_id / task_id / api_call_count / interrupted / failed / turn_exit_reason. Wired via _notify_context_engine_turn_complete(); fail-open, and the base no-op is short-circuited so non-implementing engines (incl. the built-in compressor) pay nothing per turn.
• Tests cover both verbs: no-op defaults, snapshot + metadata forwarding, shallow-copy semantics, base-no-op skip, and fail-open.

So the PR now subsumes both halves of #41918 — prepare_request_messages → select_context(), and on_turn_complete() — rather than treating #41918 as fully covered while only doing one. I updated the PR body to say so honestly.

On attribution: I've added you via Co-authored-by: on the on_turn_complete commit (the observation design follows your #41918). I used johnnykor82@users.noreply.github.com as a placeholder — send me the exact trailer address you want and I'll amend it.

Agreed on naming flexibility too: if maintainers prefer after_turn() / observe_turn() for the observation verb (or a different signature), I'm happy to rename. The thing I care about is keeping both verbs first-class: selection before the request, observation after the turn.

[arminanton]

Thanks for consolidating the selection surface here. I want to flag a complementary PR so the maintainers can see how they relate, since the title lists #50053 under "consolidates" and I think the overlap is narrower than that implies.

#50053 adds a grounding/enforcement surface to the base ContextEngine, three verbs: capabilities() (feature negotiation), pre_send() (last-mile evidence injection before dispatch), and enforce_response() (audit the model's reply after it returns: accept, regenerate with a correction, or replace/refuse).

Mapping the two PRs against each other:

• pre_send() and your select_context() genuinely overlap. Both are pre-request message-list transforms. select_context() is the more general verb (selection/routing of the whole list), so I'm happy to concede that slot to it and drop pre_send() in favor of select_context().
• on_turn_complete() and enforce_response() are different lifecycle points, not the same verb. on_turn_complete() is post-turn observation: it records what happened so the next selection can use it, but it cannot change the answer that ships. enforce_response() is an active gate before the reply is shipped: it can force a citation-corrected regeneration or refuse a hallucinated answer outright. An observation hook can't do that.

So I don't think #50053 is subsumed here. The clean split looks like: select_context() owns pre-request selection, on_turn_complete() owns post-turn observation, and enforce_response() rides on top as the optional enforcement layer for retrieval-first / grounded engines that need to gate the reply. If this lands, I'll re-scope #50053 to just the enforce_response() verb so there's no duplication on the pre-send side.

Happy to align on naming/signatures with whatever you and the maintainers settle on for the selection and observation verbs. Does keeping enforcement as a separate third verb on top of your two sound right to you?

[johnnykor82]

Thanks for adding on_turn_complete() here — that covers the main concern I raised above, and I appreciate the quick follow-up + attribution.

One smaller but important contract detail from #41918: the hook should receive the actual provider usage metadata when it is available. I noticed the current integration supports a usage parameter in the hook/helper signature, but the production call from turn_finalizer.finalize_turn() appears to pass usage=None.

I think it would be worth preserving the original behavior here and forwarding the real usage dict for the completed turn/API call where possible.

Why this matters:

• on_turn_complete() is not only “here are the messages”; it is the engine’s post-turn observation point.
• For an external context engine, the usage metadata tells it how expensive / large the selected context actually was: prompt tokens, completion tokens, total tokens, cache read/write tokens, reasoning tokens, etc.
• That lets the engine make better decisions on the next select_context() call: whether the previous selection was too large, whether cache behavior was healthy, whether a turn should be summarized/indexed aggressively, or whether the engine should tighten routing/selection.
• It also avoids forcing engines to re-tokenize or estimate cost themselves, which can be provider/model-specific and less accurate than the host’s canonical usage accounting.
• For long sessions, this becomes useful telemetry: the engine can correlate “what happened in this turn” with “what it cost”, instead of only seeing the transcript text.

usage=None is fine for paths where the host genuinely does not have usage data, such as early failures or interrupted turns before a provider response. But on normal completed turns, if the loop already has canonical usage from update_from_response() / response handling, I’d strongly prefer passing that through to on_turn_complete().

So the request is: could this PR preserve the #41918 contract by forwarding actual usage metadata when available, and only falling back to None when it is genuinely unavailable?

[alt-glitch]

This was generated by AI during triage.

Related (competing/consolidating cluster): this PR consolidates RFC #36765 against #41918 (on_turn_complete + prepare_request_messages), #47109 (lifecycle hooks), and #50053. The author has acknowledged select_context() overlaps the request-selection half of #41918 while on_turn_complete() is distinct. Flagging the cluster so a maintainer can pick the canonical hook surface.

[chaos-xxl]

@johnnykor82 done — you're right that usage=None defeated the point of the observation hook. Forwarded the real usage, preserving the #41918 contract.

Pushed in 2e8b9e4:

• The loop now stashes the most recent provider response's usage dict on the agent as _last_turn_usage — the same canonical shape already fed to update_from_response(): prompt_tokens / completion_tokens / total_tokens plus the canonical input_tokens / output_tokens / cache_read_tokens / cache_write_tokens / reasoning_tokens buckets. No re-tokenizing, no estimation — the engine gets the host's canonical accounting.
• finalize_turn() forwards _last_turn_usage instead of None.
• It's reset to None at turn start, so the None fallback is preserved exactly for the cases you called out — early failure / interrupt before any provider response. Those forward None, not a stale prior turn's usage.
• Tests cover both ends through the real finalize_turn path: a completed turn forwards the full canonical bucket set intact; a no-response turn forwards None.

One semantics call I made, flag if you'd prefer otherwise: a turn can make several API calls (tool loop), and I forward the latest call's usage rather than a per-turn sum. Reasoning: the signal you described — "how large/expensive was the selected context" — is best read from the final assembled request's prompt_tokens; summing prompt tokens across tool-loop calls would double-count the context. Happy to switch to an accumulated per-turn total (or forward both) if your engine wants aggregate cost instead.

I also documented the usage param contract on the ABC hook so it's not just implied by the signature.

On attribution: you're co-authored on this commit too. I'm still using the johnnykor82@users.noreply.github.com placeholder — send me the exact trailer address you want and I'll amend across the commits.

[chaos-xxl]

@arminanton agreed on all counts — and you're right that "consolidates #50053" was over-broad. I've corrected it.

Your split is exactly the right mental model:

• select_context() — pre-request selection
• on_turn_complete() — post-turn observation
• enforce_response() — post-generation enforcement, riding on top

enforce_response() is genuinely a different lifecycle point, not a redrawing of either of my two verbs. on_turn_complete() is a passive observer — it records what happened so the next selection is better-informed, but by construction it cannot change the reply that already shipped. An active gate that can force a citation-corrected regeneration or refuse a hallucinated answer is a distinct capability, and an observation hook should not pretend to cover it. So yes — keeping enforcement as a separate third verb on top of these two is right, and I don't think it should be folded into this PR.

What I changed in the PR:

• The consolidation table now marks feat(context-engine): additive grounding hook points on the base ContextEngine #50053 as partial — pre_send() → select_context() covered, enforce_response() explicitly not subsumed.
• Added a scope note stating this PR consolidates only the per-turn selection + observation surface, and that enforce_response() is orthogonal and out of scope here.

Thanks for conceding the pre_send() slot to select_context() — that leaves a clean, non-overlapping three-verb picture for the maintainers: selection + observation land here as the base, and you re-scope #50053 to just enforce_response() as the enforcement layer on top. I'll align with whatever names/signatures you and the maintainers settle on for the selection/observation verbs so the third verb composes cleanly.

[arminanton]

@chaos-xxl @johnnykor82 [partially-related-external-topic] When you have some time, could you review the implementation in #51565? I have been using and testing it a lot for the past 3 weeks, but would love to hear your perspectives/suggestions.

[chaos-xxl]

Status update for whoever picks this up:

• Rebased onto latest main — now even (0 behind), no conflicts. The two touched hot files (conversation_loop.py, turn_finalizer.py) still apply cleanly on current main.
• Mergeable, additive only (no deletions), no-op by default + fail-open, tests included. CI hasn't run yet — happy for it to be kicked whenever a maintainer can approve the workflow.
• Cluster has converged on the two verbs: select_context() (pre-request selection) + on_turn_complete() (post-turn observation). @johnnykor82's feat(context-engine): add request preparation hooks #41918 feedback is folded in (real usage now forwarded to on_turn_complete; co-authored), and @arminanton agreed to scope feat(context-engine): additive grounding hook points on the base ContextEngine #50053 down — pre_send() → select_context(), with enforce_response() left as a separate third verb on top rather than subsumed here.

Nothing needed from reviewers to unblock beyond a look / CI run. Happy to adjust naming (after_turn()/observe_turn()) or anything else if it helps it land.

[chaos-xxl]

Re: the auto-applied sweeper:risk-session-state label — flagging for reviewers that this change is designed specifically not to touch persisted session/context state, and the tests pin it:

• select_context() is request-only. The host helper returns a replacement for the per-request api_messages list only — the persisted conversation history / session store is never written. Covered by test_persisted_history_not_mutated.
• on_turn_complete() gets a shallow copy. The transcript is passed as [dict(m) for m in messages] and the return value is ignored, so an engine can't mutate the persisted transcript. Covered by the shallow-copy assertion (captured["messages"] is not HISTORY).
• No-op default + fail-open. Both hooks default to return None, and every host call site swallows exceptions / invalid returns — so for the built-in compressor and any non-implementing engine there's zero behavioral change.
• The only new per-turn state is agent._last_turn_usage (a usage dict), reset to None at the start of each turn.

So the risk surface only exists for an engine that explicitly opts in, and even then the host contract prevents persisted-state mutation. Happy to add an assertion or docstring note if a reviewer wants it more explicit.

[100yenadmin]

@Teknium can you check this out?

[chaos-xxl]

Re the #24949 sweeper review's three concerns for a per-request rewrite hook — #51226 already satisfies them:

• Ordering: select_context() runs before cache-control and every request sanitizer (conversation_loop.py:1008, ahead of cache-control :1022, _sanitize_api_messages, _drop_thinking_only_and_merge_users). So a replacement can't bypass validation, and a role-unusual list is normalized downstream rather than reaching the provider.
• Cache invariant (AGENTS.md): the no-op default leaves the request byte-identical → zero cache impact for the built-in compressor and any non-implementing engine. An opt-in replacement changes only its own prefix; cache-control breakpoints re-derive on the selected list.
• Persistence isolation + fail-open: already covered by tests.

Just pushed the explicit ordering/cache contract to the docstring plus two tests (no-op byte-stability for cache-control; role-unusual replacement → downstream sanitizers). Happy to align the naming/shape as the canonical surface if useful.

[johnnykor82]

@chaos-xxl A fresh review on #41918 surfaced one fail-open edge case that
also appears to be present in the current #51226 diff, so I checked the
consolidated implementation directly.

The current validation is:

if isinstance(selected, list) and all(isinstance(m, dict) for m in selected):
    return selected

This correctly rejects non-list values and lists containing non-dicts, but an
empty list is still accepted because all([]) evaluates to True.

That matters for the fail-open contract inherited from #41918. If an engine
accidentally returns [] due to a backend failure or selection bug, the host
will replace a valid assembled request with an empty message list. The
downstream sanitizers cannot restore the original request, so this may reach
the provider as an invalid/empty request instead of falling back to the
unmodified api_messages.

The minimal fix would be to require a non-empty list:

if (
    isinstance(selected, list)
    and selected
    and all(isinstance(m, dict) for m in selected)
):
    return selected

It would also be useful to add an empty list keeps the original request
regression test alongside the existing non-list and list-of-non-dicts tests.

One smaller documentation point from the same review: #51226 documents the
new hooks in code, but the public ContextEngine guide at
website/docs/developer-guide/context-engine-plugin.md still describes the
old optional-method/lifecycle contract. Since select_context() and
on_turn_complete() are public plugin APIs, could that page be updated as
part of the consolidated PR as well?

The on_turn_complete() path and real usage forwarding otherwise look
preserved correctly in the current implementation.

[chaos-xxl]

Good catch @johnnykor82 — all([]) being True did let an empty list through, which breaks the fail-open contract exactly as you describe. Fixed:

• _apply_context_engine_selection now requires a non-empty list of dicts (... and selected and ...); [] falls open to the unmodified request.
• Added the empty list keeps the original request regression test alongside the non-list / list-of-non-dicts ones.
• Updated the public guide (website/docs/developer-guide/context-engine-plugin.md) to document select_context() / on_turn_complete() — it was still describing only the old should_compress/compress contract.

Pushed in 98c053a. Thanks for reviewing the consolidated diff directly.

[teknium1]

Thanks for consolidating the request-selection and observation surfaces and for following up on the empty-list and public-documentation feedback.

Problems

• agent/conversation_loop.py:548-552 passes live messages and messages[current_turn_user_idx] to the new hook. A context engine can mutate either object in place, altering persisted transcript state despite the request-only contract. The proposed persistence test at tests/agent/test_context_engine_select_context.py:199-215 does not perform such a mutation.
• agent/turn_finalizer.py:420-431 is not reached by direct terminal returns such as agent/conversation_loop.py:1137-1150 and :1567-1576. Those paths persist and return without on_turn_complete(), so early failure/interruption does not receive the documented usage=None observation.

Suggested changes

• Snapshot the hook inputs and add mutation-regression coverage.
• Route all terminal paths through a shared notification/finalization seam, with an integration test for a real early-return path.

Automated hermes-sweeper review.

[teknium1]

This notification only runs for paths that reach finalize_turn. conversation_loop.py still has direct terminal returns after persistence (for example its no-fallback rate-limit path), so those completed failures never invoke on_turn_complete() or receive the documented usage=None behavior. Route terminal exits through a shared notification seam or notify before each such return.

[chaos-xxl]

Consolidated status on the two review points (commits 55de4a3, a0cc33f):

1. Live hook inputs → done. select_context() now receives shallow copies of conversation_messages and incoming_message, so an engine mutating them in place can't touch persisted state. Regression test added that intentionally mutates both inputs and asserts the persisted history + incoming message are unchanged.

2. Terminal-path coverage → scoped, with the gap documented and the seam left as a follow-up.

• Docs are now consistent: both the ABC docstring and the public plugin guide state that on_turn_complete() is best-effort — it fires from the standard finalization seam, and abnormal early-return paths (content-policy block, provider terminal failure, etc.) bypass finalization and do not emit it.
• Added a test pinning the testable half of that contract: the finalization seam emits on_turn_complete() with usage=None and the interrupted flag on an interrupted turn.
• I deliberately did not route all terminal paths through a shared seam in this PR: there are ~26 direct early returns in run_conversation, so unifying them is a control-flow refactor of a core path rather than part of this additive hook. Happy to do that as a dedicated follow-up — or to wire the specific named paths (content-policy / no-fallback rate-limit) if you'd prefer partial coverage in-tree now.

Both hooks remain no-op-default and fail-open, so existing engines and the built-in compressor are unaffected.

[teknium1]

Landed on main via salvage PR #70458 (rebase-merge b55bb2c) — all seven of your commits with authorship preserved, @chaos-xxl, plus @johnnykor82's Co-authored-by. Outstanding work: the RFC, the dogfooded consumer that surfaced the gap, and the fastest defect-fix turnaround we've seen in a review thread.

Two maintainer commits rode on top: (1) a base-method identity short-circuit at the select_context call site — the hasattr() check meant the built-in compressor paid per-request history copies + a hook call every turn (the ABC default always exists); now the default path is byte-identical, pinned by tests; (2) scoping docs in the plugin guide — select_context() is for engines that must REPLACE per-request context; observation-only plugins should implement a MemoryProvider (sync_turn) instead; non-no-op selection affects the prompt-cache prefix, so engines should return stable selections when nothing changed.

The should_compress()=True backdoor era is over. Thanks for driving the whole cluster to consensus.