Use internal serialization format instead of pickle

Author: tcdent

src/agentexec/core/queue.py

@@ -70,6 +70,7 @@ async def research(agent_id: UUID, context: ResearchContext):
 
 
 async def dequeue(
+    *,
     queue_name: str | None = None,
     timeout: int = 1,
 ) -> dict[str, Any] | None:

src/agentexec/core/results.py

@@ -1,28 +1,32 @@
-"""Task result storage and retrieval."""
-
 from __future__ import annotations
 
 import asyncio
 import time
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel
 
 from agentexec import state
 
 if TYPE_CHECKING:
     from agentexec.core.task import Task
 
 
-async def get_result(task: Task, timeout: float = 300) -> Any:
+DEFAULT_TIMEOUT: int = 300  # TODO improve this polling approach
+
+
+async def get_result(task: Task, timeout: int = DEFAULT_TIMEOUT) -> BaseModel:
     """Poll for a task result.
 
     Waits for a task to complete and returns its result.
+    Uses automatic type reconstruction from serialized class information.
 
     Args:
         task: The Task instance to wait for
         timeout: Maximum seconds to wait for result
 
     Returns:
-        The task's return value
+        Deserialized result as BaseModel instance
 
     Raises:
         TimeoutError: If result not available within timeout
@@ -38,22 +42,23 @@ async def get_result(task: Task, timeout: float = 300) -> Any:
     raise TimeoutError(f"Result for {task.agent_id} not available within {timeout}s")
 
 
-async def gather(*tasks: Task) -> tuple[Any, ...]:
+async def gather(*tasks: Task, timeout: int = DEFAULT_TIMEOUT) -> tuple[BaseModel, ...]:
     """Wait for multiple tasks and return their results.
 
     Similar to asyncio.gather, but for background tasks.
 
     Args:
         *tasks: Task instances to wait for
+        timeout: Maximum seconds to wait for each result
 
     Returns:
-        Tuple of results in the same order as input tasks
+        Tuple of deserialized results as BaseModel instances
 
     Example:
         brand = await ax.enqueue("brand_research", ctx)
         market = await ax.enqueue("market_research", ctx)
 
         brand_result, market_result = await ax.gather(brand, market)
     """
-    results = await asyncio.gather(*[get_result(task) for task in tasks])
+    results = await asyncio.gather(*[get_result(task, timeout) for task in tasks])
     return tuple(results)

src/agentexec/core/task.py

@@ -2,7 +2,7 @@
 
 import asyncio
 import inspect
-from typing import Any, Callable, Protocol, TypedDict, Unpack, get_type_hints
+from typing import Any, Coroutine, Protocol, TypedDict, Union, Unpack, get_type_hints
 from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, PrivateAttr, field_serializer
@@ -22,13 +22,18 @@ class TaskHandlerKwargs(TypedDict):
 
 
 class TaskHandler(Protocol):
-    """Protocol for task handler functions.
+    """Protocol for task handler functions (sync or async).
 
-    Handlers accept **kwargs matching HandlerKwargs structure.
-    Return value is ignored. Can be sync or async.
+    Handlers accept **kwargs matching TaskHandlerKwargs structure.
+    Must return a Pydantic BaseModel (or Coroutine that resolves to BaseModel).
     """
 
-    def __call__(self, **kwargs: Unpack[TaskHandlerKwargs]) -> None: ...
+    __name__: str  # All functions have __name__ attribute
+
+    def __call__(
+        self,
+        **kwargs: Unpack[TaskHandlerKwargs],
+    ) -> Union[BaseModel, Coroutine[Any, Any, BaseModel]]: ...
 
 
 class TaskDefinition:
@@ -51,10 +56,12 @@ async def research(agent_id: UUID, context: ResearchContext):
     """
 
     name: str
-    handler: Callable[..., Any]
+    handler: TaskHandler
     context_class: type[BaseModel]
+    # TODO we handle this with serialize/deserialize when writing the result so this can probably go away
+    result_class: type[BaseModel]
 
-    def __init__(self, name: str, handler: Callable[..., Any]):
+    def __init__(self, name: str, handler: TaskHandler):
         """Initialize task definition.
 
         Args:
@@ -63,12 +70,14 @@ def __init__(self, name: str, handler: Callable[..., Any]):
 
         Raises:
             TypeError: If handler doesn't have a typed 'context' parameter with BaseModel subclass
