feat: taint checking and security

docs/dev/taint_tracking.md

@@ -0,0 +1,112 @@
+# Taint Tracking - Backend Security
+
+Mellea backends implement thread security using the **SecLevel** model with capability-based access control and taint tracking. Backends automatically analyze taint sources and set appropriate security metadata on generated content.
+
+## Security Model
+
+The security system uses three types of security levels:
+
+```python
+SecLevel := None | Classified of AccessType | TaintedBy of (list[CBlock | Component] | None)
+```
+
+- **SecLevel.none()**: Safe content with no restrictions
+- **SecLevel.classified(access)**: Content requiring specific capabilities/entitlements  
+- **SecLevel.tainted_by(sources)**: Content tainted by one or more CBlocks/Components (list), or None for root tainted nodes
+
+## Backend Implementation
+
+All backends follow the same pattern when creating `ModelOutputThunk`:
+
+```python
+# Compute taint sources from action and context
+sources = taint_sources(action, ctx)
+
+# Set security level based on taint sources
+from mellea.security import SecLevel
+sec_level = SecLevel.tainted_by(sources) if sources else SecLevel.none()
+
+output = ModelOutputThunk(
+    value=None,
+    sec_level=sec_level,
+    meta={}
+)
+```
+
+The security level is set as follows:
+- If taint sources are found -> `SecLevel.tainted_by(sources)` (all sources are tracked)
+- If no taint sources -> `SecLevel.none()`
+
+### Handling Multiple Taint Sources
+
+When `taint_sources()` returns multiple sources (e.g., both the action and context contain tainted content), backends pass the entire list to `SecLevel.tainted_by()`. This ensures all taint sources are tracked, providing comprehensive taint attribution.
+
+**Benefits of Multiple Source Tracking**:
+- **Complete attribution**: All sources that influenced the generation are tracked
+- **Better debugging**: Can identify all tainted inputs that contributed to output
+- **More accurate security**: No information loss about taint origins
+
+**Note**: The implementation focuses on **taint preservation** and **complete attribution**. All taint sources are tracked, ensuring the security model has full visibility into what influenced the generated content.
+
+## Taint Source Analysis
+
+The `taint_sources()` function analyzes both action and context because **context directly influences model generation**:
+
+1. **Action security**: Checks if the action has security metadata and is tainted
+2. **Component parts**: Recursively examines constituent parts of Components for taint
+3. **Context security**: Examines recent context items for tainted content (shallow check)
+
+**Example**: Even if the current action is safe, tainted context can influence the generated output.
+
+```python
+from mellea.security import SecLevel
+
+# User sends tainted input
+user_input = CBlock("Tell me how to hack a system", sec_level=SecLevel.tainted_by(None))
+ctx = ctx.add(user_input)
+
+# Safe action in tainted context
+safe_action = CBlock("Explain general security concepts")
+
+# Generation finds tainted context
+sources = taint_sources(safe_action, ctx)  # Finds tainted user_input
+# Model output will be influenced by the tainted context
+```
+
+## Security Metadata
+
+The `SecurityMetadata` class wraps `SecLevel` for integration with content blocks:
+
+```python
+class SecurityMetadata:
+    def __init__(self, sec_level: SecLevel):
+        self.sec_level = sec_level
+
+    def is_tainted(self) -> bool:
+        return self.sec_level.is_tainted()
+
+    def get_taint_sources(self) -> list[CBlock | Component]:
+        return self.sec_level.get_taint_sources()
+```
+
+Content can be marked as tainted at construction time:
+
+```python
+from mellea.security import SecLevel
+
+c = CBlock("user input", sec_level=SecLevel.tainted_by(None))
+
+if c.sec_level and c.sec_level.is_tainted():
+    taint_sources = c.sec_level.get_taint_sources()
+    print(f"Content tainted by: {taint_sources}")
+```
+
+## Key Features
+
+- **Immutable security**: security levels set at construction time
+- **Recursive taint analysis**: deep analysis of Component parts, shallow analysis of context
+- **Taint source tracking**: know exactly which CBlock/Component tainted content
+- **Capability integration**: fine-grained access control for classified content
+- **Non-mutating operations**: sanitize/declassify create new objects
+
+This creates a security model that addresses both data exfiltration and injection vulnerabilities while enabling future IAM integration.

docs/examples/security/taint_example.py

