Skip to content

Commit 70ab0c1

Browse files
committed
feat: specify hook payload write protection
Signed-off-by: Frederico Araujo <frederico.araujo@ibm.com>
1 parent 2bac8d2 commit 70ab0c1

2 files changed

Lines changed: 285 additions & 27 deletions

File tree

docs/dev/hook_system.md

Lines changed: 149 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ class MyPlugin(MelleaPlugin):
3737
) -> PluginResult | None
3838
```
3939

40-
- **`payload`**: Mutable, strongly-typed data specific to the hook point
40+
- **`payload`**: Immutable (frozen), strongly-typed data specific to the hook point. Plugins use `model_copy(update={...})` to propose modifications
4141
- **`context`**: Read-only shared context with session metadata and utilities
4242
- **`mode`**: `"enforce"` (default), `"permissive"`, or `"fire_and_forget"` — controls execution behavior (see Execution Mode below)
4343
- **`priority`**: Lower numbers execute first (default: 50)
@@ -75,6 +75,7 @@ The hook system is backed by a lightweight plugin framework built as a Mellea de
7575
- Exposes `PluginSet` for grouping related hooks/plugins into composable, reusable units
7676
- Exposes `register()` for global plugin registration and `block()` as a convenience for returning blocking `PluginResult`s
7777
- Implements a plugin manager that loads, registers, and governs the execution of plugins
78+
- Enforces per-hook-type payload policies via `HookPayloadPolicy`, accepting only writable-field changes from plugins
7879

7980
The public API surface:
8081

@@ -108,27 +109,31 @@ result = await plugin_manager.invoke_hook(hook_type, payload, context)
108109
The caller (the base class method) is responsible for both invoking the hook and processing the result. Processing means checking the result for one of three possible outcomes:
109110

110111
1. **Continue with original payload**: — `PluginResult(continue_processing=True)` with no `modified_payload`. The caller proceeds unchanged.
111-
2. **Continue with modified payload**: — `PluginResult(continue_processing=True, modified_payload=...)`. The caller uses the modified payload fields in place of the originals.
112+
2. **Continue with modified payload**: — `PluginResult(continue_processing=True, modified_payload=...)`. The plugin manager applies the hook's payload policy, accepting only changes to writable fields and discarding unauthorized modifications. The caller uses the policy-filtered payload in place of the original.
112113
3. **Block execution**`PluginResult(continue_processing=False, violation=...)`. The caller raises or returns early with structured error information.
113114

114115
Hooks cannot redirect control flow, jump to arbitrary code, or alter the calling method's logic beyond these outcomes. This is enforced by the `PluginResult` type.
115116

116117
### Payload Design Principles
117118

118-
Hook payloads follow five design principles:
119+
Hook payloads follow six design principles:
119120

120121
1. **Strongly typed** — Each hook has a dedicated payload dataclass (not a generic dict). This enables IDE autocompletion, static analysis, and clear documentation of what each hook receives.
121122
2. **Sufficient (maximize-at-boundary)** — Each payload includes everything available at that point in time. Post-hooks include the pre-hook fields plus results. This avoids forcing plugins to maintain their own state across pre/post pairs.
122-
3. **Immutable context**`PluginContext` fields are read-only; only the `payload` is mutable. This separates "what the plugin can observe" from "what the plugin can change."
123-
4. **Serializable** — Payloads should be serializable for external (MCP-based) plugins that run out-of-process. All payload fields use types that can round-trip through JSON or similar formats.
124-
5. **Versioned** — Payload schemas carry a `payload_version` so plugins can detect incompatible changes at registration time rather than at runtime.
123+
3. **Frozen (immutable)** — Payloads are frozen Pydantic models (`model_config = ConfigDict(frozen=True)`). Plugins cannot mutate payload attributes in place. To propose changes, plugins must call `payload.model_copy(update={...})` and return the copy via `PluginResult.modified_payload`. This ensures every modification is explicit and flows through the policy system.
124+
4. **Policy-controlled** — Each hook type declares a `HookPayloadPolicy` specifying which fields are writable. The plugin manager applies the policy after each plugin returns, accepting only changes to writable fields and silently discarding unauthorized modifications. This separates "what the plugin can observe" from "what the plugin can change"and enforces it at the framework level. See [Hook Payload Policies](#hook-payload-policies) for the full policy table.
125+
5. **Serializable** — Payloads should be serializable for external (MCP-based) plugins that run out-of-process. All payload fields use types that can round-trip through JSON or similar formats.
126+
6. **Versioned** — Payload schemas carry a `payload_version` so plugins can detect incompatible changes at registration time rather than at runtime.
125127

126128
## 2. Common Payload Fields
127129

128130
All hook payloads inherit these base fields:
129131

130132
```python
131133
class BasePayload(PluginPayload):
134+
"""Frozen base — all payloads are immutable by design."""
135+
model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
136+
132137
session_id: str | None = None # Session identifier (None for functional API calls)
133138
request_id: str # Unique ID for this execution chain
134139
timestamp: datetime # When the event fired
@@ -168,6 +173,131 @@ class BasePayload(PluginPayload):
168173
| `context_prune` | Context Operations | Context | When context is trimmed |
169174
| `error_occurred` | Error Handling | Cross-cutting | When an unrecoverable error occurs |
170175

