diff --git a/docs/more/settings.mdx b/docs/more/settings.mdx index 6ea808af7..e6ea0f584 100644 --- a/docs/more/settings.mdx +++ b/docs/more/settings.mdx @@ -81,7 +81,7 @@ These control how the server listens when running with an HTTP transport. ## Tasks (Docket) -Task settings (the `FASTMCP_DOCKET_` variables) moved to the optional `fastmcp-tasks` package. See [server tasks](/servers/tasks) for configuration. +Task settings (the `FASTMCP_DOCKET_` and `FASTMCP_TASKS_` variables) live in the optional `fastmcp-tasks` package. See [server tasks](/servers/tasks) for configuration, including `FASTMCP_TASKS_ENCRYPTION_KEY` for [encrypting task snapshots at rest](/servers/tasks#credentials-at-rest). ## Security diff --git a/docs/servers/tasks.mdx b/docs/servers/tasks.mdx index b22a7b91d..7b374473d 100644 --- a/docs/servers/tasks.mdx +++ b/docs/servers/tasks.mdx @@ -162,6 +162,7 @@ mcp.add_extension(TasksExtension(url="redis://localhost:6379/0", concurrency=20) | `FASTMCP_DOCKET_URL` | `memory://` | Backend URL (`memory://` or `redis://host:port/db`) | | `FASTMCP_DOCKET_NAME` | `fastmcp` | Queue name. Servers and workers sharing a name and URL share a queue. | | `FASTMCP_DOCKET_CONCURRENCY` | `10` | Maximum concurrent tasks per worker. | +| `FASTMCP_TASKS_ENCRYPTION_KEY` | (unset) | Encrypts [task context snapshots at rest](#credentials-at-rest). Every server and worker sharing a queue must set the same key. | ## Backends @@ -193,6 +194,28 @@ mcp.add_extension(TasksExtension(url="redis://localhost:6379/0")) - **Fast**: Single-digit millisecond task pickup latency - **Scalable**: Add workers to distribute load across processes or machines +### Credentials at Rest + +A background task runs long after the request that submitted it has ended, but it still needs to know who asked for the work. FastMCP captures that identity at submission time in a **task context snapshot**: the caller's access token and every inbound HTTP header, including `Authorization`. The worker restores the snapshot before the tool body runs, so `get_access_token()` and `get_http_headers()` return the submitting caller. + +That snapshot lives in the backend for the task's TTL. With `memory://` it never leaves the process. With Redis or Valkey it is a stored value, and by default it is stored as plaintext JSON. A `rediss://` URL encrypts the connection, not the data the backend holds. Anyone who can read the backend can read the tokens. + +Set `FASTMCP_TASKS_ENCRYPTION_KEY` to encrypt the snapshot before it is written: + +```bash +export FASTMCP_TASKS_ENCRYPTION_KEY=$(python -c "import secrets; print(secrets.token_urlsafe(32))") +``` + + +Every server and worker on the same queue must set the same key. The process that restores a snapshot is rarely the one that captured it, and a worker with the wrong key cannot recover the caller. + + +With a key configured, restore **fails closed**: a worker that cannot decrypt a snapshot fails the task instead of running the tool with no identity. This matters for a tool whose behavior depends on the caller: running it as an anonymous user is worse than not running it. The failure is reported to the client as a task error, and the server log names the key mismatch. + +Two consequences of failing closed are worth planning for. Tasks submitted before the key was set fail when a worker with the key picks them up, so drain the queue before you roll a key out. Rotating a key does the same to tasks in flight under the old one. + +The key protects the snapshot only. Tool arguments and any answers a task gathers through [mid-task input](#gathering-input-mid-task) are still stored as plaintext, so treat the backend as sensitive regardless. + ## Workers Every FastMCP server with task-enabled tools automatically starts an **embedded worker**. You do not need to start a separate worker process for tasks to execute. diff --git a/fastmcp_tasks/fastmcp_tasks/context.py b/fastmcp_tasks/fastmcp_tasks/context.py index 6a4c8b514..64d6f474b 100644 --- a/fastmcp_tasks/fastmcp_tasks/context.py +++ b/fastmcp_tasks/fastmcp_tasks/context.py @@ -16,6 +16,7 @@ from contextvars import ContextVar from dataclasses import dataclass from typing import TYPE_CHECKING +from fastmcp_tasks.encryption import SnapshotDecryptionError, snapshot_codec from fastmcp_tasks.keys import ( leg_number_from_key, parse_task_key, @@ -133,6 +134,28 @@ def get_task_leg_number() -> int: return 1 +def _snapshot_redis_key(docket: Docket, task_scope: str | None, task_id: str) -> str: + """The Redis key holding a task's context snapshot.""" + return docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:snapshot") + + +async def refresh_snapshot_ttl( + docket: Docket, task_scope: str | None, task_id: str, ttl_seconds: int +) -> None: + """Slide the snapshot key's TTL alongside the task's routing keys. + + An actively polled task refreshes its metadata and leg pointers on every + ``tasks/get``, and the snapshot must live just as long: a re-entered leg + restores the submitting caller from it. Without the refresh, a task parked + on input past the snapshot's creation-time TTL loses the caller, which + means an unauthenticated run without encryption and a failed task with it. + """ + async with docket.redis() as redis: + await redis.expire( + _snapshot_redis_key(docket, task_scope, task_id), ttl_seconds + ) + + @dataclass(frozen=True, slots=True) class TaskContextSnapshot: """All context data snapshotted at task-submission time. @@ -226,10 +249,17 @@ class TaskContextSnapshot: task_id: str, ttl_seconds: int, ) -> None: - """Store this snapshot as a single Redis key.""" - key = docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:snapshot") + """Store this snapshot as a single Redis key. + + The stored value is encrypted when a ``FASTMCP_TASKS_ENCRYPTION_KEY`` is + configured: this payload carries the caller's bearer token and headers, + and a distributed backend keeps it where the backend's operators can + read it (#4747). + """ + key = _snapshot_redis_key(docket, task_scope, task_id) + payload = snapshot_codec().encode(self.to_json()) async with docket.redis() as redis: - await redis.set(key, self.to_json(), ex=ttl_seconds) + await redis.set(key, payload, ex=ttl_seconds) # Cache keyed by task_id so stale entries from previous tasks in the same @@ -285,6 +315,12 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None: backend work transparently (#3897). Failures are non-fatal: the task still runs, and sync helpers return ``None`` as they would have before the snapshot was captured. + + Configuring an encryption key changes that contract. The operator asked for + fail-closed protection, so any failure to retrieve, decrypt, parse, or apply + the snapshot, including a snapshot that is simply missing, escapes this + dependency and fails the task, rather than running the tool without the + submitting caller's identity (#4747). """ try: parts = parse_task_key(key) @@ -295,6 +331,11 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None: from fastmcp.server.dependencies import get_server 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 branches below read + # `codec.protected` to pick between the fail-open and fail-closed contracts. + codec = snapshot_codec() + try: docket = get_server()._docket except RuntimeError: @@ -302,24 +343,53 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None: if docket is None: docket = _current_docket.get() if docket is None: + if codec.protected: + raise RuntimeError( + "No Docket backend is available to retrieve the protected " + "task snapshot, so the submitting caller cannot be recovered." + ) return task_scope = parts["task_scope"] task_id = parts["client_task_id"] try: async with docket.redis() as redis: - raw = await redis.get( - docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:snapshot") - ) + raw = await redis.get(_snapshot_redis_key(docket, task_scope, task_id)) if raw is None: - return - snapshot = TaskContextSnapshot.from_json(raw) + 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." + ) + 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 # the worker, exactly as a normal request would. _apply_snapshot_to_context(snapshot) + except SnapshotDecryptionError: + # Docket reports this to the client as a generic dependency-resolution + # failure, so name the cause here. A key mismatch across servers and + # workers is the likely reason and is not guessable from the wire error. + _logger.error( + "Failed to decrypt the task snapshot for %s. Every server and worker " + "on this queue must share the same FASTMCP_TASKS_ENCRYPTION_KEY. The " + "task will fail rather than run without the submitting caller's " + "identity.", + key, + ) + raise except Exception: + 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 " + "identity.", + key, + exc_info=True, + ) + raise _logger.warning("Failed to restore task snapshot for %s", key, exc_info=True) diff --git a/fastmcp_tasks/fastmcp_tasks/encryption.py b/fastmcp_tasks/fastmcp_tasks/encryption.py new file mode 100644 index 000000000..34b4fec12 --- /dev/null +++ b/fastmcp_tasks/fastmcp_tasks/encryption.py @@ -0,0 +1,171 @@ +"""Encryption of the task-context snapshot at rest. + +The snapshot a task carries holds the submitting caller's access token and every +inbound HTTP header, and it lives in the Docket backend for the task's TTL. A +distributed backend therefore keeps bearer credentials in Redis, where a +``rediss://`` URL protects the wire but not the stored value. + +Setting ``FASTMCP_TASKS_ENCRYPTION_KEY`` turns the stored snapshot into a +Fernet token. The same key must reach every server and worker on the queue, +because the process 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 + +from fastmcp.utilities.logging import get_logger +from fastmcp_tasks.settings import tasks_settings + +logger = get_logger(__name__) + +# Domain separation: FASTMCP_TASKS_ENCRYPTION_KEY may protect other task-owned +# state over time, and each use derives its own Fernet key from this material. +_SNAPSHOT_KEY_SALT = "fastmcp-task-snapshot-key" + +# Below this, warn: the keyspace is small enough that the offline attacker this +# feature defends against can search it even through PBKDF2. Matches the OAuth +# proxy's threshold for its signing-key material. +_SHORT_KEY_WARNING_LENGTH = 12 + +# Every Fernet token starts with the version byte 0x80, which base64url encodes +# (together with the leading zero bytes of its 64-bit timestamp) as "gAAAAA". +# A plaintext snapshot is a JSON object starting with "{", so the prefix cannot +# collide with a legitimately unencrypted value. +_FERNET_PREFIX = "gAAAAA" + + +class SnapshotDecryptionError(Exception): + """A stored snapshot is encrypted but cannot be read by this process. + + Raised for a wrong key, a tampered value, a plaintext value written before + the key was configured, or an encrypted value read by a process with no key + configured at all. The restore path lets this escape so the task fails, + rather than running the tool as an anonymous caller. + """ + + +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. + + It still refuses to decode a Fernet envelope: an encrypted snapshot + reaching a keyless process means the submitter configured a key this + process lacks (a partial rollout, or a lost setting), and passing the + ciphertext through would end in a swallowed parse error and an anonymous + run instead of the configured fail-closed behavior. + """ + + protected = False + + def encode(self, payload: str) -> str: + return payload + + def decode(self, stored: str | bytes) -> str: + text = stored.decode() if isinstance(stored, bytes) else stored + if text.startswith(_FERNET_PREFIX): + raise SnapshotDecryptionError( + "The stored task snapshot is encrypted, but this process has " + "no FASTMCP_TASKS_ENCRYPTION_KEY configured." + ) + return text + + +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 + comes from PBKDF2, never from HKDF. The stretch costs about a second, paid + once per process (see ``_codec_for``). + """ + + protected = True + + def __init__(self, material: str) -> None: + from cryptography.fernet import Fernet + + from fastmcp.server.auth.jwt_issuer import derive_jwt_key + + if not material: + raise ValueError( + "FASTMCP_TASKS_ENCRYPTION_KEY must not be empty. Unset it to store " + "task snapshots as plaintext, or set at least 32 random " + "characters." + ) + if len(material) < _SHORT_KEY_WARNING_LENGTH: + logger.warning( + "The configured encryption key is shorter than %d characters; " + "use at least 32 random characters.", + _SHORT_KEY_WARNING_LENGTH, + ) + key = derive_jwt_key(low_entropy_material=material, salt=_SNAPSHOT_KEY_SALT) + + self._fernet = Fernet(key=key) + + def encode(self, payload: str) -> str: + """Return the encrypted form of a serialized snapshot.""" + return self._fernet.encrypt(payload.encode()).decode() + + def decode(self, stored: str | bytes) -> str: + """Return the serialized snapshot a stored value holds. + + Raises ``SnapshotDecryptionError`` if the value was not produced by this + key, including when it is unencrypted. + """ + from cryptography.fernet import InvalidToken + + raw = stored.encode() if isinstance(stored, str) else stored + try: + return self._fernet.decrypt(raw).decode() + except InvalidToken as e: + raise SnapshotDecryptionError( + "The stored task snapshot could not be decrypted with the " + "configured FASTMCP_TASKS_ENCRYPTION_KEY." + ) from e + + +_PLAINTEXT_CODEC = PlaintextCodec() + + +@lru_cache(maxsize=4) +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 EncryptedCodec(material) + + +def snapshot_codec() -> SnapshotCodec: + """The codec for the configured key; the plaintext codec when none is set.""" + key = tasks_settings.encryption_key + if key is None: + return _PLAINTEXT_CODEC + return _codec_for(key.get_secret_value()) + + +def clear_codec_cache() -> None: + """Drop the cached codecs, so a changed key takes effect.""" + _codec_for.cache_clear() diff --git a/fastmcp_tasks/fastmcp_tasks/handlers.py b/fastmcp_tasks/fastmcp_tasks/handlers.py index 5193d5b84..ae9deb996 100644 --- a/fastmcp_tasks/fastmcp_tasks/handlers.py +++ b/fastmcp_tasks/fastmcp_tasks/handlers.py @@ -32,7 +32,7 @@ from fastmcp.exceptions import NotFoundError from fastmcp.tools.base import InputRequiredToolResult, Tool, ToolResult from fastmcp.utilities.tasks import DEFAULT_POLL_INTERVAL_MS from fastmcp.utilities.versions import VersionSpec -from fastmcp_tasks.context import get_task_scope +from fastmcp_tasks.context import get_task_scope, refresh_snapshot_ttl from fastmcp_tasks.creation import ( TASK_MAPPING_TTL_BUFFER_SECONDS, enqueue_task_leg, @@ -165,6 +165,10 @@ async def _lookup_task( await redis.expire(created_at_key, refresh_ttl) await redis.expire(poll_key, refresh_ttl) await refresh_current_leg_ttl(docket, task_scope, task_id, refresh_ttl) + # The snapshot must outlive the routing keys it serves: a re-entered leg + # restores the submitting caller from it, and with encryption configured a + # missing snapshot fails the task instead of degrading to an anonymous run. + await refresh_snapshot_ttl(docket, task_scope, task_id, refresh_ttl) created_at = created_at_bytes.decode("utf-8") if created_at_bytes else None diff --git a/fastmcp_tasks/fastmcp_tasks/settings.py b/fastmcp_tasks/fastmcp_tasks/settings.py index b836b0f67..8dcc96504 100644 --- a/fastmcp_tasks/fastmcp_tasks/settings.py +++ b/fastmcp_tasks/fastmcp_tasks/settings.py @@ -13,7 +13,7 @@ import os from datetime import timedelta from typing import Annotated -from pydantic import Field +from pydantic import Field, SecretStr from pydantic_settings import BaseSettings, SettingsConfigDict # Load the same dotenv source as core FastMCP settings, so a deployment that @@ -129,6 +129,38 @@ class DocketSettings(BaseSettings): docket_settings = DocketSettings() +class TasksSettings(BaseSettings): + """Settings for the task engine itself, as opposed to its Docket backend.""" + + model_config = SettingsConfigDict( + env_prefix="FASTMCP_TASKS_", + env_file=_ENV_FILE, + extra="ignore", + ) + + encryption_key: Annotated[ + SecretStr | None, + Field( + description=inspect.cleandoc( + """ + Key used to encrypt task context snapshots at rest. The snapshot + carries the submitting caller's access token and HTTP headers, + and it is written to the Docket backend for the task's TTL. + Every server and worker sharing a task queue must set the same + key; a worker that cannot decrypt a snapshot fails the task + rather than running it as an anonymous caller. When unset, the + snapshot is stored as plaintext JSON. The Fernet key is derived + from this value with PBKDF2, so any non-empty string works, but + use at least 32 random characters. + """ + ), + ), + ] = None + + +tasks_settings = TasksSettings() + + class TasksClientSettings(BaseSettings): """Client-side settings for driving background tasks. diff --git a/fastmcp_tasks/pyproject.toml b/fastmcp_tasks/pyproject.toml index a491aaaa8..6a1de018f 100644 --- a/fastmcp_tasks/pyproject.toml +++ b/fastmcp_tasks/pyproject.toml @@ -53,6 +53,9 @@ fallback-version = "0.0.0" [tool.hatch.metadata.hooks.uv-dynamic-versioning] dependencies = [ "fastmcp-slim[server]=={{ version }}", + # Fernet and the PBKDF2 key derivation behind FASTMCP_TASKS_ENCRYPTION_KEY, + # which encrypts task context snapshots at rest. + "cryptography>=43.0.0", "pydocket>=0.20.0", # burner-redis 0.1.7's Windows build crashes the interpreter (native fault, # no Python traceback) running the memory:// backend under pytest-xdist — diff --git a/tests/tasks/server/test_snapshot_encryption.py b/tests/tasks/server/test_snapshot_encryption.py new file mode 100644 index 000000000..0c35b6fe4 --- /dev/null +++ b/tests/tasks/server/test_snapshot_encryption.py @@ -0,0 +1,437 @@ +"""Tests for encryption of the task-context snapshot at rest (#4747). + +The snapshot carries the submitting caller's access token and every inbound HTTP +header, and it is written to the Docket backend for the task's TTL. With a +distributed backend those credentials sit in Redis where the backend's operators +can read them. Setting ``FASTMCP_TASKS_ENCRYPTION_KEY`` makes the snapshot a Fernet +token instead, and makes a worker that cannot decrypt one fail the task rather +than run it as an anonymous caller. +""" + +from __future__ import annotations + +import json +import logging +from collections.abc import Iterator +from unittest.mock import patch + +import pytest +from fastmcp_tasks.context import TaskContextSnapshot +from fastmcp_tasks.encryption import ( + EncryptedCodec, + PlaintextCodec, + SnapshotDecryptionError, + clear_codec_cache, + snapshot_codec, +) +from fastmcp_tasks.keys import task_redis_prefix +from fastmcp_tasks.settings import TasksSettings, tasks_settings +from pydantic import SecretStr + +from fastmcp import FastMCP +from fastmcp.server.dependencies import get_access_token +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + get_task, + make_access_token, + running_task_server, + submit_task, + wait_for_task, +) + +KEY = "a-test-encryption-key-for-snapshots" +OTHER_KEY = "a-different-test-encryption-key-entirely" + + +@pytest.fixture +def encryption_key() -> Iterator[str]: + """Configure the tasks encryption key for the duration of a test.""" + clear_codec_cache() + previous = tasks_settings.encryption_key + tasks_settings.encryption_key = SecretStr(KEY) + try: + yield KEY + finally: + tasks_settings.encryption_key = previous + clear_codec_cache() + + +@pytest.fixture +def no_encryption_key() -> Iterator[None]: + """Guarantee no key is configured, whatever the ambient environment holds.""" + clear_codec_cache() + previous = tasks_settings.encryption_key + tasks_settings.encryption_key = None + try: + yield + finally: + tasks_settings.encryption_key = previous + clear_codec_cache() + + +@pytest.fixture +def sensitive_snapshot() -> TaskContextSnapshot: + """A snapshot carrying a bearer token and an Authorization header.""" + token = make_access_token("client-a", "user-1") + return TaskContextSnapshot( + access_token_json=token.model_dump_json(), + http_headers={"authorization": f"Bearer {token.token}", "x-trace-id": "abc"}, + origin_request_id="req-1", + session_id="session-1", + owning_tool_name="peek", + owning_tool_version="1.0", + ) + + +class TestSnapshotCodec: + def test_round_trips_a_payload(self): + codec = EncryptedCodec(KEY) + assert codec.decode(codec.encode('{"a": 1}')) == '{"a": 1}' + + def test_encoded_payload_hides_the_credentials( + self, sensitive_snapshot: TaskContextSnapshot + ): + 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 = EncryptedCodec(OTHER_KEY).encode('{"a": 1}') + with pytest.raises(SnapshotDecryptionError): + 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): + 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"): + EncryptedCodec("") + + def test_decode_accepts_bytes(self): + """Redis hands back bytes on some backends.""" + 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_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}' + + def test_plaintext_codec_refuses_an_encrypted_payload(self): + """A keyless process must not pass ciphertext through as plaintext. + + Passing it through would end in a swallowed parse error and an + anonymous run, defeating the submitter's fail-closed configuration. + """ + encrypted = EncryptedCodec(KEY).encode('{"a": 1}') + with pytest.raises( + SnapshotDecryptionError, match="no FASTMCP_TASKS_ENCRYPTION_KEY" + ): + PlaintextCodec().decode(encrypted) + + +class TestTasksSettings: + def test_encryption_key_defaults_to_none(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("FASTMCP_TASKS_ENCRYPTION_KEY", raising=False) + + assert TasksSettings().encryption_key is None + + def test_encryption_key_env_var(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("FASTMCP_TASKS_ENCRYPTION_KEY", "s3kr1t-material") + + key = TasksSettings().encryption_key + assert key is not None + assert key.get_secret_value() == "s3kr1t-material" + + def test_encryption_key_is_not_printable(self, monkeypatch: pytest.MonkeyPatch): + """A settings dump must never carry the key into a log.""" + monkeypatch.setenv("FASTMCP_TASKS_ENCRYPTION_KEY", "s3kr1t-material") + + assert "s3kr1t-material" not in repr(TasksSettings()) + + +class TestSnapshotSerialization: + def test_json_round_trip_preserves_every_field( + self, sensitive_snapshot: TaskContextSnapshot + ): + assert ( + TaskContextSnapshot.from_json(sensitive_snapshot.to_json()) + == sensitive_snapshot + ) + + +async def _read_stored_snapshot(mcp: FastMCP, task_scope: str, task_id: str) -> str: + """Return the raw stored value of a task's snapshot key.""" + docket = mcp._docket + assert docket is not None + key = docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:snapshot") + async with docket.redis() as redis: + raw = await redis.get(key) + assert raw is not None + return raw.decode() if isinstance(raw, bytes) else str(raw) + + +async def _write_stored_snapshot( + mcp: FastMCP, task_scope: str, task_id: str, payload: str +) -> None: + """Overwrite a task's stored snapshot value.""" + docket = mcp._docket + assert docket is not None + key = docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:snapshot") + async with docket.redis() as redis: + await redis.set(key, payload) + + +async def _delete_stored_snapshot(mcp: FastMCP, task_scope: str, task_id: str) -> None: + """Remove a task's stored snapshot, as a TTL expiry would.""" + docket = mcp._docket + assert docket is not None + key = docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:snapshot") + async with docket.redis() as redis: + await redis.delete(key) + + +@pytest.fixture +def echo_token_server() -> FastMCP: + """A task server whose one tool reports the caller it restored.""" + mcp = FastMCP("snapshot-encryption-test") + mcp.add_extension(TasksExtension()) + + @mcp.tool(task=True) + async def whoami() -> str: + token = get_access_token() + return token.token if token else "no-token" + + return mcp + + +class TestEncryptedSnapshotRoundTrip: + async def test_worker_still_sees_the_submitting_caller( + self, echo_token_server: FastMCP, encryption_key: str + ): + token = make_access_token("client-a", "user-1") + + async with running_task_server(echo_token_server): + created = await submit_task( + echo_token_server, "whoami", {}, access_token=token + ) + final = await wait_for_task( + echo_token_server, created.task_id, access_token=token + ) + + assert final.status == "completed" + assert final.result is not None + assert final.result["structuredContent"] == {"result": token.token} + + async def test_stored_value_is_not_readable( + self, echo_token_server: FastMCP, encryption_key: str + ): + token = make_access_token("client-a", "user-1") + + async with running_task_server(echo_token_server): + created = await submit_task( + echo_token_server, "whoami", {}, access_token=token + ) + stored = await _read_stored_snapshot( + echo_token_server, "client-a|user-1", created.task_id + ) + await wait_for_task(echo_token_server, created.task_id, access_token=token) + + assert token.token not in stored + assert "authorization" not in stored + with pytest.raises(json.JSONDecodeError): + json.loads(stored) + + async def test_undecryptable_snapshot_fails_the_task( + self, + echo_token_server: FastMCP, + encryption_key: str, + caplog: pytest.LogCaptureFixture, + ): + """Fail closed: a worker that cannot recover the caller must not run. + + Running anyway would execute the tool as an anonymous caller, which for + an authorization-sensitive tool is worse than not running at all. Docket + surfaces this on the wire as a generic dependency failure, so the named + cause has to come from the log. + """ + token = make_access_token("client-a", "user-1") + 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): + created = await submit_task( + echo_token_server, "whoami", {}, access_token=token + ) + await _write_stored_snapshot( + echo_token_server, "client-a|user-1", created.task_id, tampered + ) + final = await wait_for_task( + echo_token_server, + created.task_id, + access_token=token, + target_states=frozenset({"failed"}), + ) + + assert final.status == "failed" + assert final.error is not None + assert "FASTMCP_TASKS_ENCRYPTION_KEY" in caplog.text + + async def test_missing_snapshot_fails_the_task( + self, echo_token_server: FastMCP, encryption_key: str + ): + """Fail closed extends to a snapshot that is gone, not just unreadable. + + A missing snapshot is reachable in production through TTL expiry, and + it loses the caller just as completely as a wrong key does. + """ + token = make_access_token("client-a", "user-1") + + async with running_task_server(echo_token_server): + created = await submit_task( + echo_token_server, "whoami", {}, access_token=token + ) + await _delete_stored_snapshot( + echo_token_server, "client-a|user-1", created.task_id + ) + final = await wait_for_task( + echo_token_server, + created.task_id, + access_token=token, + target_states=frozenset({"failed"}), + ) + + assert final.status == "failed" + + async def test_unparseable_snapshot_fails_the_task( + self, echo_token_server: FastMCP, encryption_key: str + ): + """Fail closed extends past decryption: a parse failure also loses the + caller, so it must not degrade to an anonymous run.""" + token = make_access_token("client-a", "user-1") + + def boom(*_args, **_kwargs): + raise RuntimeError("simulated deserialization failure") + + async with running_task_server(echo_token_server): + with patch.object(TaskContextSnapshot, "from_json", boom): + created = await submit_task( + echo_token_server, "whoami", {}, access_token=token + ) + final = await wait_for_task( + echo_token_server, + created.task_id, + access_token=token, + target_states=frozenset({"failed"}), + ) + + assert final.status == "failed" + + async def test_keyless_worker_fails_the_encrypted_task( + self, echo_token_server: FastMCP, encryption_key: str + ): + """A worker whose key was lost mid-rollout must not run anonymously. + + The submitter wrote an encrypted snapshot; the restoring process has no + key at all, so its plaintext codec would otherwise pass the ciphertext + through to a parse failure the fail-open path swallows. + """ + token = make_access_token("client-a", "user-1") + + async with running_task_server(echo_token_server): + created = await submit_task( + echo_token_server, "whoami", {}, access_token=token + ) + tasks_settings.encryption_key = None + clear_codec_cache() + final = await wait_for_task( + echo_token_server, + created.task_id, + access_token=token, + target_states=frozenset({"failed"}), + ) + + assert final.status == "failed" + + +class TestUnencryptedByDefault: + async def test_snapshot_stays_plaintext_without_a_key( + self, echo_token_server: FastMCP, no_encryption_key: None + ): + """No key configured is the pre-existing contract, unchanged.""" + token = make_access_token("client-a", "user-1") + + async with running_task_server(echo_token_server): + created = await submit_task( + echo_token_server, "whoami", {}, access_token=token + ) + stored = await _read_stored_snapshot( + echo_token_server, "client-a|user-1", created.task_id + ) + final = await wait_for_task( + echo_token_server, created.task_id, access_token=token + ) + + assert json.loads(stored)["access_token_json"] is not None + assert final.status == "completed" + + async def test_unreadable_snapshot_is_nonfatal_without_a_key( + self, echo_token_server: FastMCP, no_encryption_key: None + ): + """Without encryption a corrupt snapshot still only degrades the caller.""" + token = make_access_token("client-a", "user-1") + + async with running_task_server(echo_token_server): + created = await submit_task( + echo_token_server, "whoami", {}, access_token=token + ) + await _write_stored_snapshot( + echo_token_server, "client-a|user-1", created.task_id, "not json" + ) + final = await wait_for_task( + echo_token_server, created.task_id, access_token=token + ) + + assert final.status == "completed" + assert final.result is not None + assert final.result["structuredContent"] == {"result": "no-token"} + + +class TestTaskStillResolvesAfterFailure: + async def test_failed_task_reports_an_error( + self, echo_token_server: FastMCP, encryption_key: str + ): + """A fail-closed task is still a well-formed `tasks/get` result.""" + token = make_access_token("client-a", "user-1") + tampered = EncryptedCodec(OTHER_KEY).encode(TaskContextSnapshot().to_json()) + + async with running_task_server(echo_token_server): + created = await submit_task( + echo_token_server, "whoami", {}, access_token=token + ) + await _write_stored_snapshot( + echo_token_server, "client-a|user-1", created.task_id, tampered + ) + await wait_for_task( + echo_token_server, + created.task_id, + access_token=token, + target_states=frozenset({"failed"}), + ) + fetched = await get_task( + echo_token_server, created.task_id, access_token=token + ) + + assert fetched.status == "failed" diff --git a/tests/tasks/server/test_task_ttl.py b/tests/tasks/server/test_task_ttl.py index 8fa9f670e..edffda25a 100644 --- a/tests/tasks/server/test_task_ttl.py +++ b/tests/tasks/server/test_task_ttl.py @@ -96,3 +96,30 @@ async def test_poll_refreshes_routing_key_ttl(): async with docket.redis() as redis: # Refreshed well past the shrunk 5s, back toward the full window. assert await redis.ttl(key) > 60 + + +async def test_poll_refreshes_snapshot_ttl(): + """A poll extends the context snapshot's TTL alongside the routing keys. + + A re-entered leg restores the submitting caller from the snapshot, so an + actively polled task must never outlive it: without encryption an expired + snapshot degrades the leg to an anonymous run, and with encryption it fails + the task. After shrinking the snapshot's TTL, a `tasks/get` restores it. + """ + from fastmcp_tasks.context import _snapshot_redis_key + + mcp = _ttl_server() + async with running_task_server(mcp): + created = await submit_task(mcp, "slow_task", {}) + docket = mcp._docket + assert docket is not None + key = _snapshot_redis_key(docket, None, created.task_id) + + async with docket.redis() as redis: + await redis.expire(key, 5) + assert await redis.ttl(key) <= 5 + + await get_task(mcp, created.task_id) + + async with docket.redis() as redis: + assert await redis.ttl(key) > 60 diff --git a/uv.lock b/uv.lock index 58fe0e134..a3d4ce109 100644 --- a/uv.lock +++ b/uv.lock @@ -1089,6 +1089,7 @@ name = "fastmcp-tasks" source = { editable = "fastmcp_tasks" } dependencies = [ { name = "burner-redis", marker = "sys_platform == 'win32'" }, + { name = "cryptography" }, { name = "fastmcp-slim", extra = ["server"] }, { name = "pydocket" }, ] @@ -1096,6 +1097,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "burner-redis", marker = "sys_platform == 'win32'", specifier = "<0.1.7" }, + { name = "cryptography", specifier = ">=43.0.0" }, { name = "fastmcp-slim", extras = ["server"], editable = "fastmcp_slim" }, { name = "pydocket", specifier = ">=0.20.0" }, ]