Add result persistence (#184)

## Summary

Adds automatic result persistence for task executions using
`py-key-value-aio` storage, plus critical bug fixes for resource cleanup
and improved test coverage.

## ✨ New Feature: Result Persistence

Tasks can now store return values and exceptions for later retrieval:

```python
async def calculate() -> int:
    return 42

execution = await docket.add(calculate)()
await worker.run_until_finished()

# Retrieve result (waits if task still running)
result = await execution.get_result()  # Returns 42
```

### Key Features

- **Automatic serialization**: Uses cloudpickle for any Python object
- **Exception storage**: Failed tasks store exceptions, which are
re-raised on `get_result()`
- **Smart skipping**: Tasks returning `None` skip persistence
- **TTL management**: Results expire with `execution_ttl` (default: 1
hour)
- **Timeout support**: `get_result(timeout=...)` with graceful timeout
handling
- **Pub/sub integration**: Waits for completion via existing state
subscription

### Implementation Details

- `Docket.result_storage`: `RedisStore` or `MemoryStore` backend
- `Execution.result_key`: Tracks where result is stored
- `Execution.get_result()`: Retrieves results, waiting via pub/sub if
needed
- Worker captures return values/exceptions after task execution
- Base64-encoded JSON for storage compatibility

### Storage Backend

Uses `py-key-value-aio` with pluggable storage:
- **Redis**: `RedisStore` for production (separate connection pool)
- **Memory**: `MemoryStore` for `memory://` URLs (testing)

## 📚 API Changes

### New Public Methods
```python
execution.get_result(timeout=...)  # New method
```

### Docket Configuration
```python
Docket(
    result_storage=custom_storage  # Optional: provide custom AsyncKeyValue
)
```

Closes #166

---

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude <[email protected]>

Author: desertaxle

.github/workflows/ci.yml

@@ -17,7 +17,7 @@ jobs:
         backend:
           - name: "Redis 6.2, redis-py <5"
             redis-version: "6.2"
-            redis-py-version: ">=4.6,<5"
+            redis-py-version: ">=5,<6"
           - name: "Redis 7.4, redis-py >=5"
             redis-version: "7.4"
             redis-py-version: ">=5"
@@ -27,13 +27,6 @@ jobs:
           - name: "Memory (in-memory backend)"
             redis-version: "memory"
             redis-py-version: ">=5"
-        exclude:
-          # Python 3.10 + Redis 6.2 + redis-py <5 combination is skipped
-          - python-version: "3.10"
-            backend:
-              name: "Redis 6.2, redis-py <5"
-              redis-version: "6.2"
-              redis-py-version: ">=4.6,<5"
         include:
           - python-version: "3.10"
             cov-threshold: 100

pyproject.toml

@@ -28,8 +28,9 @@ dependencies = [
     "opentelemetry-api>=1.30.0",
     "opentelemetry-exporter-prometheus>=0.51b0",
     "prometheus-client>=0.21.1",
+    "py-key-value-aio[memory,redis]>=0.2.8",
     "python-json-logger>=3.2.1",
-    "redis>=4.6",
+    "redis>=5",
     "rich>=13.9.4",
     "typer>=0.15.1",
     "typing_extensions>=4.12.0",

src/docket/cli.py

@@ -914,7 +914,7 @@ def create_display_layout() -> Layout:
                 )
 
                 # Add worker if available
-                if worker_name:
+                if worker_name:  # pragma: no branch
                     info_lines.append(f"[bold]Worker:[/bold] {worker_name}")
 
                 # Add error if failed

src/docket/docket.py

@@ -24,6 +24,7 @@
     overload,
 )
 
+from key_value.aio.stores.base import BaseContextManagerStore
 from typing_extensions import Self
 
 import redis.exceptions
@@ -41,6 +42,9 @@
     StrikeList,
     TaskFunction,
 )
