feat: hook system and plugin support for Mellea

[araujof]

Summary

This PR introduces a plugin and hook system for Mellea, built on top of the cpex framework. It adds an extensibility layer that allows users to register handlers at defined points in the session, generation, sampling, component, tool, and validation pipelines.

Changes

New: mellea/plugins/ package

• types.py — Defines HookType enum (30 hook points across 7 categories) and PluginMode enum (SEQUENTIAL, CONCURRENT, AUDIT, FIRE_AND_FORGET). Registers all hook types with the ContextForge hook registry.
• base.py — MelleaBasePayload (frozen Pydantic model shared by all payloads), MelleaPlugin base class with typed accessors for backend/context/session and context manager support, and PluginViolationError for blocked-execution errors.
• decorators.py — @hook(hook_type, mode, priority) for standalone functions or class methods; @plugin(name, priority) for plain classes; both inject context manager support (with/async with).
• policies.py — HookPayloadPolicy per hook type declaring which fields plugins may write; DefaultHookPolicy enum; MELLEA_HOOK_PAYLOAD_POLICIES registry.
• registry.py — register(), deregister_session_plugins(), shutdown_plugins(), has_plugins() (accepts optional HookType filter), and init_plugins().
• manager.py — Singleton PluginManager wrapper; invoke_hook() dispatches payloads, applies HookPayloadPolicy to filter plugin-proposed modifications to writable fields only, raises PluginViolationError when a plugin blocks in enforce mode, and routes FIRE_AND_FORGET hooks to background tasks.
• pluginset.py — PluginSet for grouping multiple plugins or hook functions with a shared priority and context manager support.
• context.py — build_global_context() assembles the PluginContext passed to every hook invocation.
• hooks/ — Typed payload classes for each hook category:
  • session.py: SessionPreInitPayload, SessionPostInitPayload, SessionResetPayload, SessionCleanupPayload
  • generation.py: GenerationPreCallPayload, GenerationPostCallPayload, GenerationStreamChunkPayload
  • sampling.py: SamplingLoopStartPayload, SamplingIterationPayload, SamplingRepairPayload, SamplingLoopEndPayload
  • component.py: ComponentPreCreatePayload, ComponentPostCreatePayload, ComponentPreExecutePayload, ComponentPostSuccessPayload, ComponentPostErrorPayload
  • tool.py: ToolPreInvokePayload, ToolPostInvokePayload
  • validation.py: ValidationPreCheckPayload, ValidationPostCheckPayload

Modified: stdlib integration

• mellea/stdlib/functional.py — invoke_hook() calls added at: generation pre/post call, sampling loop start/iteration/repair/end, component pre-create/post-create/pre-execute/post-success/post-error.
• mellea/stdlib/sampling/base.py — Hook call sites added to the sampling loop; passes sample_contexts[best_failed_index] (not original context) on failure path.
• mellea/stdlib/session.py — Session lifecycle hooks (SESSION_PRE_INIT, SESSION_POST_INIT, SESSION_RESET, SESSION_CLEANUP) added; session state threaded into global plugin context.
• mellea/core/backend.py — Backend reference surfaced in plugin context.

Documentation

• docs/dev/hook_system.md — Full specification: hook types, payload schemas, write-protection rules, execution modes, plugin scoping (global vs session), policy enforcement, known limitations.
• docs/dev/hook_system_implementation_plan.md — Phased implementation plan.

Examples

• docs/examples/plugins/ — Six example scripts covering: standalone hook functions, class-based plugins (@plugin), session-scoped plugins, plugin set composition, concurrent hooks, and tool hooks.

Tests

test/plugins/:

• All payload types and their field shapes (test_all_payloads.py)
• Blocking behavior and PluginViolationError (test_blocking.py)
• Global context construction (test_build_global_context.py)
• Decorator behavior (test_decorators.py)
• Example script execution (test_example_scripts.py)
• Execution modes: sequential, concurrent, audit, fire-and-forget (test_execution_modes.py)
• Hook call sites in stdlib (test_hook_call_sites.py)
• Plugin manager lifecycle (test_manager.py)
• MelleaPlugin base class (test_mellea_plugin.py)
• Payload memory reference behavior (test_payload_memory_references.py)
• Payload field access (test_payloads.py)
• PluginSet (test_pluginset.py)
• Policy definitions (test_policies.py)
• Policy enforcement (write-protection, field filtering) (test_policy_enforcement.py)
• Priority ordering (test_priority_ordering.py)
• Plugin scoping (global vs session isolation) (test_scoping.py)

