Detach OAuth continuation replay from token-exchange invoke

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

@@ -4,9 +4,10 @@
 """
 
 from __future__ import annotations
+import asyncio
 import logging
 from contextlib import nullcontext
-from copy import copy
+from copy import copy, deepcopy
 from functools import partial
 
 import re
@@ -101,6 +102,7 @@ def __init__(
         self._route_list = _RouteList[StateT]()
         self._internal_before_turn = []
         self._internal_after_turn = []
+        self._background_tasks: set[asyncio.Task[None]] = set()
 
         configuration = kwargs
 
@@ -835,14 +837,14 @@ async def _on_turn(self, context: TurnContext):
                         )
                         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)
+                                self._start_auth_continuation(
+                                    context, continuation_activity
+                                )
                             return
 
                     logger.debug("Running before turn middleware")
@@ -866,6 +868,46 @@ async def _on_turn(self, context: TurnContext):
             )
             await self._on_error(context, err)
 
+    def _start_auth_continuation(
+        self, context: TurnContext, continuation_activity: Activity
+    ) -> None:
+        """Replay a completed OAuth flow without holding its invoke request open."""
+        if context.identity is None:
+            raise ApplicationError(
+                "Cannot replay an OAuth continuation without a claims identity."
+            )
+
+        replay_activity = deepcopy(continuation_activity)
+        original_activity_id = replay_activity.id
+        replay_activity.apply_conversation_reference(
+            context.activity.get_conversation_reference(), is_incoming=True
+        )
+        replay_activity.id = original_activity_id
+
+        async def replay_turn(replay_context: TurnContext) -> None:
+            self._auth._copy_cached_tokens(context, replay_context)
+            await self._on_turn(replay_context)
+
+        async def replay() -> None:
+            await context.adapter.continue_conversation_with_claims(
+                context.identity,
+                replay_activity,
+                replay_turn,
+            )
+
+        task = asyncio.create_task(replay())
+        self._background_tasks.add(task)
+        task.add_done_callback(self._auth_continuation_done)
+
+    def _auth_continuation_done(self, task: asyncio.Task[None]) -> None:
+        self._background_tasks.discard(task)
+        if task.cancelled():
+            return
+        try:
+            task.result()
+        except Exception:  # noqa: BLE001
+            logger.exception("OAuth continuation replay failed")
+
     def _remove_mentions(self, context: TurnContext):
         if (
             self.options.remove_recipient_mention

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

@@ -201,6 +201,15 @@ def _cache_token(
         key = Authorization._cache_key(context, handler_id)
         context.turn_state[key] = token_response
 
+    def _copy_cached_tokens(
+        self, source_context: TurnContext, target_context: TurnContext
+    ) -> None:
+        """Copy completed OAuth tokens into a new context for the same user turn."""
+        for handler_id in self._handlers:
+            token_response = Authorization._get_cached_token(source_context, handler_id)
+            if token_response is not None:
+                Authorization._cache_token(target_context, handler_id, token_response)
+
     @staticmethod
     def _delete_cached_token(context: TurnContext, handler_id: str) -> None:
         key = Authorization._cache_key(context, handler_id)

tests/hosting_core/app/test_agent_application.py

@@ -4,11 +4,11 @@
 """
 
 import asyncio
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
-from microsoft_agents.activity import Activity, ActivityTypes
+from microsoft_agents.activity import Activity, ActivityTypes, TokenResponse
 from microsoft_agents.hosting.core import MemoryStorage
 from microsoft_agents.hosting.core.app import (
     AgentApplication,
@@ -77,6 +77,98 @@ async def send_activity(self, activity):
         self.responded = True
 
 
+@pytest.mark.asyncio
+async def test_token_exchange_invoke_returns_before_continuation_replay_completes():
+    """Token exchange must not keep the invoke open while the message is replayed."""
+    app = _make_integration_app()
+    replay_started = asyncio.Event()
+    release_replay = asyncio.Event()
+    replay_completed = asyncio.Event()
+
+    continuation = Activity(
+        type=ActivityTypes.message,
+        id="message-id",
+        text="hello",
+        channel_id="msteams:COPILOT",
+        conversation={"id": "conversation"},
+        from_property={"id": "user"},
+        recipient={"id": "agent"},
+        service_url="https://message.example.org",
+    )
+    invoke = Activity(
+        type=ActivityTypes.invoke,
+        id="invoke-id",
+        name="signin/tokenExchange",
+        channel_id="msteams:COPILOT",
+        conversation={"id": "conversation"},
+        from_property={"id": "user"},
+        recipient={"id": "agent"},
+        service_url="https://invoke.example.org",
+    )
+    context = StubTurnContext(invoke)
+    context.identity = MagicMock()
+    token_response = TokenResponse(token="access-token")
+    app._auth._handlers["handler"] = MagicMock()
+    Authorization._cache_token(context, "handler", token_response)
+    turn_state = MagicMock()
+    call_order = []
+
+    async def save_turn_state(_context):
+        call_order.append("save")
+
+    turn_state.save = AsyncMock(side_effect=save_turn_state)
+
+    async def replay_turn(replay_context):
+        assert (
+            Authorization._get_cached_token(replay_context, "handler") is token_response
+        )
+        replay_started.set()
+        await release_replay.wait()
+        replay_completed.set()
+
+    async def continue_conversation(identity, replay_activity, callback):
+        call_order.append("continue")
+        replay_context = StubTurnContext(replay_activity, context.adapter)
+        replay_context.identity = identity
+        await callback(replay_context)
+
+    context.adapter.continue_conversation_with_claims = AsyncMock(
+        side_effect=continue_conversation
+    )
+
+    original_on_turn = app._on_turn
+    app.on_turn = AsyncMock(side_effect=replay_turn)
+    app._on_turn = AsyncMock(side_effect=replay_turn)
+    app._auth._on_turn_auth_intercept = AsyncMock(return_value=(True, continuation))
+
+    with patch.object(
+        app, "_initialize_state", new_callable=AsyncMock, return_value=turn_state
+    ):
+        await asyncio.wait_for(original_on_turn(context), timeout=0.1)
+
+    await asyncio.wait_for(replay_started.wait(), timeout=0.1)
+    assert not replay_completed.is_set()
+    turn_state.save.assert_awaited_once_with(context)
+    context.adapter.continue_conversation_with_claims.assert_awaited_once()
+    replay_args = context.adapter.continue_conversation_with_claims.call_args.args
+    assert replay_args[0] is context.identity
+    replay_activity = replay_args[1]
+    assert replay_args[2] != app._on_turn
+    assert replay_activity is not continuation
+    assert replay_activity.id == "message-id"
+    assert replay_activity.type == ActivityTypes.message
+    assert replay_activity.channel_id == "msteams:COPILOT"
+    assert replay_activity.service_url == "https://invoke.example.org"
+    assert context.activity.type == ActivityTypes.invoke
+    assert context.activity.id == "invoke-id"
+    assert context.activity.channel_id == "msteams:COPILOT"
+    assert call_order[:2] == ["save", "continue"]
+    app.on_turn.assert_not_awaited()
+
+    release_replay.set()
+    await asyncio.wait_for(replay_completed.wait(), timeout=0.1)
+
+
 def make_auth():
     return Authorization(
         storage=MemoryStorage(),