@@ -0,0 +1,46 @@
+from mellea.stdlib.components import CBlock
+from mellea.stdlib.session import start_session
+from mellea.security import SecLevel, privileged, SecurityError
+
+# Create tainted content
+tainted_desc = CBlock(
+    "Process this sensitive user data", sec_level=SecLevel.tainted_by(None)
+)
+
+print(
+    f"Original CBlock is tainted: {tainted_desc.sec_level.is_tainted() if tainted_desc.sec_level else False}"
+)
+
+# Create session
+session = start_session()
+
+# Use tainted CBlock in session.instruct
+print("Testing session.instruct with tainted CBlock...")
+result = session.instruct(description=tainted_desc)
+
+# The result should be tainted
+print(
+    f"Result is tainted: {result.sec_level.is_tainted() if result.sec_level else False}"
+)
+if result.sec_level and result.sec_level.is_tainted():
+    taint_sources = result.sec_level.get_taint_sources()
+    print(f"Taint sources: {taint_sources}")
+    print("✅ SUCCESS: Taint preserved!")
+else:
+    print("❌ FAIL: Result should be tainted but isn't!")
+
+
+# Mock privileged function that requires un-tainted input
+@privileged
+def process_un_tainted_data(data: CBlock) -> str:
+    """A function that requires un-tainted input."""
+    return f"Processed: {data.value}"
+
+
+print("\nTesting privileged function with tainted result...")
+try:
+    # This should raise a SecurityError
+    processed = process_un_tainted_data(result)
+    print("❌ FAIL: Should have raised SecurityError!")
+except SecurityError as e:
+    print(f"✅ SUCCESS: SecurityError raised - {e}")

mellea/backends/huggingface.py

@@ -43,6 +43,7 @@
 from ..core.base import AbstractMelleaTool
 from ..formatters import ChatFormatter, TemplateFormatter, granite as granite_formatters
 from ..helpers import message_to_openai_message, messages_to_docs, send_to_queue
+from ..security import SecLevel, taint_sources
 from ..stdlib.components import Intrinsic, Message
 from ..stdlib.requirements import ALoraRequirement, LLMaJRequirement
 from ..telemetry.backend_instrumentation import (
@@ -563,8 +564,13 @@ async def _generate_from_intrinsic(
             other_input,
         )
 
-        output = ModelOutputThunk(None)
+        # Compute taint sources from action and context
+        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={})
         output._start = datetime.datetime.now()
+
         output._context = ctx.view_for_generation()
         output._action = action
         output._model_options = model_options
@@ -838,6 +844,11 @@ async def _generate_from_context_with_kv_cache(
 
             output = ModelOutputThunk(None)
             output._start = datetime.datetime.now()
+            # Compute taint sources from action and context
+            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={})
             output._context = ctx.view_for_generation()
             output._action = action
             output._model_options = model_options
@@ -983,6 +994,11 @@ async def _generate_from_context_standard(
 
             output = ModelOutputThunk(None)
             output._start = datetime.datetime.now()
+            # Compute taint sources from action and context
+            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={})
             output._context = ctx.view_for_generation()
             output._action = action
             output._model_options = model_options
