feat: taint checking and security

[ambrishrawat]

This PR introduces a minimal proof-of-concept for taint and security propagation across CBlock, ModelOutputThunk, and session flows, as discussed in generative-computing/mellea#189
.

[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)(?:\(.+\))?:

[ambrishrawat]

@nrfulton quick clarifications -

1. What’s the best way for expose taint configuration to devs? e.g. when a description includes a user variable like summarise the following {{email_body}}, should taint be inferred automatically or something they can configure?
2. Would it make sense to have a global strictness setting to toggle between warnings and exceptions for taint violations? Is blocify the best place for this?

[nrfulton]

What’s the best way for expose taint configuration to devs? e.g. when a description includes a user variable like summarise the following {{email_body}}, should taint be inferred automatically or something they can configure?

We should infer automatically where-ever possible. I nthis case, I'm not sure how you would infer taint. I guess you assumption here is that email_boy -- or any user_variable input -- should entail taint?

[ambrishrawat]

Yes, that was the thinking. Making it configurable may make more sense for taint. Any thoughts on the best way to expose that? Would love your take on the code too.

[davidcox]

If there is a tainted variable in the context, everything downstream should get tainted. As for how variables get tainted in the first place, a common way people do this is to define sources, sinks, and (optionally) washers. These are wrappers around interfaces that produce sensitive data (e.g. HR database api), or where it enters an unsafe place (e.g. sending to a UI).

[nrfulton]

1. CBlocks should be immutable.
2. Naming a variable Ccmponent and assigning it to a CBlock is confusing.
3. The cyclic reference here is a bit confusing and invites buggy code. Use tainted_by(None) instead of tainted_by(self) for the root node.

c = CBlock("user input", sec_level=SecLevel.tained_by(None))

[ambrishrawat]

Updated this; tainted_by(None) for root now

[nrfulton]

Why not c.sec_level?

[ambrishrawat]

Defined it as a property and now this works as c.sec_level.is_tainted()

[nrfulton]

Unless the example critically depends on using a particular model, always use session = start_session() instead. This makes the examples easier to maintain.

[nrfulton]

We should use is_tainted instead of is_safe. The meaning of safe is very ambiguous.

[ambrishrawat]

Removed all instances of is_safe

[nrfulton]

Especially in a module called security.core, we should avoid use of magic strings.

[ambrishrawat]

Created SecLevelType enum

[nrfulton]

Instead us something like:

match action:
     case Component...

    case CBlock...

(If type(action) :> Component then check is not necessary because the Component protocol has a parts() method. )

[ambrishrawat]

Updated it to use match/case

[ambrishrawat]

Thanks for the review @nrfulton !
I have incorporated your suggestions. Appreciate another pass when you get the chance

[guicho271828]

hi ambrish!

[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)(?:\(.+\))?:

[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)(?:\(.+\))?:

[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)(?:\(.+\))?:

[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)(?:\(.+\))?:

[planetf1]

I'm wondering if there's an issue with here in not looking recursively (ie via taint_sources() )? the args may be ok, but we could have a tainted Component?

[ambrishrawat]

Ah I see that, yes need to call taint_sources here and do a recursive check for classified sources as well. I will update this.

[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)(?:\(.+\))?:

[planetf1]

Only CBlocks or any subclass? If the latter, we return a CBlock which would lose the subclass info?
For example could it be an ImageBlock?
Also note that ImageBlock never calls super().init() which means the value that is being copied does not exist == crash?

[ambrishrawat]

Right, this particular declassify is merely illustrative, hence only applicable to CBlocks. One would need bespoke declassification methods for applications to sparsify taint information across the code. That was my initial thinking

[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)(?:\(.+\))?:

[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)(?:\(.+\))?:

[ambrishrawat]

@planetf1 Updated the logic for privileged. It recursively checks for both taint and classified.
cc: @nrfulton @jakelorocco

[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)(?:\(.+\))?:

[planetf1]

I think there are also some

• uses of Any where the types should be stricter
• incorrect types
• missing type annotations

I know we have omissions elsewhere, but it feels that the focus on security here makes it worth paying close attention.

Need to have a fiddle with mypy/py (the default rules are not strict enough) to be sure

[planetf1]

Suggested change

super().__init__(value, meta)

I think this is left over from conflict resolution as it's a double init() ?

[ambrishrawat]

Added sec_level here

[planetf1]

We'll now return none -- inconsistent with annotations, and likely to cause a crash?

[ambrishrawat]

Corrected

[planetf1]

left over from conflict since we're assigning a few lines later. And what about assigning _start (not checked if we do that much further down). We'd lose it?

[planetf1]

is this another dup ? similar -- worth checking for this pattern through the code?

[planetf1]

As per Nathan's comment above there's another example. I think we should be explicit - especially given it's security. May be worth doing a more thorough check on type annotations and use of Any?

[planetf1]

what else could happen here? Do we need more checks for other exceptions? Or log something at least? I guess we're trying to be non-intrusive (need to understand code better to suggest exact changes!)

[planetf1]

so any contribs/older code outside mellea would be impacted -- should we ensure we capture in release notes/doc as a breaking change

[planetf1]