176+
## 3b. Hook Payload Policies
177+
178+
Each hook type declares a `HookPayloadPolicy` that specifies which payload fields plugins are allowed to modify. The plugin manager enforces these policies after each plugin returns: only changes to writable fields are accepted; all other modifications are silently discarded.
179+
180+
Hooks not listed in the policy table are **observe-only** — plugins can read the payload but cannot modify any fields.
181+
182+
### Policy Types
183+
184+
```python
185+
from dataclasses import dataclass
186+
from enum import Enum
187+
188+
class DefaultHookPolicy(str, Enum):
189+
"""Controls behavior for hooks without an explicit policy."""
190+
ALLOW = "allow" # Accept all modifications (backwards-compatible)
191+
DENY = "deny" # Reject all modifications (strict mode, default for Mellea)
192+
193+
@dataclass(frozen=True)
194+
class HookPayloadPolicy:
195+
"""Defines which payload fields plugins may modify."""
196+
writable_fields: frozenset[str]
197+
```
198+
199+
### Policy Enforcement
200+
201+
When a plugin returns `PluginResult(modified_payload=...)`, the plugin manager applies `apply_policy()`:
202+
203+
```python
204+
def apply_policy(
205+
original: BaseModel,
206+
modified: BaseModel,
207+
policy: HookPayloadPolicy,
208+
) -> BaseModel | None:
209+
"""Accept only changes to writable fields; discard all others.
210+
211+
Returns an updated payload via model_copy(update=...), or None
212+
if the plugin made no effective (allowed) changes.
213+
"""
214+
updates: dict[str, Any] = {}
215+
for field in policy.writable_fields:
216+
old_val = getattr(original, field, _SENTINEL)
217+
new_val = getattr(modified, field, _SENTINEL)
218+
if new_val is not _SENTINEL and new_val != old_val:
219+
updates[field] = new_val
220+
return original.model_copy(update=updates) if updates else None
221+
```
222+
223+
### Policy Table
224+
225+
| Hook Point | Writable Fields |
226+
|------------|----------------|
227+
| **Session Lifecycle** | |
228+
| `session_pre_init` | `backend_name`, `model_id`, `model_options`, `backend_kwargs` |
229+
| `session_post_init` | *(observe-only)* |
230+
| `session_reset` | *(observe-only)* |
231+
| `session_cleanup` | *(observe-only)* |
232+
| **Component Lifecycle** | |
233+
| `component_pre_create` | `description`, `images`, `requirements`, `icl_examples`, `grounding_context`, `user_variables`, `prefix`, `template_id` |
234+
| `component_post_create` | `component` |
235+
| `component_pre_execute` | `action`, `context`, `context_view`, `requirements`, `model_options`, `format`, `strategy`, `tool_calls_enabled` |
236+
| `component_post_success` | `result` |
237+
| `component_post_error` | *(observe-only)* |
238+
| **Generation Pipeline** | |
239+
| `generation_pre_call` | `model_options`, `tools`, `format`, `formatted_prompt` |
240+
| `generation_post_call` | `processed_output`, `model_output` |
241+
| `generation_stream_chunk` | `chunk`, `accumulated` |
242+
| **Validation** | |
243+
| `validation_pre_check` | `requirements`, `model_options` |
244+
| `validation_post_check` | `results`, `all_passed` |
245+
| **Sampling Pipeline** | |
246+
| `sampling_loop_start` | `loop_budget` |
247+
| `sampling_iteration` | *(observe-only)* |
248+
| `sampling_repair` | `repair_action`, `repair_context` |
249+
| `sampling_loop_end` | `final_result` |
250+
| **Tool Execution** | |
251+
| `tool_pre_invoke` | `tool_args` |
252+
| `tool_post_invoke` | `tool_output` |
253+
| **Backend Adapter Ops** | |
254+
| `adapter_pre_load` | *(observe-only)* |
255+
| `adapter_post_load` | *(observe-only)* |
256+
| `adapter_pre_unload` | *(observe-only)* |
257+
| `adapter_post_unload` | *(observe-only)* |
258+
| **Context Operations** | |
259+
| `context_update` | *(observe-only)* |
260+
| `context_prune` | *(observe-only)* |
261+
| **Error Handling** | |
262+
| `error_occurred` | *(observe-only)* |
263+
264+
### Default Policy
265+
266+
Mellea uses `DefaultHookPolicy.DENY` as the default for hooks without an explicit policy. This means:
267+
268+
- **Hooks with an explicit policy**: Only writable fields are accepted; other changes are discarded.
269+
- **Hooks without a policy** (observe-only): All modifications are rejected with a warning log.
270+
- **Custom hooks**: Custom hooks registered by users default to `DENY`. To allow modifications, pass a `HookPayloadPolicy` when registering the custom hook type.
271+
272+
### Modification Pattern
273+
274+
Because payloads are frozen, plugins must use `model_copy(update={...})` to create a modified copy:
275+
276+
```python
277+
@hook("generation_pre_call", mode="enforce", priority=10)
278+
async def enforce_budget(payload, ctx):
279+
if (payload.estimated_tokens or 0) > 4000:
280+
return block("Token budget exceeded")
281+
282+
# Modify a writable field — use model_copy, not direct assignment
283+
modified = payload.model_copy(update={"model_options": {**payload.model_options, "max_tokens": 4000}})
284+
return PluginResult(continue_processing=True, modified_payload=modified)
285+
```
286+
287+
Attempting to set attributes directly (e.g., `payload.model_options = {...}`) raises a `FrozenModelError`.
288+
289+
### Chaining
290+
291+
When multiple plugins modify the same hook's payload, modifications are chained:
292+
293+
1. Plugin A receives the original payload, returns a modified copy.
294+
2. The policy filters Plugin A's changes to writable fields only.
295+
3. Plugin B receives the policy-filtered result from Plugin A.
296+
4. The policy filters Plugin B's changes.
297+
5. The final policy-filtered payload is returned to the caller.
298+
299+
This ensures each plugin sees the cumulative effect of prior plugins, and all modifications pass through the policy filter.
300+
171301
## 4. Hook Definitions
172302