+            TypeError: If handler doesn't have a return type annotation with BaseModel subclass
         """
         self.name = name
         self.handler = handler
         self.context_class = self._infer_context_class(handler)
+        self.result_class = self._infer_result_class(handler)
 
-    def _infer_context_class(self, handler: Callable[..., Any]) -> type[BaseModel]:
+    def _infer_context_class(self, handler: TaskHandler) -> type[BaseModel]:
         """Infer context class from handler's type annotations.
 
         Looks for a 'context' parameter with a Pydantic BaseModel type hint.
@@ -98,6 +107,36 @@ def _infer_context_class(self, handler: Callable[..., Any]) -> type[BaseModel]:
 
         return context_type
 
+    def _infer_result_class(self, handler: TaskHandler) -> type[BaseModel]:
+        """Infer result class from handler's return type annotation.
+
+        Looks for a return annotation with a Pydantic BaseModel type hint.
+
+        Args:
+            handler: The task handler function
+
+        Returns:
+            Result class (BaseModel subclass)
+
+        Raises:
+            TypeError: If return annotation is missing or not a BaseModel subclass
+        """
+        hints = get_type_hints(handler)
+        if "return" not in hints:
+            raise TypeError(
+                f"Task handler '{handler.__name__}' must have a return type "
+                f"annotation with a BaseModel subclass"
+            )
+
+        return_type = hints["return"]
+        if not (inspect.isclass(return_type) and issubclass(return_type, BaseModel)):
+            raise TypeError(
+                f"Task handler '{handler.__name__}' return type must be a "
+                f"BaseModel subclass, got {return_type}"
+            )
+
+        return return_type
+
 
 class Task(BaseModel):
     """Represents a background task instance.
@@ -187,17 +226,16 @@ def create(cls, task_name: str, context: BaseModel) -> Task:
             agent_id=agent_id,
         )
 
-    async def execute(self) -> Any:
+    async def execute(self) -> BaseModel | None:
         """Execute the task using its bound definition's handler.
 
         Manages task lifecycle: marks started, runs handler, marks completed/errored.
 
         Returns:
-            Handler return value
+            Handler return value, or None if handler raised an exception
 
         Raises:
             RuntimeError: If task has not been bound to a definition
-            Exception: Re-raises any exception from the handler after marking errored
         """
         if self._definition is None:
             raise RuntimeError("Task must be bound to a definition before execution")
@@ -214,12 +252,12 @@ async def execute(self) -> Any:
                 "context": self.context,
             }
 
+            result: BaseModel
             if asyncio.iscoroutinefunction(self._definition.handler):
                 result = await self._definition.handler(**kwargs)
             else:
-                result = self._definition.handler(**kwargs)
+                result = self._definition.handler(**kwargs)  # type: ignore[assignment]
 
-            # Store result for pipeline coordination
             await state.aset_result(
                 self.agent_id,
                 result,
@@ -239,3 +277,4 @@ async def execute(self) -> Any:
                 message=CONF.activity_message_error.format(error=e),
                 status=activity.Status.ERROR,
             )
+            return None

src/agentexec/state/__init__.py

@@ -1,9 +1,12 @@
 # cspell:ignore acheck
 
-from typing import AsyncGenerator, Coroutine
+from typing import AsyncGenerator, Coroutine, cast
 from uuid import UUID
 
+from pydantic import BaseModel
+
 from agentexec.config import CONF
+from agentexec.state.backend import StateBackend
 
 KEY_RESULT = (CONF.key_prefix, "result")
 KEY_EVENT = (CONF.key_prefix, "event")
@@ -12,10 +15,11 @@
 
 match CONF.state_backend:
     case "redis":
-        from agentexec.state import redis_backend as backend
+        from agentexec.state import redis_backend as _backend
     case _:
         raise RuntimeError(f"Unsupported state backend: {CONF.state_backend}.")
 
+backend: StateBackend = cast(StateBackend, _backend)
 
 __all__ = [
     "backend",
@@ -34,42 +38,50 @@
 ]
 
 
-def get_result(agent_id: UUID | str) -> object | None:
+def get_result(agent_id: UUID | str) -> BaseModel | None:
     """Get result for an agent (sync).
 
+    Returns deserialized BaseModel instance with automatic type reconstruction.
+
     Args:
         agent_id: Unique agent identifier (UUID or string)
 
     Returns:
-        Deserialized result object or None if not found
+        Deserialized BaseModel or None if not found
     """
     data = backend.get(backend.format_key(*KEY_RESULT, str(agent_id)))
     return backend.deserialize(data) if data else None
 
 
-def aget_result(agent_id: UUID | str) -> Coroutine[None, None, object | None]:
+def aget_result(agent_id: UUID | str) -> Coroutine[None, None, BaseModel | None]:
     """Get result for an agent (async).
 
+    Returns deserialized BaseModel instance with automatic type reconstruction.
+
     Args:
         agent_id: Unique agent identifier (UUID or string)
 
     Returns:
-        Coroutine that resolves to deserialized result object or None if not found
+        Coroutine that resolves to deserialized BaseModel or None if not found
     """
 
-    async def _get() -> object | None:
+    async def _get() -> BaseModel | None:
         data = await backend.aget(backend.format_key(*KEY_RESULT, str(agent_id)))
         return backend.deserialize(data) if data else None
 
     return _get()
 
 
-def set_result(agent_id: UUID | str, data: object, ttl_seconds: int | None = None) -> bool:
+def set_result(
+    agent_id: UUID | str,
+    data: BaseModel,
+    ttl_seconds: int | None = None,
+) -> bool:
     """Set result for an agent (sync).
 
     Args:
         agent_id: Unique agent identifier (UUID or string)
-        data: Result data to store
+        data: Result data (must be Pydantic BaseModel)
         ttl_seconds: Optional time-to-live in seconds
 
     Returns:
@@ -84,14 +96,14 @@ def set_result(agent_id: UUID | str, data: object, ttl_seconds: int | None = Non
 
 def aset_result(
     agent_id: UUID | str,
-    data: object,
+    data: BaseModel,
     ttl_seconds: int | None = None,
 ) -> Coroutine[None, None, bool]:
     """Set result for an agent (async).
 
     Args:
         agent_id: Unique agent identifier (UUID or string)
-        data: Result data to store
+        data: Result data (must be Pydantic BaseModel)
         ttl_seconds: Optional time-to-live in seconds
 
     Returns:

src/agentexec/state/backend.py

@@ -2,6 +2,8 @@
 
 from typing import AsyncGenerator, Coroutine, Optional, Protocol
 
+from pydantic import BaseModel
+
 
 class StateBackend(Protocol):
     """Protocol defining the state backend interface.
@@ -147,7 +149,7 @@ def publish(self, channel: str, message: str) -> None:
         """
         ...
 
-    async def subscribe(self, channel: str) -> AsyncGenerator[str, None]:
+    def subscribe(self, channel: str) -> AsyncGenerator[str, None]:
         """Subscribe to a channel and yield messages.
 
         Args:
@@ -171,35 +173,38 @@ def format_key(self, *args: str) -> str:
         ...
 
     # Serialization
-    def serialize(self, obj: object) -> bytes:
-        """Serialize a Python object to bytes.
+    def serialize(self, obj: BaseModel) -> bytes:
+        """Serialize a Pydantic BaseModel to bytes.
 
-        The serialization method is backend-specific. For example, Redis might
-        use pickle, while other backends might use JSON or msgpack.
+        Stores the fully qualified class name alongside the data to enable
+        automatic type reconstruction during deserialization.
 
         Args:
-            obj: Python object to serialize
+            obj: Pydantic BaseModel instance to serialize
 
         Returns:
             Serialized bytes
 
         Raises:
-            Exception: If serialization fails
+            TypeError: If obj is not a BaseModel instance
         """
         ...
 
-    def deserialize(self, data: bytes) -> object:
-        """Deserialize bytes back to a Python object.
+    def deserialize(self, data: bytes) -> BaseModel:
+        """Deserialize bytes back to a Pydantic BaseModel instance.
 
-        Must be compatible with the serialize() method for this backend.
+        Uses the stored class information to dynamically import and reconstruct
+        the original type.
 
         Args:
             data: Serialized bytes
 
         Returns:
-            Deserialized Python object
+            Deserialized BaseModel instance
 
         Raises:
-            Exception: If deserialization fails
+            ImportError: If the class module cannot be imported
+            AttributeError: If the class does not exist in the module
+            ValueError: If the data is invalid
         """
         ...