+from key_value.aio.protocols.key_value import AsyncKeyValue
+from key_value.aio.stores.redis import RedisStore
+from key_value.aio.stores.memory import MemoryStore
 
 from .instrumentation import (
     REDIS_DISRUPTIONS,
@@ -147,6 +151,7 @@ def __init__(
         heartbeat_interval: timedelta = timedelta(seconds=2),
         missed_heartbeats: int = 5,
         execution_ttl: timedelta = timedelta(hours=1),
+        result_storage: AsyncKeyValue | None = None,
     ) -> None:
         """
         Args:
@@ -171,6 +176,14 @@ def __init__(
         self.execution_ttl = execution_ttl
         self._cancel_task_script = None
 
+        self.result_storage: AsyncKeyValue
+        if url.startswith("memory://"):
+            self.result_storage = MemoryStore()
+        else:
+            self.result_storage = RedisStore(
+                url=url, default_collection=f"{name}:results"
+            )
+
     @property
     def worker_group_name(self) -> str:
         return "docket-workers"
@@ -220,6 +233,10 @@ async def __aenter__(self) -> Self:
             if "BUSYGROUP" not in repr(e):
                 raise
 
+        if isinstance(self.result_storage, BaseContextManagerStore):
+            await self.result_storage.__aenter__()
+        else:
+            await self.result_storage.setup()
         return self
 
     async def __aexit__(
@@ -228,6 +245,9 @@ async def __aexit__(
         exc_value: BaseException | None,
         traceback: TracebackType | None,
     ) -> None:
+        if isinstance(self.result_storage, BaseContextManagerStore):
+            await self.result_storage.__aexit__(exc_type, exc_value, traceback)
+
         del self.tasks
         del self.strike_list
 

src/docket/execution.py

@@ -1,4 +1,6 @@
 import abc
+import asyncio
+import base64
 import enum
 import inspect
 import json
@@ -237,7 +239,7 @@ async def sync(self) -> None:
                 self.message = None
                 self.updated_at = None
 
-    async def _delete(self) -> None:
+    async def delete(self) -> None:
         """Delete the progress data from Redis.
 
         Called internally when task execution completes.
@@ -326,6 +328,7 @@ def __init__(
         self.started_at: datetime | None = None
         self.completed_at: datetime | None = None
         self.error: str | None = None
+        self.result_key: str | None = None
         self.progress: ExecutionProgress = ExecutionProgress(docket, key)
         self._redis_key = f"{docket.name}:runs:{key}"
 
@@ -632,37 +635,47 @@ async def claim(self, worker: str) -> None:
             }
         )
 
-    async def mark_as_completed(self) -> None:
+    async def mark_as_completed(self, result_key: str | None = None) -> None:
         """Mark task as completed successfully.
 
+        Args:
+            result_key: Optional key where the task result is stored
+
         Sets TTL on state data (from docket.execution_ttl) and deletes progress data.
         """
         completed_at = datetime.now(timezone.utc).isoformat()
         async with self.docket.redis() as redis:
+            mapping: dict[str, str] = {
+                "state": ExecutionState.COMPLETED.value,
+                "completed_at": completed_at,
+            }
+            if result_key is not None:
+                mapping["result_key"] = result_key
             await redis.hset(
                 self._redis_key,
-                mapping={
-                    "state": ExecutionState.COMPLETED.value,
-                    "completed_at": completed_at,
-                },
+                mapping=mapping,
             )
             # Set TTL from docket configuration
             await redis.expire(
                 self._redis_key, int(self.docket.execution_ttl.total_seconds())
             )
         self.state = ExecutionState.COMPLETED
+        self.result_key = result_key
         # Delete progress data
-        await self.progress._delete()
+        await self.progress.delete()
         # Publish state change event
         await self._publish_state(
             {"state": ExecutionState.COMPLETED.value, "completed_at": completed_at}
         )
 
-    async def mark_as_failed(self, error: str | None = None) -> None:
+    async def mark_as_failed(
+        self, error: str | None = None, result_key: str | None = None
+    ) -> None:
         """Mark task as failed.
 
         Args:
             error: Optional error message describing the failure
+            result_key: Optional key where the exception is stored
 
         Sets TTL on state data (from docket.execution_ttl) and deletes progress data.
         """
@@ -674,14 +687,17 @@ async def mark_as_failed(self, error: str | None = None) -> None:
             }
             if error:
                 mapping["error"] = error