173303
### A. Session Lifecycle Hooks
@@ -1039,20 +1169,26 @@ class ContextSnapshot:
10391169
Hooks can return different result types to control execution:
10401170

10411171
1. **Continue (no-op)**`PluginResult(continue_processing=True)` with no `modified_payload`. Execution proceeds with the original payload unchanged.
1042-
2. **Continue with modification**`PluginResult(continue_processing=True, modified_payload=...)`. Execution proceeds with the modified payload fields in place of the originals.
1172+
2. **Continue with modification**`PluginResult(continue_processing=True, modified_payload=...)`. The plugin manager applies the hook's `HookPayloadPolicy`, accepting only changes to writable fields. Execution proceeds with the policy-filtered payload.
10431173
3. **Block execution**`PluginResult(continue_processing=False, violation=...)`. Execution halts with structured error information via `PluginViolation`.
10441174

10451175
These three outcomes are exhaustive. Hooks cannot redirect control flow, throw arbitrary exceptions, or alter the calling method's logic beyond these outcomes. This is enforced by the `PluginResult` type — there is no escape hatch. The `violation` field provides structured error information but does not influence which code path runs next.
10461176

1177+
Because payloads are frozen, the `modified_payload` in option 2 must be a new object created via `payload.model_copy(update={...})`not a mutated version of the original.
1178+
10471179
### Modify Payload
10481180

10491181
```python
1182+
# Create an immutable copy with only the desired changes
1183+
modified = payload.model_copy(update={"model_options": new_options})
10501184
return PluginResult(
1051-
continue_processing=True
1052-
modified_payload=modified_payload,
1185+
continue_processing=True,
1186+
modified_payload=modified,
10531187
)
10541188
```
10551189

1190+
> **Note**: Only changes to fields listed in the hook's `HookPayloadPolicy.writable_fields` will be accepted. Changes to other fields are silently discarded by the policy enforcement layer.
1191+
10561192
### Block Execution
10571193

10581194
```python
@@ -1358,15 +1494,15 @@ class PIIRedactor:
13581494
async def redact_input(self, payload, ctx):
13591495
redacted = self._redact(payload.description)
13601496
if redacted != payload.description:
1361-
payload.description = redacted
1362-
return PluginResult(continue_processing=True, modified_payload=payload)
1497+
modified = payload.model_copy(update={"description": redacted})
1498+
return PluginResult(continue_processing=True, modified_payload=modified)
13631499

13641500
@hook("generation_post_call")
13651501
async def redact_output(self, payload, ctx):
13661502
redacted = self._redact(payload.processed_output)
13671503
if redacted != payload.processed_output:
1368-
payload.processed_output = redacted
1369-
return PluginResult(continue_processing=True, modified_payload=payload)
1504+
modified = payload.model_copy(update={"processed_output": redacted})
1505+
return PluginResult(continue_processing=True, modified_payload=modified)
13701506

13711507
def _redact(self, text: str) -> str:
13721508
for pattern in self.patterns:

0 commit comments

Comments
 (0)