Dependencies

• cpex added as an optional dependency under the cpex extra.

Notes

• The plugin framework is optional. All plugin paths guard with try/except ImportError; stdlib call sites are no-ops when cpex is not installed.
• PluginViolationError is raised only when a plugin explicitly sets block=True in its result and the hook's policy is DefaultHookPolicy.SEQUENTIAL.
• Plugin identity in the registry uses the function __name__ attribute; dynamically generated hooks with identical names will collide with "Plugin already registered" errors.

[mergify]

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Enforce conventional commit

Wonderful, this rule succeeded.

Make sure that we follow https://www.conventionalcommits.org/en/v1.0.0/

• title ~= ^(fix|feat|docs|style|refactor|perf|test|build|ci|chore|revert|release)(?:\(.+\))?:

[jakelorocco]

It's looking good! I will take a deeper look at the new tests and make sure our existing tests haven't broken.

[jakelorocco]

We should change this to generate_from_context and have the other backend classes implement a generate function that gets called here.

@nrfulton @HendrikStrobelt, we should also make a note that generate_from_raw won't have these hooks. Or we should decide now if we want them to and quickly create separate hooks for generate from raw. It will mostly be the same interface except actions can be an array.

[araujof]

done.

[jakelorocco]

We briefly mentioned this before, but sampling hooks will only impact sampling strategies derived from BaseSampling strategy that don't override the sample method. I think this okay for now, but we should eventually refactor our sampling classes to handle this better. @nrfulton and @HendrikStrobelt, do you agree?

[jakelorocco]

Will talk with the team; one option is to disable plugins if python <= 3.10. Putting our version download numbers here for future reference:

  ┌────────────────┬───────────┬───────┐
  │ Python Version │ Downloads │ Share │
  ├────────────────┼───────────┼───────┤
  │ 3.12           │ 76,864    │ 59.0% │
  ├────────────────┼───────────┼───────┤
  │ 3.13           │ 27,919    │ 21.4% │
  ├────────────────┼───────────┼───────┤
  │ 3.11           │ 11,428    │ 8.8%  │
  ├────────────────┼───────────┼───────┤
  │ 3.14           │ 7,416     │ 5.7%  │
  ├────────────────┼───────────┼───────┤
  │ 3.10           │ 4,106     │ 3.2%  │
  ├────────────────┼───────────┼───────┤
  │ 3.9            │ 50        │ ~0%   │
  └────────────────┴───────────┴───────┘

[jakelorocco]

Decided that we will move to stop explicitly supporting python 3.10. I think there are changes to do so (ie change testing setup and some linting rules) that are outside the scope of this PR. I think we can continue with this PR even if "breaks" python 3.10 and make the version changes in a separate PR.

[jakelorocco]

Adding a comment here so that we remember to update this before the final merge.

[jakelorocco]

I believe you can remove this. We run our examples as part of our testing pipeline already.

[araujof]

Comment in the summary. I think it's good to be able to run it as part of dev pytest. Let me know if I'm missing any other dev tooling otherwise.

[araujof]

Thanks for the thorough review, @jakelorocco!

I've pushed changes addressing most of your comments. Here's the breakdown:

Changes

Timing removed from generate_from_context_with_hooks (comments on backend.py)
Removed time.monotonic() and set latency_ms=0. The MOT is uncomputed (lazy) at this point so timing was meaningless. Added a docstring + TODO about moving the post-call hook into ModelOutputThunk.astream (after post_process) where latency can be measured accurately.

Sync instruct()/chat() now fire component creation hooks (base.py, functional.py)
COMPONENT_PRE_CREATE and COMPONENT_POST_CREATE were only firing in the async ainstruct()/achat() path. The sync instruct() and chat() now delegate to _run_async_in_thread(ainstruct(...)) and _run_async_in_thread(achat(...)) respectively, so hooks fire in both sync and async paths.

@hook decorator validates async (decorators.py)
Added inspect.iscoroutinefunction check — raises TypeError with a clear message if the decorated function is not async def.

