Split the snapshot codec into plaintext and encrypted variants

snapshot_codec() now always returns a codec, so save encodes
unconditionally and the restore path reads codec.protected to pick
between the fail-open and fail-closed contracts, instead of testing the
codec for None.

🤖 Generated with Claude Code

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

fastmcp_tasks/fastmcp_tasks/context.py

@@ -235,10 +235,7 @@ class TaskContextSnapshot:
         read it (#4747).
         """
         key = docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:snapshot")
-        codec = snapshot_codec()
-        payload = self.to_json()
-        if codec is not None:
-            payload = codec.encode(payload)
+        payload = snapshot_codec().encode(self.to_json())
         async with docket.redis() as redis:
             await redis.set(key, payload, ex=ttl_seconds)
 
@@ -313,8 +310,8 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None:
     from fastmcp_tasks.dependencies import _current_docket
 
     # Resolved before anything can fail: a misconfigured key (e.g. an empty
-    # string) raises here and fails the task, and the except blocks below read
-    # it to pick between the fail-open and fail-closed contracts.
+    # string) raises here and fails the task, and the branches below read
+    # `codec.protected` to pick between the fail-open and fail-closed contracts.
     codec = snapshot_codec()
 
     try:
@@ -324,7 +321,7 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None:
     if docket is None:
         docket = _current_docket.get()
     if docket is None:
-        if codec is not None:
+        if codec.protected:
             raise RuntimeError(
                 "No Docket backend is available to retrieve the protected "
                 "task snapshot, so the submitting caller cannot be recovered."
@@ -339,15 +336,13 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None:
                 docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:snapshot")
             )
         if raw is None:
-            if codec is None:
+            if not codec.protected:
                 return
             raise RuntimeError(
                 "The task's context snapshot is missing (its TTL may have "
                 "expired), so the submitting caller cannot be recovered."
             )
-        if codec is not None:
-            raw = codec.decode(raw)
-        snapshot = TaskContextSnapshot.from_json(raw)
+        snapshot = TaskContextSnapshot.from_json(codec.decode(raw))
         _remember_snapshot(task_id, snapshot)
         # Restore the ambient request context (auth token, headers) so core's
         # get_access_token()/get_http_headers() see the submitting caller inside
@@ -365,7 +360,7 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None:
         )
         raise
     except Exception:
-        if codec is not None:
+        if codec.protected:
             _logger.error(
                 "Failed to restore the protected task snapshot for %s. The task "
                 "will fail rather than run without the submitting caller's "

fastmcp_tasks/fastmcp_tasks/encryption.py

@@ -12,7 +12,9 @@ that restores a snapshot is rarely the one that captured it.
 
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
 from functools import lru_cache
+from typing import ClassVar
 
 import fastmcp
 from fastmcp.utilities.logging import get_logger
@@ -38,8 +40,39 @@ class SnapshotDecryptionError(Exception):
     """
 
 
-class SnapshotCodec:
-    """Encrypts and decrypts snapshot payloads with a key derived from material.
+class SnapshotCodec(ABC):
+    """Transforms snapshot payloads on their way to and from the backend.
+
+    ``protected`` tells the restore path which failure contract applies: a
+    protected snapshot that cannot be restored fails the task, an unprotected
+    one degrades to an anonymous run with a warning.
+    """
+
+    protected: ClassVar[bool]
+
+    @abstractmethod
+    def encode(self, payload: str) -> str:
+        """Return the stored form of a serialized snapshot."""
+
+    @abstractmethod
+    def decode(self, stored: str | bytes) -> str:
+        """Return the serialized snapshot a stored value holds."""
+
+
+class PlaintextCodec(SnapshotCodec):
+    """Stores snapshots as-is; the contract when no encryption key is set."""
+
+    protected = False
+
+    def encode(self, payload: str) -> str:
+        return payload
+
+    def decode(self, stored: str | bytes) -> str:
+        return stored.decode() if isinstance(stored, bytes) else stored
+
+
+class EncryptedCodec(SnapshotCodec):
+    """Encrypts snapshot payloads with a key derived from material.
 
     The material is a string from the environment, and nothing about a string
     proves it is random, so it is always treated as low-entropy: the Fernet key
@@ -47,6 +80,8 @@ class SnapshotCodec:
     once per process (see ``_codec_for``).
     """
 
+    protected = True
+
     def __init__(self, material: str) -> None:
         from cryptography.fernet import Fernet
 
@@ -90,21 +125,24 @@ class SnapshotCodec:
             ) from e
 
 
+_PLAINTEXT_CODEC = PlaintextCodec()
+
+
 @lru_cache(maxsize=4)
-def _codec_for(material: str) -> SnapshotCodec:
+def _codec_for(material: str) -> EncryptedCodec:
     """One codec per key, so the derivation cost is paid once per process.
 
     The PBKDF2 stretch takes about a second, and every task submission and
     every restore needs a codec.
     """
-    return SnapshotCodec(material)
+    return EncryptedCodec(material)
 
 
-def snapshot_codec() -> SnapshotCodec | None:
-    """The codec for the configured key, or ``None`` when none is configured."""
+def snapshot_codec() -> SnapshotCodec:
+    """The codec for the configured key; the plaintext codec when none is set."""
     key = fastmcp.settings.encryption_key
     if key is None:
-        return None
+        return _PLAINTEXT_CODEC
     return _codec_for(key.get_secret_value())
 
 

tests/tasks/server/test_snapshot_encryption.py

@@ -18,7 +18,8 @@ from unittest.mock import patch
 import pytest
 from fastmcp_tasks.context import TaskContextSnapshot
 from fastmcp_tasks.encryption import (
-    SnapshotCodec,
+    EncryptedCodec,
+    PlaintextCodec,
     SnapshotDecryptionError,
     clear_codec_cache,
     snapshot_codec,
@@ -84,41 +85,49 @@ def sensitive_snapshot() -> TaskContextSnapshot:
 
 class TestSnapshotCodec:
     def test_round_trips_a_payload(self):
-        codec = SnapshotCodec(KEY)
+        codec = EncryptedCodec(KEY)
         assert codec.decode(codec.encode('{"a": 1}')) == '{"a": 1}'
 
     def test_encoded_payload_hides_the_credentials(
         self, sensitive_snapshot: TaskContextSnapshot
     ):
-        encoded = SnapshotCodec(KEY).encode(sensitive_snapshot.to_json())
+        encoded = EncryptedCodec(KEY).encode(sensitive_snapshot.to_json())
         assert "token-client-a-user-1" not in encoded
         assert "authorization" not in encoded
 
     def test_decode_rejects_another_keys_payload(self):
-        encoded = SnapshotCodec(OTHER_KEY).encode('{"a": 1}')
+        encoded = EncryptedCodec(OTHER_KEY).encode('{"a": 1}')
         with pytest.raises(SnapshotDecryptionError):
-            SnapshotCodec(KEY).decode(encoded)
+            EncryptedCodec(KEY).decode(encoded)
 
     def test_decode_rejects_plaintext(self):
         """A snapshot written before the key was set must not be trusted."""
         with pytest.raises(SnapshotDecryptionError):
-            SnapshotCodec(KEY).decode('{"access_token_json": null}')
+            EncryptedCodec(KEY).decode('{"access_token_json": null}')
 
     def test_empty_material_is_rejected(self):
         """An empty key would derive a universally reproducible Fernet key."""
         with pytest.raises(ValueError, match="must not be empty"):
-            SnapshotCodec("")
+            EncryptedCodec("")
 
     def test_decode_accepts_bytes(self):
         """Redis hands back bytes on some backends."""
-        codec = SnapshotCodec(KEY)
+        codec = EncryptedCodec(KEY)
         assert codec.decode(codec.encode('{"a": 1}').encode()) == '{"a": 1}'
 
     def test_same_key_reuses_one_codec(self, encryption_key: str):
         assert snapshot_codec() is snapshot_codec()
 
-    def test_no_codec_without_a_key(self, no_encryption_key: None):
-        assert snapshot_codec() is None
+    def test_plaintext_codec_without_a_key(self, no_encryption_key: None):
+        codec = snapshot_codec()
+        assert isinstance(codec, PlaintextCodec)
+        assert not codec.protected
+
+    def test_plaintext_codec_is_a_pass_through(self):
+        codec = PlaintextCodec()
+        assert codec.encode('{"a": 1}') == '{"a": 1}'
+        assert codec.decode('{"a": 1}') == '{"a": 1}'
+        assert codec.decode(b'{"a": 1}') == '{"a": 1}'
 
 
 class TestSnapshotSerialization:
@@ -227,7 +236,7 @@ class TestEncryptedSnapshotRoundTrip:
         cause has to come from the log.
         """
         token = make_access_token("client-a", "user-1")
-        tampered = SnapshotCodec(OTHER_KEY).encode(TaskContextSnapshot().to_json())
+        tampered = EncryptedCodec(OTHER_KEY).encode(TaskContextSnapshot().to_json())
 
         with caplog.at_level(logging.ERROR, logger="fastmcp_tasks.context"):
             async with running_task_server(echo_token_server):
@@ -348,7 +357,7 @@ class TestTaskStillResolvesAfterFailure:
     ):
         """A fail-closed task is still a well-formed `tasks/get` result."""
         token = make_access_token("client-a", "user-1")
-        tampered = SnapshotCodec(OTHER_KEY).encode(TaskContextSnapshot().to_json())
+        tampered = EncryptedCodec(OTHER_KEY).encode(TaskContextSnapshot().to_json())
 
         async with running_task_server(echo_token_server):
             created = await submit_task(