use proper state loading condition

libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/state/agent_state.py

@@ -5,13 +5,16 @@
 
 from abc import abstractmethod
 from copy import deepcopy
+import logging
 from typing import Callable, Type
 
 from microsoft_agents.hosting.core.storage import Storage, StoreItem
 
 from .state_property_accessor import StatePropertyAccessor
 from ..turn_context import TurnContext
 
+logger = logging.getLogger(__name__)
+
 
 class CachedAgentState(StoreItem):
     """
@@ -47,6 +50,10 @@ def store_item_to_json(self) -> dict:
         }
         return serialized
 
+    def clear(self) -> None:
+        self.state = {}
+        self.hash = ""
+
     @staticmethod
     def from_json_to_store_item(json_data: dict) -> StoreItem:
         return CachedAgentState(json_data)
@@ -87,16 +94,23 @@ def __init__(self, storage: Storage, context_service_key: str):
         self._context_service_key = context_service_key
         self._cached_state: CachedAgentState | None = None
 
-    def get_cached_state(self, turn_context: TurnContext) -> CachedAgentState:
+    def get_cached_state(
+        self, turn_context: TurnContext | None = None
+    ) -> CachedAgentState:
         """
         Gets the cached agent state instance that wraps the raw cached data for this "AgentState"
         from the turn context.
 
-        :param turn_context: The context object for this turn.
+        :param turn_context: Deprecated. The context object for this turn. This parameter is no
+            longer required and will be removed in a future release.
         :type turn_context: :class:`microsoft_agents.hosting.core.turn_context.TurnContext`
         :return: The cached agent state instance.
         """
-        return turn_context.turn_state.get(self._context_service_key)
+        if turn_context:
+            logger.warning(
+                "AgentState.get_cached_state(): turn_context is deprecated and no longer required for get_cached_state(); it will be removed in a future release"
+            )
+        return self._cached_state
 
     def create_property(self, name: str) -> StatePropertyAccessor:
         """
@@ -113,10 +127,20 @@ def create_property(self, name: str) -> StatePropertyAccessor:
             )
         return BotStatePropertyAccessor(self, name)
 
-    def get(self, turn_context: TurnContext) -> dict[str, StoreItem]:
-        cached = self.get_cached_state(turn_context)
+    def get(self, turn_context: TurnContext | None = None) -> dict[str, StoreItem]:
+        """
+        Gets the raw cached state dictionary for this "AgentState".
 
-        return getattr(cached, "state", None)
+        :param turn_context: Deprecated. The context object for this turn. This parameter is no
+            longer required and will be removed in a future release.
+        :type turn_context: :class:`microsoft_agents.hosting.core.turn_context.TurnContext`
+        :return: The cached state dictionary.
+        """
+        if turn_context:
+            logger.warning(
+                "AgentState.get(): turn_context is deprecated and no longer required for get(); it will be removed in a future release"
+            )
+        return getattr(self._cached_state, "state", None)
 
     async def load(self, turn_context: TurnContext, force: bool = False) -> None:
         """
@@ -129,12 +153,27 @@ async def load(self, turn_context: TurnContext, force: bool = False) -> None:
         """
         storage_key = self.get_storage_key(turn_context)
 
-        if force or not self._cached_state:
+        if self._should_load(turn_context, force):
             items = await self._storage.read([storage_key], target_cls=CachedAgentState)
             val = items.get(storage_key, CachedAgentState())
             self._cached_state = val
             turn_context.turn_state[self._context_service_key] = val
 
+    def _should_load(self, turn_context: TurnContext, force: bool = False) -> bool:
+        """
+        Determines whether the state should be loaded from storage.
+
+        :param turn_context: The context object for this turn
+        :type turn_context: :class:`microsoft_agents.hosting.core.turn_context.TurnContext`
+        :param force: Optional, true to bypass the cache
+        :type force: bool
+        :return: True if the state should be loaded from storage, False otherwise
+        :rtype: bool
+        """
+
+        self._cached_state = turn_context.turn_state.get(self._context_service_key)
+        return force or self._cached_state is None
+
     async def save(self, turn_context: TurnContext, force: bool = False) -> None:
         """
         Saves the state cached in the current context for this turn.
@@ -146,17 +185,21 @@ async def save(self, turn_context: TurnContext, force: bool = False) -> None:
         :type force: bool
         """
 
-        if force or (self._cached_state is not None and self._cached_state.is_changed):
+        # Avoid self._cached_state updating during the save operation
+        cached_state = self._cached_state
+
+        if force or (cached_state is not None and cached_state.is_changed):
             storage_key = self.get_storage_key(turn_context)
-            changes: dict[str, StoreItem] = {storage_key: self._cached_state}
+            changes: dict[str, StoreItem] = {storage_key: cached_state}
             await self._storage.write(changes)
-            self._cached_state.hash = self._cached_state.compute_hash()
+            cached_state.hash = cached_state.compute_hash()
 
-    def clear(self, turn_context: TurnContext):
+    def clear(self, turn_context: TurnContext | None = None) -> None:
         """
         Clears any state currently stored in this state scope.
 
-        :param turn_context: The context object for this turn
+        :param turn_context: Deprecated. The context object for this turn. This parameter is no
+            longer required and will be removed in a future release.
         :type turn_context: :class:`microsoft_agents.hosting.core.turn_context.TurnContext`
 
         :return: None
@@ -165,9 +208,14 @@ def clear(self, turn_context: TurnContext):
             This function must be called in order for the cleared state to be persisted to the underlying store.
         """
         #  Explicitly setting the hash will mean IsChanged is always true. And that will force a Save.