Payload field docstrings with expected types (all hooks/*.py)
Added Attributes docstrings to all 14 payload classes with expected types (e.g. Component, Context, ModelOutputThunk) and writable/observe-only annotations so the types are visible without digging into the code.

class_plugin.py example fixes

• Fixed redact_output to handle uncomputed MOT (model_output.value may be None when the hook fires). Added docstring explaining the limitation.
• Added shared-state warning about redaction_count and thread-safety if hooks ever run in parallel threads.

pyproject.toml fixes

• Restored click<8.2.0 — removal was accidental (collateral from a different change).
• Added grpcio>=1.78.0 to the cpex optional dependency group to fix the grpcio version mismatch error.

Renamed all_valid field
Renamed to all_validations_passed. Happy to rename if the team agrees on a different name.

Misc
Fixed a few mypy false-positives.

Items Deferred (need team discussion)

Comment	Status
Rename generate_from_context_with_hooks / restructure Backend inheritance	Deferred — major refactor across all Backend subclasses, needs alignment with @nrfulton and @HendrikStrobelt
generate_from_raw hook support	Design decision — should we create separate hooks or document that generate_from_raw is hook-free?
Tool hooks location (_call_tool vs MelleaTool.call)	Keeping in _call_tool for now since it has access to backend/context; open to moving if we adopt context managers for threading backend/session (ref: #400)
Frozen payloads and mutable references	Acknowledged — payloads hold live object refs. However, cpex wraps then in a copy-on-write frozen model that transionally wraps dict and list objects in copy-on-write couterparts. We do this to avoid deep copies on every payload. Check here to see how we handle this in the cpex framework.
Sampling hooks only fire for BaseSamplingStrategy subclasses	Acknowledged limitation, deferred refactor
Duplicate context in invoke_hook kwargs vs payload	The global context serves as a transport for objects (backend, session) that aren't part of the payload schema. The payload's context field and the kwarg will always be the same object at the call site
Remove test_example_scripts.py	Kept for now — docs/examples/conftest.py handles live collection but skips when Ollama isn't available. The mock-based tests in test_example_scripts.py ensure examples run in CI without any backend
cpex git dependency	Will update to a released version before final merge
Python <=3.10 plugin support	Already gated with python_version >= '3.11' in the optional dep. We should be good here supporting >=3.11, and we will reach EoL of Python 3.10 soon anyway.

[ajbozarth]

Would it make sense to group this into the telemetry group? This definitely feels like it makes sense in that group (especially since I plan to update the telemetry to use the hooks)

[araujof]

Hi @ajbozarth ! Good timing :)
We may want to keep it separate, at least for now. The way I see it, telemetry will be built as a plugin extension, but there are many other capabilities you could build with the new hook system that aren't observability-related. Maybe, in the future we will end-up with a superset group, but for now, I'd just keep them separate. I added cpex to the all group in my latest commit.

[ajbozarth]

in that case I'll just make telemetry group dependent on this one once I implement that then.

I do wonder about the name though, cpex is very specific and most users might not know what it is, maybe we should call it hooks?

[jakelorocco]

that's a good call; yes, we should change the group name to hooks / plugins or something of that nature

[ajbozarth]

also worth noting you'll need to rebase/merge once the merge queue is finished. I just reworked the pyproject file to clean up dependencies in #588

[araujof]

good call, let's rename the group. Are we good with hooks/plugins?

[ajbozarth]

I think hooks for this because it enables hooks, whereas we might have a plugins group in the future which contains dependencies for actual plugin implementations

[jakelorocco]

Thanks for the thorough review, @jakelorocco!

I've pushed changes addressing most of your comments. Here's the breakdown:

Changes

Timing removed from generate_from_context_with_hooks (comments on backend.py) Removed time.monotonic() and set latency_ms=0. The MOT is uncomputed (lazy) at this point so timing was meaningless. Added a docstring + TODO about moving the post-call hook into ModelOutputThunk.astream (after post_process) where latency can be measured accurately.

IMO, this must be moved before merging if we want to support this hook.

Sync instruct()/chat() now fire component creation hooks (base.py, functional.py) COMPONENT_PRE_CREATE and COMPONENT_POST_CREATE were only firing in the async ainstruct()/achat() path. The sync instruct() and chat() now delegate to _run_async_in_thread(ainstruct(...)) and _run_async_in_thread(achat(...)) respectively, so hooks fire in both sync and async paths.

I would prefer we just move this hook to Component init. We can change Component from being a protocol, but I don't think this hook makes sense if it doesn't run for every component. I also don't like re-routing instruct / chat.

class_plugin.py example fixes

• Fixed redact_output to handle uncomputed MOT (model_output.value may be None when the hook fires). Added docstring explaining the limitation.
• Added shared-state warning about redaction_count and thread-safety if hooks ever run in parallel threads.

I think these fixes don't resolve the redaction issue alone.

Items Deferred (need team discussion)

Comment	Status
Rename generate_from_context_with_hooks / restructure Backend inheritance	Deferred — major refactor across all Backend subclasses, needs alignment with @nrfulton and @HendrikStrobelt

I don't think this needs to be a large refactor. We can simply add an intermediary function. I think this is also a must have / must fix before merging.

Frozen payloads and mutable references	Acknowledged — payloads hold live object refs. However, cpex wraps then in a copy-on-write frozen model that transionally wraps dict and list objects in copy-on-write couterparts. We do this to avoid deep copies on every payload. Check here to see how we handle this in the cpex framework.

Is this saying that we do actually prevent modification of subfields? I might have misread the code previously, but I think it only prevents it for known object types.

Duplicate context in invoke_hook kwargs vs payload	The global context serves as a transport for objects (backend, session) that aren't part of the payload schema. The payload's context field and the kwarg will always be the same object at the call site

If this is the case, why do we duplicate this information?

Remove test_example_scripts.py	Kept for now — docs/examples/conftest.py handles live collection but skips when Ollama isn't available. The mock-based tests in test_example_scripts.py ensure examples run in CI without any backend

I would prefer we don't keep this. Any solution to this should be done for the whole code base, not a one off solution. We will make sure examples work (mocked or not).

---

Also, please install and run our pre-commit hooks to make sure formatting is okay. To prevent having to fix each commit individually, you should be able to run pre-commit run --all-files.

[araujof]

Hi @jakelorocco, thanks for the follow-up feedback. Here's a summary of the latest round of changes (commits after 88cca3a):

Addressed from your feedback

1. Renamed generate_from_context_with_hooks → hooks now injected via __init_subclass__ (backend.py)
No more separate wrapper method. Backend.__init_subclass__ now wraps generate_from_context at class definition time for any concrete subclass. This means:

• All backends get hooks automatically — no need to call a different method.
• aact() now calls backend.generate_from_context(...) directly (the wrapper is transparent).
• No changes required in backend subclasses.

2. generation_post_call hook moved into MOT lifecycle (backend.py, base.py)
The post-call hook now fires via an _on_computed callback on the ModelOutputThunk:

• Lazy MOTs (normal path): fires inside ModelOutputThunk.astream after post_process, so latency_ms now reflects the real time from generate_from_context call to value materialization.
• Already-computed MOTs (cached): fires inline before returning. MOT replacement is supported in both paths.

3. Component hooks: removed COMPONENT_PRE_CREATE / COMPONENT_POST_CREATE (hooks/component.py, functional.py)
Rather than re-routing instruct()/chat() through async, we dropped the pre/post create hooks entirely for now. The COMPONENT_PRE_EXECUTE and COMPONENT_POST_EXECUTE hooks (which fire in aact()) remain and cover the component lifecycle at the execution boundary. Moving creation hooks into Component.__init__ is a better design but requires making Component a concrete class. This is now documented in the spec. We'll tackle that separately.

4. Sync instruct() / chat() no longer re-route through async
Reverted to direct act() delegation (constructing the Instruction/Message inline), which is simpler and avoids the _run_async_in_thread(ainstruct(...)) indirection you flagged.

5. Dependency group renamed from cpex to hooks (pyproject.toml)
Per @ajbozarth and your suggestion, mellea[hooks] is now the install target.

6. Removed test_example_scripts.py
Dropped per your request.

New additions

7. unregister() function (registry.py)
Per @HendrikStrobelt suggestion, new public API to remove globally-registered plugins. Accepts the same items as register(): @hook functions, Plugin instances, PluginSet instances, or lists thereof.

8. modify() helper (registry.py)
Convenience function for returning modifying PluginResults from hooks. It mirrors block() for the modification case:

return modify(payload, model_output=new_mot)

9. Instruction.description property (components/instruction.py)
Added getter/setter so plugins can read/modify the description text on created Instruction objects.

Housekeeping

• Lint, formatting, and mypy fixes across the board
• Updated specs and implementation plan to reflect the current state
• Updated lock file

Still open / deferred

Item	Status
generate_from_raw hook support	Needs design decision — separate hooks or documented as hook-free?
Tool hooks in _call_tool vs MelleaTool.call	Keeping in _call_tool for now; revisit with #400
Frozen payload mutable references	cpex uses copy-on-write wrappers for known collection types — deeper object graphs are not deep-copied. Happy to discuss if we need stronger guarantees.
Duplicate context in invoke_hook kwargs vs payload	I can see the overlap here. Maybe we should drop it from the global_context. I need to think about it a bit more.
cpex git dependency → released version	Will update before final merge. The package is already published in PyPi

[ajbozarth]

Nit: should move this up into the relevant section above (note section comments)