+            if result_key is not None:
+                mapping["result_key"] = result_key
             await redis.hset(self._redis_key, mapping=mapping)
             # Set TTL from docket configuration
             await redis.expire(
                 self._redis_key, int(self.docket.execution_ttl.total_seconds())
             )
         self.state = ExecutionState.FAILED
+        self.result_key = result_key
         # Delete progress data
-        await self.progress._delete()
+        await self.progress.delete()
         # Publish state change event
         state_data = {
             "state": ExecutionState.FAILED.value,
@@ -691,6 +707,80 @@ async def mark_as_failed(self, error: str | None = None) -> None:
             state_data["error"] = error
         await self._publish_state(state_data)
 
+    async def get_result(self, *, timeout: datetime | None = None) -> Any:
+        """Retrieve the result of this task execution.
+
+        If the execution is not yet complete, this method will wait using
+        pub/sub for state updates until completion.
+
+        Args:
+            timeout: Optional absolute datetime when to stop waiting.
+                    If None, waits indefinitely.
+
+        Returns:
+            The result of the task execution, or None if the task returned None.
+
+        Raises:
+            Exception: If the task failed, raises the stored exception
+            TimeoutError: If timeout is reached before execution completes
+        """
+        # Wait for execution to complete if not already done
+        if self.state not in (ExecutionState.COMPLETED, ExecutionState.FAILED):
+            # Calculate timeout duration if absolute timeout provided
+            timeout_seconds = None
+            if timeout is not None:
+                timeout_seconds = (timeout - datetime.now(timezone.utc)).total_seconds()
+                if timeout_seconds <= 0:
+                    raise TimeoutError(
+                        f"Timeout waiting for execution {self.key} to complete"
+                    )
+
+            try:
+
+                async def wait_for_completion():
+                    async for event in self.subscribe():  # pragma: no branch
+                        if event["type"] == "state":
+                            state = ExecutionState(event["state"])
+                            if state in (
+                                ExecutionState.COMPLETED,
+                                ExecutionState.FAILED,
+                            ):
+                                # Sync to get latest data including result key
+                                await self.sync()
+                                break
+
+                # Use asyncio.wait_for to enforce timeout
+                await asyncio.wait_for(wait_for_completion(), timeout=timeout_seconds)
+            except asyncio.TimeoutError:
+                raise TimeoutError(
+                    f"Timeout waiting for execution {self.key} to complete"
+                )
+
+        # If failed, retrieve and raise the exception
+        if self.state == ExecutionState.FAILED:
+            if self.result_key:
+                # Retrieve serialized exception from result_storage
+                result_data = await self.docket.result_storage.get(self.result_key)
+                if result_data and "data" in result_data:
+                    # Base64-decode and unpickle
+                    pickled_exception = base64.b64decode(result_data["data"])
+                    exception = cloudpickle.loads(pickled_exception)  # type: ignore[arg-type]
+                    raise exception
+            # If no stored exception, raise a generic error with the error message
+            error_msg = self.error or "Task execution failed"
+            raise Exception(error_msg)
+
+        # If completed successfully, retrieve result if available
+        if self.result_key:
+            result_data = await self.docket.result_storage.get(self.result_key)
+            if result_data is not None and "data" in result_data:
+                # Base64-decode and unpickle
+                pickled_result = base64.b64decode(result_data["data"])
+                return cloudpickle.loads(pickled_result)  # type: ignore[arg-type]
+
+        # No result stored - task returned None
+        return None
+
     async def sync(self) -> None:
         """Synchronize instance attributes with current execution data from Redis.
 
@@ -720,13 +810,17 @@ async def sync(self) -> None:
                     else None
                 )
                 self.error = data[b"error"].decode() if b"error" in data else None
+                self.result_key = (
+                    data[b"result_key"].decode() if b"result_key" in data else None
+                )
             else:
                 # No data exists - reset to defaults
                 self.state = ExecutionState.SCHEDULED
                 self.worker = None
                 self.started_at = None
                 self.completed_at = None
                 self.error = None
+                self.result_key = None
 
         # Sync progress data
         await self.progress.sync()