-        cache_value = CachedAgentState()
-        cache_value.hash = ""
-        self._cached_state = cache_value
+        if turn_context:
+            logger.warning(
+                "AgentState.clear(): turn_context is deprecated and no longer required for clear(); it will be removed in a future release"
+            )
+        if self._cached_state is None:
+            logger.warning("AgentState.clear(): No cached state to clear.")
+            return
+        self._cached_state.clear()
 
     async def delete(self, turn_context: TurnContext) -> None:
         """
@@ -246,7 +294,7 @@ def delete_value(self, property_name: str) -> None:
         """
         if not property_name:
             raise TypeError(
-                "AgentState.delete_property(): property_name cannot be None."
+                "AgentState.delete_value(): property_name cannot be None or empty."
             )
 
         if self._cached_state.state.get(property_name):
@@ -265,7 +313,7 @@ def set_value(self, property_name: str, value: StoreItem) -> None:
         """
         if not property_name:
             raise TypeError(
-                "AgentState.delete_property(): property_name cannot be None."
+                "AgentState.set_value(): property_name cannot be None or empty."
             )
         self._cached_state.state[property_name] = value
 

libraries/microsoft-agents-hosting-dialogs/microsoft_agents/hosting/dialogs/dialog_manager.py

@@ -83,28 +83,10 @@ async def on_turn(self, context: TurnContext) -> DialogManagerResult:
         # Register DialogManager with TurnState.
         context.turn_state[DialogManager.__name__] = self
 
-        # Resolve ConversationState
-        conversation_state_name = ConversationState.__name__
-        if self.conversation_state is None:
-            if conversation_state_name not in context.turn_state:
-                raise Exception(
-                    f"Unable to get an instance of {conversation_state_name} from turn_context. "
-                    f"Please ensure ConversationState is available in turn_state."
-                )
-            self.conversation_state = cast(
-                ConversationState, context.turn_state[conversation_state_name]
-            )
-        else:
-            context.turn_state[conversation_state_name] = self.conversation_state
-
-        # Resolve UserState (optional)
-        user_state_name = UserState.__name__
-        if self.user_state is None:
-            self.user_state = cast(
-                UserState | None, context.turn_state.get(user_state_name, None)
-            )
-        else:
-            context.turn_state[user_state_name] = self.user_state
+        if self.conversation_state is not None:
+            await self.conversation_state.load(context, False)
+            if self.user_state is not None:
+                await self.user_state.load(context, False)
 
         # Create property accessors
         last_access_property = self.conversation_state.create_property(self.last_access)

tests/hosting_core/state/test_agent_state.py

@@ -384,6 +384,7 @@ async def test_cached_state_hash_computation(self):
         # Make a change
         prop_accessor = self.user_state.create_property("test_prop")
         await prop_accessor.set(self.context, _MockTestDataItem("test_value"))
+        cached_state = self.user_state.get_cached_state(self.context)
 
         # State should now be changed
         assert cached_state.is_changed

tests/hosting_dialogs/memory/scopes/test_memory_scopes.py

@@ -213,7 +213,7 @@ async def test_conversation_memory_scope_should_return_conversation_state(self):
         # Create test context
         adapter = DialogTestAdapter()
         context = TurnContext(adapter, _begin_message)
-        context.turn_state["ConversationState"] = conversation_state
+        await conversation_state.load(context)
 
         dialog_context = await dialogs.create_context(context)
 
@@ -285,7 +285,7 @@ async def test_user_memory_scope_should_return_state_once_loaded(self):
 
         # Create a DialogState property, DialogSet and register the dialogs.
         conversation_state = ConversationState(storage)
-        context.turn_state["ConversationState"] = conversation_state
+        await conversation_state.load(context)
         dialog_state = conversation_state.create_property("dialogs")
         dialogs = DialogSet(dialog_state)
         dialog = _TestDialog("test", "test message")

tests/hosting_dialogs/test_dialog_extensions.py

@@ -121,6 +121,7 @@ async def capture_eoc(
             await DialogExtensions.run_dialog(
                 dialog, context, convo_state.create_property("DialogState")
             )
+            await convo_state.save(context)
 
         return DialogTestAdapter(logic)
 

tests/hosting_dialogs/test_dialog_manager.py

@@ -120,13 +120,17 @@ async def logic(context: TurnContext):
             if test_case != SkillFlowTestCase.root_bot_only:
                 # Create a skill ClaimsIdentity and put it in turn_state so isSkillClaim() returns True.
                 claims_identity = ClaimsIdentity({}, False)
-                claims_identity.claims["ver"] = (
-                    "2.0"  # AuthenticationConstants.VersionClaim
-                )
-                claims_identity.claims["aud"] = (
+                claims_identity.claims[
+                    "ver"
+                ] = "2.0"  # AuthenticationConstants.VersionClaim
+                claims_identity.claims[
+                    "aud"
+                ] = (
                     SimpleComponentDialog.skill_bot_id
                 )  # AuthenticationConstants.AudienceClaim
-                claims_identity.claims["azp"] = (
+                claims_identity.claims[
+                    "azp"
+                ] = (
                     SimpleComponentDialog.parent_bot_id
                 )  # AuthenticationConstants.AuthorizedParty
                 context._identity = claims_identity