OAuth continuation fix

libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/agent_application.py

@@ -4,6 +4,8 @@
 """
 
 from __future__ import annotations
+
+import asyncio
 import logging
 from contextlib import nullcontext
 from copy import copy
@@ -39,13 +41,13 @@
 
 from .state import TurnState
 from ..channel_service_adapter import ChannelServiceAdapter
-from .oauth import Authorization
+from .oauth.authorization import Authorization, _AuthInterceptResult
 from .typing_indicator import TypingIndicator
 from .telemetry import spans
 
 from ._type_defs import RouteHandler, RouteSelector
 from ._routes import _RouteList, _Route, RouteRank, _agentic_selector
-from .proactive import Proactive, ProactiveOptions
+from .proactive import Proactive
 
 logger = logging.getLogger(__name__)
 
@@ -827,23 +829,16 @@ async def _on_turn(self, context: TurnContext):
                         ActivityTypes.invoke,
                     ]:
 
-                        (
-                            auth_intercepts,
-                            continuation_activity,
-                        ) = await self._auth._on_turn_auth_intercept(
-                            context, turn_state
+                        auth_intercept_result = (
+                            await self._auth._on_turn_auth_intercept(
+                                context, turn_state
+                            )
                         )
-                        if auth_intercepts:
-                            if continuation_activity:
-                                new_context = copy(context)
-                                new_context.activity = continuation_activity
-                                logger.info(
-                                    "Resending continuation activity %s",
-                                    continuation_activity.text,
-                                )
-                                await self.on_turn(new_context)
-                                await turn_state.save(context)
-                            return
+                        if auth_intercept_result.should_skip_turn:
+                            await self._handle_turn_skip(
+                                context, turn_state, auth_intercept_result
+                            )
+                            return  # allows immediate sending of invoke response
 
                     logger.debug("Running before turn middleware")
                     if not await self._run_before_turn_middleware(context, turn_state):
@@ -1020,6 +1015,61 @@ async def _start_long_running_call(
 
         return await func(context)
 
+    async def _handle_turn_skip(
+        self, context: TurnContext, turn_state: StateT, result: _AuthInterceptResult
+    ):
+        """Handle the case where the turn should be skipped due to an auth intercept result."""
+
+        async def __replay_turn(replay_context: TurnContext):
+
+            for key, val in context.turn_state.items():
+                if replay_context.turn_state.get(key, None) is None:
+                    replay_context.turn_state[key] = val
+
+            await self._on_turn(replay_context)
+
+        async def __replay(act: Activity):
+
+            await self._adapter.continue_conversation_with_claims(
+                context.identity,
+                act,
+                __replay_turn,
+            )
+
+        def __log_task_result(task: asyncio.Task):
+            try:
+                task.result()
+            except Exception as e:
+                logger.error(
+                    f"Error occurred while replaying the turn.",
+                    exc_info=True,
+                )
+
+        if result.should_replay:
+
+            if not context.identity:
+                logger.error(
+                    "Cannot replay turn because context.identity is not set. Ensure that the adapter is configured to set the identity."
+                )
+                raise ApplicationError(
+                    "Cannot replay turn because context.identity is not set. Ensure that the adapter is configured to set the identity."
+                )
+
+            await turn_state.save(context)
+
+            task: asyncio.Task
+
+            activity = result.continuation_activity or context.activity.model_copy(
+                deep=True
+            )
+            logger.info(
+                "Replaying the turn with current or saved continuation activity"
+            )
+            task = asyncio.create_task(__replay(activity))
+            task.add_done_callback(__log_task_result)
+
+        return
+
     async def _on_error(self, context: TurnContext, err: ApplicationError) -> None:
         if self._error:
             logger.info(

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

@@ -230,6 +230,16 @@ async def _handle_flow_response(
                     ).model_dump(exclude_unset=True),
                 )
             )
+        elif (
+            flow_state.tag == _FlowStateTag.COMPLETE
+            and context.activity.type == ActivityTypes.invoke
+        ):
+            await context.send_activity(
+                Activity(
+                    type=ActivityTypes.invoke_response,
+                    value=InvokeResponse(status=200).model_dump(exclude_unset=True),
+                )
+            )
 
     async def _sign_in(
         self,

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

@@ -5,6 +5,7 @@
 
 import logging
 from typing import Optional, Callable, Awaitable, cast
+from dataclasses import dataclass
 
 from microsoft_agents.activity import Activity, Channels, SignInConstants, TokenResponse
 from microsoft_agents.activity.activity_types import ActivityTypes
@@ -33,6 +34,15 @@
 }
 
 
+@dataclass
+class _AuthInterceptResult:
+    """Dataclass representing the result of an authentication intercept."""
+
+    should_skip_turn: bool
+    should_replay: bool
+    continuation_activity: Activity | None = None
+
+
 class Authorization:
     """Class responsible for managing authorization flows."""
 
@@ -305,20 +315,17 @@ async def sign_out(
 
     async def _on_turn_auth_intercept(
         self, context: TurnContext, state: TurnState
-    ) -> tuple[bool, Optional[Activity]]:
+    ) -> _AuthInterceptResult:
         """Intercepts the turn to check for active authentication flows.
 
