fix: deepcopy AgentState

Author: Unshure

src/strands/agent/state.py

@@ -1,5 +1,6 @@
 """Agent state management."""
 
+import copy
 import json
 from typing import Any, Dict, Optional
 
@@ -13,12 +14,12 @@ class AgentState:
     - Get/set/delete operations
     """
 
-    def __init__(self, initial_state: Optional[Dict[str, Dict[str, Any]]] = None):
+    def __init__(self, initial_state: Optional[Dict[str, Any]] = None):
         """Initialize AgentState."""
         self._state: Dict[str, Dict[str, Any]]
         if initial_state:
             self._validate_json_serializable(initial_state)
-            self._state = initial_state.copy()
+            self._state = copy.deepcopy(initial_state)
         else:
             self._state = {}
 
@@ -35,7 +36,7 @@ def set(self, key: str, value: Any) -> None:
         self._validate_key(key)
         self._validate_json_serializable(value)
 
-        self._state[key] = value
+        self._state[key] = copy.deepcopy(value)
 
     def get(self, key: Optional[str] = None) -> Any:
         """Get a value or entire state.
@@ -47,10 +48,10 @@ def get(self, key: Optional[str] = None) -> Any:
             The stored value, entire state dict, or None if not found
         """
         if key is None:
-            return self._state.copy()
+            return copy.deepcopy(self._state)
         else:
             # Return specific key
-            return self._state.get(key)
+            return copy.deepcopy(self._state.get(key))
 
     def delete(self, key: str) -> None:
         """Delete a specific key from the state.

tests/strands/agent/test_agent.py

@@ -1,5 +1,6 @@
 import copy
 import importlib
+import json
 import os
 import textwrap
 import threading
@@ -1326,3 +1327,37 @@ def test_non_json_serializable_state_throws_error():
     with pytest.raises(ValueError, match="Value is not JSON serializable"):
         agent = Agent(state={"object": object()})
         print(agent.state)
+
+
+def test_agent_state_breaks_dict_reference():
+    ref_dict = {"hello": "world"}
+    agent = Agent(state=ref_dict)
+    ref_dict["hello"] = object()
+
+    json.dumps(agent.state.get())
+
+
+def test_agent_state_breaks_deep_dict_reference():
+    ref_dict = {"world": "!"}
+    init_dict = {"hello": ref_dict}
+    agent = Agent(state=init_dict)
+    ref_dict["world"] = object()
+
+    json.dumps(agent.state.get())
+
+
+def test_agent_state_set_breaks_dict_reference():
+    agent = Agent()
+    ref_dict = {"hello": "world"}
+    agent.state.set("hello", ref_dict)
+    ref_dict["hello"] = object()
+
+    json.dumps(agent.state.get())
+
+
+def test_agent_state_get_breaks_deep_dict_reference():
+    agent = Agent(state={"hello": {"world": "!"}})
+    ref_state = agent.state.get()
+    ref_state["hello"]["world"] = object()
+
+    json.dumps(agent.state.get())