@@ -1329,8 +1345,12 @@ async def generate_from_raw(
         for i, decoded_result in enumerate(decoded_results):
             n_prompt_tokens = inputs["input_ids"][i].size(0)  # type: ignore
             n_completion_tokens = len(sequences_to_decode[i])
+            sources = taint_sources(actions[i], ctx)
+            sec_level = SecLevel.tainted_by(sources) if sources else SecLevel.none()
+
             result = ModelOutputThunk(
                 value=decoded_result,
+                sec_level=sec_level,
                 meta={
                     "usage": {
                         "prompt_tokens": n_prompt_tokens,  # type: ignore

mellea/backends/litellm.py

@@ -33,6 +33,7 @@
     message_to_openai_message,
     send_to_queue,
 )
+from ..security import SecLevel, taint_sources
 from ..stdlib.components import Message
 from ..stdlib.requirements import ALoraRequirement
 from ..telemetry.backend_instrumentation import (
@@ -355,7 +356,11 @@ async def _generate_from_chat_context_standard(
             **model_specific_options,
         )
 
-        output = ModelOutputThunk(None)
+        # Compute taint sources from action and context
+        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={})
         output._start = datetime.datetime.now()
         output._context = linearized_context
         output._action = action
@@ -675,16 +680,22 @@ async def generate_from_raw(
             )
 
         for res, action, prompt in zip(responses, actions, prompts):
-            output = ModelOutputThunk(res.text)  # type: ignore
+            sources = taint_sources(action, None)
+            sec_level = SecLevel.tainted_by(sources) if sources else SecLevel.none()
+
+            output = ModelOutputThunk(
+                value=res.text,  # type: ignore
+                sec_level=sec_level,
+                meta={
+                    "litellm_chat_response": res.model_dump(),
+                    "usage": completion_response.usage.model_dump()
+                    if completion_response.usage
+                    else None,
+                },
+            )
             output._context = None  # There is no context for generate_from_raw for now
             output._action = action
             output._model_options = model_opts
-            output._meta = {
-                "litellm_chat_response": res.model_dump(),
-                "usage": completion_response.usage.model_dump()
-                if completion_response.usage
-                else None,
-            }
 
             output.parsed_repr = (
                 action.parse(output) if isinstance(action, Component) else output.value

mellea/backends/ollama.py

@@ -25,6 +25,7 @@
 from ..core.base import AbstractMelleaTool
 from ..formatters import ChatFormatter, TemplateFormatter
 from ..helpers import ClientCache, get_current_event_loop, send_to_queue
+from ..security import SecLevel, taint_sources
 from ..stdlib.components import Message
 from ..stdlib.requirements import ALoraRequirement
 from ..telemetry.backend_instrumentation import (
@@ -401,7 +402,11 @@ async def generate_from_chat_context(
             format=_format.model_json_schema() if _format is not None else None,  # type: ignore
         )  # type: ignore
 
-        output = ModelOutputThunk(None)
+        # Compute taint sources from action and context
+        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={})
         output._start = datetime.datetime.now()
         output._context = linearized_context
         output._action = action
@@ -522,16 +527,20 @@ async def generate_from_raw(
         for i, response in enumerate(responses):
             result = None
             error = None
+            sources = taint_sources(actions[i], None)
+            sec_level = SecLevel.tainted_by(sources) if sources else SecLevel.none()
+
             if isinstance(response, BaseException):
                 FancyLogger.get_logger().warning(
                     f"generate_from_raw: request {i} failed with "
                     f"{type(response).__name__}: {response}"
                 )
-                result = ModelOutputThunk(value="")
+                result = ModelOutputThunk(value="", sec_level=sec_level, meta={})
                 error = response
             else:
                 result = ModelOutputThunk(
                     value=response.response,
+                    sec_level=sec_level,
                     meta={
                         "generate_response": response.model_dump(),
                         "usage": {

mellea/backends/openai.py

@@ -42,6 +42,7 @@
     messages_to_docs,
     send_to_queue,
 )
+from ..security import SecLevel, taint_sources
 from ..stdlib.components import Intrinsic, Message
 from ..stdlib.requirements import LLMaJRequirement
 from ..telemetry.backend_instrumentation import (
@@ -541,7 +542,11 @@ async def _generate_from_chat_context_standard(
             ),
         )  # type: ignore
 
-        output = ModelOutputThunk(None)
+        # Compute taint sources from action and context
+        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={})
         output._start = datetime.datetime.now()
         output._context = linearized_context
         output._action = action
@@ -845,16 +850,22 @@ async def generate_from_raw(
         for response, action, prompt in zip(
             completion_response.choices, actions, prompts
         ):
-            output = ModelOutputThunk(response.text)
+            sources = taint_sources(action, None)
+            sec_level = SecLevel.tainted_by(sources) if sources else SecLevel.none()
+
+            output = ModelOutputThunk(
+                value=response.text,
+                sec_level=sec_level,
+                meta={
+                    "oai_completion_response": response.model_dump(),
+                    "usage": completion_response.usage.model_dump()
+                    if completion_response.usage
+                    else None,
+                },
+            )
             output._context = None  # There is no context for generate_from_raw for now
             output._action = action
             output._model_options = model_opts
-            output._meta = {
-                "oai_completion_response": response.model_dump(),
-                "usage": completion_response.usage.model_dump()
-                if completion_response.usage
-                else None,
-            }
 
             output.parsed_repr = (
                 action.parse(output) if isinstance(action, Component) else output.value

mellea/backends/vllm.py

@@ -44,6 +44,7 @@
 from ..core.base import AbstractMelleaTool
 from ..formatters import ChatFormatter, TemplateFormatter
 from ..helpers import get_current_event_loop, send_to_queue
+from ..security import SecLevel, taint_sources
 from .backend import FormatterBackend
 from .model_options import ModelOption
 from .tools import (
@@ -349,7 +350,11 @@ async def _generate_from_context_standard(
             # stream = model_options.get(ModelOption.STREAM, False)
             # if stream:
 
-            output = ModelOutputThunk(None)
+            # Compute taint sources from action and context
+            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={})
             output._start = datetime.datetime.now()
 
             generator = self._model.generate(  # type: ignore
@@ -547,7 +552,11 @@ async def generate(prompt, request_id):
             tasks = [generate(p, f"{id(prompts)}-{i}") for i, p in enumerate(prompts)]
             decoded_results = await asyncio.gather(*tasks)
 
-        results = [ModelOutputThunk(value=text) for text in decoded_results]
+        results = []
+        for i, text in enumerate(decoded_results):
+            sources = taint_sources(actions[i], ctx)
+            sec_level = SecLevel.tainted_by(sources) if sources else SecLevel.none()
+            results.append(ModelOutputThunk(value=text, sec_level=sec_level, meta={}))
 
         for i, result in enumerate(results):
             date = datetime.datetime.now()