-        Returns true if the rest of the turn should be skipped because auth did not finish.
-        Returns false if the turn should continue processing as normal.
-        If auth completes and a new turn should be started, returns the continuation activity
-        from the cached _SignInState.
+        Returns an _AuthInterceptResult indicating whether the turn should be skipped, replayed, and if a continuation activity is present.
 
         :param context: The context object for the current turn.
         :type context: :class:`microsoft_agents.hosting.core.turn_context.TurnContext`
         :param state: The turn state for the current turn.
         :type state: :class:`microsoft_agents.hosting.core.app.state.turn_state.TurnState`
-        :return: A tuple indicating whether the turn should be skipped and the continuation activity if applicable.
-        :rtype: tuple[bool, Optional[:class:`microsoft_agents.activity.Activity`]]
+        :return: An _AuthInterceptResult indicating whether the turn should be skipped and if a continuation activity is present.
+        :rtype: :class:`microsoft_agents.hosting.core.app.oauth.authorization._AuthInterceptResult`
         """
         sign_in_state = await self._load_sign_in_state(context)
 
@@ -331,17 +338,33 @@ async def _on_turn_auth_intercept(
                 if sign_in_response.tag == _FlowStateTag.COMPLETE:
                     if not sign_in_state:
                         # flow just completed, no continuation activity
-                        return False, None
+                        is_invoke = context.activity.type == ActivityTypes.invoke
+                        return _AuthInterceptResult(
+                            should_skip_turn=is_invoke,
+                            should_replay=False,
+                            continuation_activity=None,
+                        )
                     assert sign_in_state.continuation_activity is not None
                     continuation_activity = (
                         sign_in_state.continuation_activity.model_copy()
                     )
                     # flow complete, start new turn with continuation activity
-                    return True, continuation_activity
+
+                    return _AuthInterceptResult(
+                        should_skip_turn=True,
+                        should_replay=True,
+                        continuation_activity=continuation_activity,
+                    )
                 # auth flow still in progress, the turn should be skipped
-                return True, None
+                return _AuthInterceptResult(
+                    should_skip_turn=True,
+                    should_replay=False,
+                    continuation_activity=None,
+                )
         # no active auth flow, continue processing
-        return False, None
+        return _AuthInterceptResult(
+            should_skip_turn=False, should_replay=False, continuation_activity=None
+        )
 
     async def get_token(
         self, context: TurnContext, auth_handler_id: Optional[str] = None

tests/hosting_core/app/_oauth/test_authorization.py

@@ -23,9 +23,7 @@
 
 from microsoft_agents.hosting.core import (
     AuthHandler,
-    Storage,
     MemoryStorage,
-    TurnContext,
 )
 
 from tests._common.storage.utils import StorageBaseline
@@ -646,12 +644,11 @@ async def test_on_turn_auth_intercept_no_intercept(
     ):
         await authorization._delete_sign_in_state(context)
 
-        intercepts, continuation_activity = await authorization._on_turn_auth_intercept(
-            context, None
-        )
+        res = await authorization._on_turn_auth_intercept(context, None)
 
-        assert not continuation_activity
-        assert not intercepts
+        assert not res.continuation_activity
+        assert not res.should_skip_turn
+        assert not res.should_replay
 
         final_state = await authorization._load_sign_in_state(context)
 
@@ -687,12 +684,11 @@ async def test_on_turn_auth_intercept_with_intercept_incomplete(
             context, copy_sign_in_state(initial_state)
         )
 
-        intercepts, continuation_activity = await authorization._on_turn_auth_intercept(
-            context, auth_handler_id
-        )
+        res = await authorization._on_turn_auth_intercept(context, auth_handler_id)
 
-        assert not continuation_activity
-        assert intercepts
+        assert not res.continuation_activity
+        assert res.should_skip_turn
+        assert not res.should_replay
 
         final_state = await authorization._load_sign_in_state(context)
         assert sign_in_state_eq(final_state, initial_state)
@@ -721,12 +717,11 @@ async def test_on_turn_auth_intercept_with_intercept_complete(
             context, copy_sign_in_state(initial_state)
         )
 
-        intercepts, continuation_activity = await authorization._on_turn_auth_intercept(
-            context, auth_handler_id
-        )
+        res = await authorization._on_turn_auth_intercept(context, auth_handler_id)
 
-        assert continuation_activity == old_activity
-        assert intercepts
+        assert res.continuation_activity == old_activity
+        assert res.should_skip_turn
+        assert not res.should_replay
 
         final_state = await authorization._load_sign_in_state(context)
         assert sign_in_state_eq(final_state, initial_state)

tests/hosting_dialogs/test_dialog_manager.py

@@ -120,17 +120,13 @@ 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