Adding Storage wrappers for internal use cases

libraries/microsoft-agents-activity/microsoft_agents/activity/__init__.py

@@ -101,7 +101,11 @@
 
 from .channel_adapter_protocol import ChannelAdapterProtocol
 from .turn_context_protocol import TurnContextProtocol
-from ._load_configuration import load_configuration_from_env
+from ._utils import (
+    load_configuration_from_env,
+    _raise_if_falsey,
+    _raise_if_none
+)
 
 __all__ = [
     "AgentsModel",
@@ -194,6 +198,8 @@
     "ConversationUpdateTypes",
     "MessageUpdateTypes",
     "load_configuration_from_env",
+    "_raise_if_falsey",
+    "_raise_if_none",
     "ChannelAdapterProtocol",
     "TurnContextProtocol",
     "TokenOrSignInResourceResponse",

libraries/microsoft-agents-activity/microsoft_agents/activity/_utils/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+from ._error_handling import (
+    _raise_if_falsey,
+    _raise_if_none,
+)
+from ._load_configuration import load_configuration_from_env
+
+__all__ = [
+    "_raise_if_falsey",
+    "_raise_if_none",
+    "load_configuration_from_env",
+]

libraries/microsoft-agents-activity/microsoft_agents/activity/_utils/_error_handling.py

@@ -0,0 +1,30 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+def _raise_if_none(label: str, **kwargs) -> None:
+    """Raises an exception if any of the provided keyword arguments are None.
+
+    :param label: A label to include in the exception message.
+    :param kwargs: The keyword arguments to check.
+    :raises ValueError: If any of the provided keyword arguments are None.
+    """
+
+    none_args = [name for name, value in kwargs.items() if value is None]
+    if none_args:
+        raise ValueError(
+            f"{label}: The following arguments must be set and non-None: {', '.join(none_args)}"
+        )
+
+def _raise_if_falsey(label: str, **kwargs) -> None:
+    """Raises an exception if any of the provided keyword arguments are falsey.
+
+    :param label: A label to include in the exception message.
+    :param kwargs: The keyword arguments to check.
+    :raises ValueError: If any of the provided keyword arguments are falsey.
+    """
+
+    falsey_args = [name for name, value in kwargs.items() if not value]
+    if falsey_args:
+        raise ValueError(
+            f"{label}: The following arguments must be set and non-falsey (cannot be None or an empty string, for example): {', '.join(falsey_args)}"
+        )

libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/_oauth/_flow_storage_client.py

@@ -3,6 +3,8 @@
 
 from typing import Optional
 
+from microsoft_agents.activity import _raise_if_falsey
+
 from ..storage import Storage
 from ._flow_state import _FlowState
 
@@ -19,10 +21,16 @@ async def delete(self, keys: list[str]) -> None:
         pass
 
 
+# class FlowStorageClient(StorageNamespace):
+#     def __init__(self, channel_id: str, user_id: str, storage: Storage):
+#         super().__init__(_FlowStorageClient(channel_id, user_id, storage))
+
+
 # this could be generalized. Ideas:
 # - CachedStorage class for two-tier storage
 # - Namespaced/PrefixedStorage class for namespacing keying
 # not generally thread or async safe (operations are not atomic)
+
 class _FlowStorageClient:
     """Wrapper around Storage that manages sign-in state specific to each user and channel.
 
@@ -33,28 +41,18 @@ def __init__(
         self,
         channel_id: str,
         user_id: str,
-        storage: Storage,
-        cache_class: Optional[type[Storage]] = None,
-    ):
-        """
-        Args:
-            channel_id: used to create the prefix
-            user_id: used to create the prefix
-            storage: the backing storage
-            cache_class: the cache class to use (defaults to DummyCache, which performs no caching).
-                This cache's lifetime is tied to the FlowStorageClient instance.
-        """
+        storage: Storage
+    ) -> None:
+        """Initializes the _FlowStorageClient.
 
-        if not user_id or not channel_id:
-            raise ValueError(
-                "FlowStorageClient.__init__(): channel_id and user_id must be set."
-            )
+        :param channel_id: The ID of the channel.
+        :param user_id: The ID of the user.
+        :param storage: The backing storage.
+        """
+        _raise_if_falsey("_FlowStorageClient.__init__", channel_id=channel_id, user_id=user_id, storage=storage)
 
         self._base_key = f"auth/{channel_id}/{user_id}/"
         self._storage = storage
-        if cache_class is None:
-            cache_class = _DummyCache
-        self._cache = cache_class()
 
     @property
     def base_key(self) -> str:

libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/oauth/_handlers/_user_authorization.py

@@ -27,6 +27,11 @@
     _FlowStorageClient,
     _FlowStateTag,
 )
+from microsoft_agents.hosting.core.storage import (
+    _ItemNamespace,
+    Namespaces
+)
+
 from .._sign_in_response import _SignInResponse
 from ._authorization_handler import _AuthorizationHandler
 
@@ -41,7 +46,7 @@ class _UserAuthorization(_AuthorizationHandler):
 
     async def _load_flow(
         self, context: TurnContext
-    ) -> tuple[_OAuthFlow, _FlowStorageClient]:
+    ) -> tuple[_OAuthFlow, _ItemNamespace[_FlowState]]:
         """Loads the OAuth flow for a specific auth handler.
 
         A new flow is created in Storage if none exists for the channel, user, and handler
@@ -72,9 +77,13 @@ async def _load_flow(
         ms_app_id = context.turn_state.get(context.adapter.AGENT_IDENTITY_KEY).claims[
             "aud"
         ]
+
+        namespace = Namespaces.USER_AUTHORIZATION.format(
+            channel_id=channel_id,
+            user_id=user_id,
+        )
+        flow_storage_client = _ItemNamespace(namespace, self._storage, _FlowState)
 
-        # try to load existing state
-        flow_storage_client = _FlowStorageClient(channel_id, user_id, self._storage)
         logger.info("Loading OAuth flow state from storage")
         flow_state: _FlowState = await flow_storage_client.read(self._id)
         if not flow_state:
@@ -86,7 +95,6 @@ async def _load_flow(
                 connection=self._handler.abs_oauth_connection_name,
                 ms_app_id=ms_app_id,
             )
-            # await flow_storage_client.write(flow_state)
 
         flow = _OAuthFlow(flow_state, user_token_client)
         return flow, flow_storage_client

libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/oauth/authorization.py

@@ -11,7 +11,11 @@
 from microsoft_agents.activity import Activity, TokenResponse
 
 from ...turn_context import TurnContext
-from ...storage import Storage
+from ...storage import (
+    Storage,
+    _ItemNamespace,
+    _Namespaces
+)
 from ...authorization import Connections
 from ..._oauth import _FlowStateTag
 from ..state import TurnState
@@ -67,6 +71,11 @@ def __init__(
 
         self._storage = storage
         self._connection_manager = connection_manager
+        self._sign_in_state_store = _ItemNamespace(
+            _Namespaces.AUTHORIZATION,
+            self._storage,
+            _SignInState,
+        )
 
         self._sign_in_success_handler: Optional[
             Callable[[TurnContext, TurnState, Optional[str]], Awaitable[None]]
@@ -117,20 +126,10 @@ def _init_handlers(self) -> None:
                 connection_manager=self._connection_manager,
                 auth_handler=auth_handler,
             )
-
+    
     @staticmethod
-    def _sign_in_state_key(context: TurnContext) -> str:
-        """Generate a unique storage key for the sign-in state based on the context.
-
-        This is the key used to store and retrieve the sign-in state from storage, and
-        can be used to inspect or manipulate the state directly if needed.
-
-        :param context: The turn context for the current turn of conversation.
-        :type context: :class:`microsoft_agents.hosting.core.turn_context.TurnContext`
-        :return: A unique (across other values of channel_id and user_id) key for the sign-in state.
-        :rtype: str
-        """
-        return f"auth:_SignInState:{context.activity.channel_id}:{context.activity.from_property.id}"
+    def _sign_in_state_vkey(context: TurnContext) -> str:
+        return f"{context.activity.channel_id}:{context.activity.from_property.id}"
 
     async def _load_sign_in_state(self, context: TurnContext) -> Optional[_SignInState]:
         """Load the sign-in state from storage for the given context.
@@ -228,7 +227,7 @@ async def _start_or_continue_sign_in(
         auth_handler_id = auth_handler_id or self._default_handler_id
 
         # check cached sign in state
-        sign_in_state = await self._load_sign_in_state(context)
+        sign_in_state = await self._sign_in_state_store.read(self._sign_in_state_vkey(context))
         if not sign_in_state:
             # no existing sign-in state, create a new one
             sign_in_state = _SignInState(active_handler_id=auth_handler_id)
@@ -243,6 +242,7 @@ async def _start_or_continue_sign_in(
         if sign_in_response.tag == _FlowStateTag.COMPLETE:
             if self._sign_in_success_handler:
                 await self._sign_in_success_handler(context, state, auth_handler_id)
+            await self._sign_in_state_store.delete(self._sign_in_state_vkey(context))
             await self._delete_sign_in_state(context)
             Authorization._cache_token(
                 context, auth_handler_id, sign_in_response.token_response

libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/storage/__init__.py

@@ -1,27 +1,33 @@
 from .store_item import StoreItem
-from .storage import Storage, AsyncStorageBase
+from .storage import Storage, _AsyncStorageBase
 from .memory_storage import MemoryStorage
-from .transcript_info import TranscriptInfo
-from .transcript_logger import (
+from ._transcript import (
+    TranscriptInfo,
     TranscriptLogger,
     ConsoleTranscriptLogger,
     TranscriptLoggerMiddleware,
     FileTranscriptLogger,
     PagedResult,
+    TranscriptStore,
+    FileTranscriptStore,
+    TranscriptMemoryStore,
+)
+from ._wrappers import (
+    _StorageNamespace
 )
-from .transcript_store import TranscriptStore
-from .transcript_file_store import FileTranscriptStore
 
 __all__ = [
+    "_StorageNamespace",
     "StoreItem",
     "Storage",
-    "AsyncStorageBase",
+    "_AsyncStorageBase",
     "MemoryStorage",
     "TranscriptInfo",
     "TranscriptLogger",
     "ConsoleTranscriptLogger",
     "TranscriptLoggerMiddleware",
     "TranscriptStore",
+    "TranscriptMemoryStore",
     "FileTranscriptLogger",
     "FileTranscriptStore",
     "PagedResult",

libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/storage/_namespaces.py

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+class _Namespaces:
+    """Storage key namespaces used by various components."""
+
+    USER_AUTHORIZATION = "auth/{channel_id}/{user_id}"
+    AUTHORIZATION = "auth/{channel_id}/{from_property_id}"

libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/storage/_transcript/__init__.py

@@ -0,0 +1,23 @@
+from .transcript_info import TranscriptInfo
+from .transcript_logger import (
+    TranscriptLogger,
+    ConsoleTranscriptLogger,
+    TranscriptLoggerMiddleware,
+    FileTranscriptLogger,
+    PagedResult,
+)
+from .transcript_store import TranscriptStore
+from .transcript_file_store import FileTranscriptStore
+from .transcript_memory_store import TranscriptMemoryStore
+
+__all__ = [
+    "TranscriptInfo",
+    "TranscriptLogger",
+    "ConsoleTranscriptLogger",
+    "TranscriptLoggerMiddleware",
+    "FileTranscriptLogger",
+    "PagedResult",
+    "TranscriptStore",
+    "TranscriptMemoryStore",
+    "FileTranscriptStore",
+]