feat(agent): Add pluggable serialization strategies for AgentState

src/strands/__init__.py

@@ -2,11 +2,17 @@
 
 from . import agent, models, telemetry, types
 from .agent.agent import Agent
+from .agent.serializers import JSONSerializer, PickleSerializer, StateSerializer
+from .agent.state import AgentState
 from .tools.decorator import tool
 from .types.tools import ToolContext
 
 __all__ = [
     "Agent",
+    "AgentState",
+    "JSONSerializer",
+    "PickleSerializer",
+    "StateSerializer",
     "agent",
     "models",
     "tool",

src/strands/agent/__init__.py

@@ -4,6 +4,7 @@
 
 - Agent: The main interface for interacting with AI models and tools
 - ConversationManager: Classes for managing conversation history and context windows
+- Serializers: Pluggable serialization strategies for agent state (JSONSerializer, PickleSerializer)
 """
 
 from .agent import Agent
@@ -14,12 +15,18 @@
     SlidingWindowConversationManager,
     SummarizingConversationManager,
 )
+from .serializers import JSONSerializer, PickleSerializer, StateSerializer
+from .state import AgentState
 
 __all__ = [
     "Agent",
     "AgentResult",
+    "AgentState",
     "ConversationManager",
+    "JSONSerializer",
     "NullConversationManager",
+    "PickleSerializer",
     "SlidingWindowConversationManager",
+    "StateSerializer",
     "SummarizingConversationManager",
 ]

src/strands/agent/serializers.py

@@ -0,0 +1,159 @@
+"""State serializers for agent state management.
+
+This module provides pluggable serialization strategies for AgentState:
+- JSONSerializer: Default serializer, backward compatible, validates on set()
+- PickleSerializer: Supports any picklable Python object, validates on set()
+- StateSerializer: Protocol for custom serializers
+"""
+
+import copy
+import json
+import pickle
+from typing import Any, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class StateSerializer(Protocol):
+    """Protocol for state serializers.
+
+    Custom serializers can implement this protocol to provide
+    alternative serialization strategies for agent state.
+    """
+
+    def serialize(self, data: dict[str, Any]) -> bytes:
+        """Serialize state dict to bytes.
+
+        Args:
+            data: Dictionary of state data to serialize
+
+        Returns:
+            Serialized state as bytes
+        """
+        ...
+
+    def deserialize(self, data: bytes) -> dict[str, Any]:
+        """Deserialize bytes back to state dict.
+
+        Args:
+            data: Serialized state bytes
+
+        Returns:
+            Deserialized state dictionary
+        """
+        ...
+
+    def validate(self, value: Any) -> None:
+        """Validate a value can be serialized.
+
+        Serializers that accept any value should implement this as a no-op.
+
+        Args:
+            value: The value to validate
+
+        Raises:
+            ValueError: If value cannot be serialized by this serializer
+        """
+        ...
+
+
+class JSONSerializer:
+    """JSON-based state serializer.
+
+    Default serializer that provides:
+    - Human-readable serialization format
+    - Validation on set() to maintain current behavior
+    - Backward compatibility with existing code
+    """
+
+    def serialize(self, data: dict[str, Any]) -> bytes:
+        """Serialize state dict to JSON bytes.
+
+        Args:
+            data: Dictionary of state data to serialize
+
+        Returns:
+            JSON serialized state as bytes
+        """
+        return json.dumps(data).encode("utf-8")
+
+    def deserialize(self, data: bytes) -> dict[str, Any]:
+        """Deserialize JSON bytes back to state dict.
+
+        Args:
+            data: JSON serialized state bytes
+
+        Returns:
+            Deserialized state dictionary
+        """
+        result: dict[str, Any] = json.loads(data.decode("utf-8"))
+        return result
+
+    def validate(self, value: Any) -> None:
+        """Validate that a value is JSON serializable.
+
+        Args:
+            value: The value to validate
+
+        Raises:
+            ValueError: If value is not JSON serializable
+        """
+        try:
+            json.dumps(value)
+        except (TypeError, ValueError) as e:
+            raise ValueError(
+                f"Value is not JSON serializable: {type(value).__name__}. "
+                f"Only JSON-compatible types (str, int, float, bool, list, dict, None) are allowed."
+            ) from e
+
+
+class PickleSerializer:
+    """Pickle-based state serializer.
+
+    Provides:
+    - Support for any picklable Python object (datetime, UUID, dataclass, Pydantic models, etc.)
+    - Validation on set() to catch unpicklable objects (DB connections, file handles, etc.)
+
+    Security Warning:
+        Pickle can execute arbitrary code during deserialization.
+        Only unpickle data from trusted sources.
+    """
+
+    def serialize(self, data: dict[str, Any]) -> bytes:
+        """Serialize state dict using pickle.
+
+        Args:
+            data: Dictionary of state data to serialize
+
+        Returns:
+            Pickle serialized state as bytes
+        """
+        return pickle.dumps(copy.deepcopy(data))
+
+    def deserialize(self, data: bytes) -> dict[str, Any]:
+        """Deserialize pickle bytes back to state dict.
+
+        Args:
+            data: Pickle serialized state bytes
+
+        Returns:
+            Deserialized state dictionary
+        """
+        result: dict[str, Any] = pickle.loads(data)  # noqa: S301
+        return result
+
+    def validate(self, value: Any) -> None:
+        """Validate that a value can be pickled.
+
+        Args:
+            value: The value to validate
+
+        Raises:
+            ValueError: If value cannot be pickled
+        """
+        try:
+            pickle.dumps(value)
+        except TypeError as e:
+            raise ValueError(
+                f"Value is not picklable: {type(value).__name__}. "
+                f"Objects like database connections, file handles, and sockets cannot be serialized."
+            ) from e

src/strands/agent/state.py

@@ -1,6 +1,147 @@
-"""Agent state management."""
+"""Agent state management.
 
-from ..types.json_dict import JSONSerializableDict
+Provides flexible state container with pluggable serialization.
+"""
 
-# Type alias for agent state
-AgentState = JSONSerializableDict
+import copy
+from typing import Any
+
+from .serializers import JSONSerializer, StateSerializer
+
+
+class AgentState:
+    """Flexible state container with pluggable serialization.
+
+    AgentState provides a key-value store for agent state with:
+    - Pluggable serialization (JSON by default, Pickle for rich types)
+    - Backward compatible API with existing code
+
+    Example:
+        Basic usage (backward compatible):
+        ```python
+        state = AgentState()
+        state.set("count", 42)
+        state.get("count")  # Returns 42
+        ```
+
+        Rich types with PickleSerializer:
+        ```python
+        from strands.agent.serializers import PickleSerializer
+        from datetime import datetime
+
+        state = AgentState(serializer=PickleSerializer())
+        state.set("created_at", datetime.now())  # Works with Pickle
+        ```
+    """
+
+    def __init__(
+        self,
+        initial_state: dict[str, Any] | None = None,
+        serializer: StateSerializer | None = None,
+    ):
+        """Initialize AgentState.
+
+        Args:
+            initial_state: Optional initial state dictionary
+            serializer: Serializer to use for state persistence.
+                Defaults to JSONSerializer for backward compatibility.
+
+        Raises:
+            ValueError: If initial_state contains non-serializable values (with JSONSerializer)
+        """
+        self._serializer = serializer if serializer is not None else JSONSerializer()
+        self._data: dict[str, Any]
+
+        if initial_state:
+            # Validate initial state
+            self._serializer.validate(initial_state)
+            self._data = copy.deepcopy(initial_state)
+        else:
+            self._data = {}
+
+    @property
+    def serializer(self) -> StateSerializer:
+        """Get the current serializer.
+
+        Returns:
+            The serializer used for state persistence
+        """
+        return self._serializer
+
+    @serializer.setter
+    def serializer(self, value: StateSerializer) -> None:
+        """Set the serializer.
+
+        Args:
+            value: New serializer to use for state persistence
+        """
+        self._serializer = value
+
+    def set(self, key: str, value: Any) -> None:
+        """Set a value in the store.
+
+        Args:
+            key: The key to store the value under
+            value: The value to store
+
+        Raises:
+            ValueError: If key is invalid, or if value is not serializable
+        """
+        self._validate_key(key)
+        self._serializer.validate(value)
+        self._data[key] = copy.deepcopy(value)
+
+    def get(self, key: str | None = None) -> Any:
+        """Get a value or entire data.
+
+        Args:
+            key: The key to retrieve (if None, returns entire data dict)
+
+        Returns:
+            The stored value, entire data dict, or None if not found
+        """
+        if key is None:
+            return copy.deepcopy(self._data)
+        else:
+            return copy.deepcopy(self._data.get(key))
+
+    def delete(self, key: str) -> None:
+        """Delete a specific key from the store.
+
+        Args:
+            key: The key to delete
+        """
+        self._validate_key(key)
+        self._data.pop(key, None)
+
+    def serialize(self) -> bytes:
+        """Serialize state.
+
+        Returns:
+            Serialized state as bytes
+        """
+        return self._serializer.serialize(self._data)
+
+    def deserialize(self, data: bytes) -> None:
+        """Deserialize state.
+
+        Args:
+            data: Serialized state bytes to restore
+        """
+        self._data = self._serializer.deserialize(data)
+
+    def _validate_key(self, key: str) -> None:
+        """Validate that a key is valid.
+
+        Args:
+            key: The key to validate
+
+        Raises:
+            ValueError: If key is invalid
+        """
+        if key is None:
+            raise ValueError("Key cannot be None")
+        if not isinstance(key, str):
+            raise ValueError("Key must be a string")
+        if not key.strip():
+            raise ValueError("Key cannot be empty")