this works on a class but fails on an instance - other parts of the code have checks for class vs instance to do the right thing, but this is omitted here

[jakelorocco]

I've started taking a look at this. Will finish my review soon.

[araujof]

Cool work, @ambrishrawat @nrfulton! I found this PR through the originating discussion (#189).

Here are my observations, including some suggestions around synergies with the plugin system that landed after this PR was originally proposed. The PR has good foundational work on the SecLevel type system and the TaintChecking protocol. The main suggestion is to keep the security module (mellea/security/) as the data model but move enforcement and propagation into the plugin system. This would reduce the diff significantly (no backend changes needed), make security opt-in, and align with how other cross-cutting concerns like telemetry and metrics are already implemented.

I would be happy to iterate on this together and collaborate on refactoring the implementation to integrate with the plugin system. Let me know if you'd like to pair on it.

Suggested refactoring

Mellea now has a plugin system with lifecycle hooks that did not exist when this PR was conceived. The overlap is significant, and I think leveraging it would simplify the implementation considerably.

A. Replace backend boilerplate with a GENERATION_POST_CALL plugin

The GENERATION_POST_CALL hook fires after every generation with access to the ModelOutputThunk. A TaintPropagationPlugin could centralize all taint computation:

class TaintPropagationPlugin(Plugin, name="taint_propagation", priority=10):
    @hook("generation_post_call", mode=PluginMode.TRANSFORM)
    async def propagate_taint(self, payload: GenerationPostCallPayload, context):
        mot = payload.model_output
        sources = taint_sources(mot._action, mot._context)
        if sources:
            mot._sec_level = SecLevel.tainted_by(sources)
        else:
            mot._sec_level = SecLevel.none()

This eliminates all 20+ identical blocks from the backends, keeps backends security-unaware, makes taint tracking opt-in, and allows users to swap in custom propagation logic.

B. @privileged as a TOOL_PRE_INVOKE guard

Rather than decorating individual functions, a security plugin could intercept all tool invocations via TOOL_PRE_INVOKE, checking arguments for taint before execution.

C. COMPONENT_PRE_EXECUTE for declarative security policies

The component_pre_execute hook fires before aact() with access to the action, context, and requirements. A security plugin could block execution based on taint thresholds, inject sanitization requirements, or log security audit trails via FIRE_AND_FORGET mode.

D. Session-scoped security via context managers

The plugin system's context-manager pattern makes security opt-in and scoped:

with TaintPropagationPlugin(scan_depth=10, strict=True):
    result = session.instruct("Process sensitive data")
    # Taint automatically tracked

Minor issues (most addressed by the refactoring above)

Identical taint logic repeated across all backends

Every backend now has this 3-line block repeated in every generation method:

sources = taint_sources(action, ctx)
sec_level = SecLevel.tainted_by(sources) if sources else SecLevel.none()
output = ModelOutputThunk(value=None, sec_level=sec_level, meta={})

If the security logic changes, all call sites must be updated in lockstep. Taint propagation could be made systematic, not sprinkled manually. See the proposal refactoring above for a cleaner alternative.

Boilerplate sec_level property on every Component (~12 classes)

Every component class got a copy-pasted block:

self._sec_level: SecLevel | None = None

@property
def sec_level(self) -> SecLevel | None:
    return self._sec_level

The _sec_level is always initialized to None with no constructor parameter to set it, so for components this appears to be dead code. The PR already adds a default sec_level property returning None on the Component protocol (base.py), which means the per-class overrides can simply be removed. Only classes that actually construct with a non-None value need an override.

Context scanning is shallow and hardcoded to 5

context_items = ctx.as_list(last_n_components=5)

The last_n_components=5 limit is arbitrary and undocumented. A tainted item at position 6 in context is silently ignored. For a security feature, this gives false confidence. The test test_taint_sources_shallow_search_limit validates this behavior as correct, but it is a soundness gap.

Suggestion: either scan the full context by default with an opt-in limit, or document this as a known limitation.

Silent except Exception: pass in security code

try:
    context_items = ctx.as_list(last_n_components=5)
    ...
except Exception:
    pass

A blanket exception swallow in security-critical code. If as_list() fails, taint sources in the context are silently ignored, producing a false negative in a security analysis. At minimum this should log a warning; better to let it propagate.

SecLevel.data is untyped Any

def __init__(self, level_type: SecLevelType | str, data: Any = None):

The data field holds either an AccessType[T], a list[CBlock | Component], or None depending on level_type. This is a discriminated union case that would benefit from stricter typing, especially given the reviewers' feedback about type annotations in security code.

mification sets a property descriptor without checking class vs instance

sec_level_prop = property(
    lambda self: getattr(self, "_sec_level", None),
)
setattr(obj, "sec_level", sec_level_prop)

The existing mification code distinguishes between classes and instances when binding methods, but the new sec_level property addition does not. When mify() is called on an instance (which it is, e.g., in functional.py:370 and several tests), setattr stores the property object as a regular attribute rather than installing it as a descriptor. In that case obj.sec_level returns the property object itself instead of calling its getter. The class-decorator path (@mify) works correctly since property descriptors on classes resolve normally.