From fb89185fba4f7f5f46505c0d1db4f10695222a5a Mon Sep 17 00:00:00 2001 From: David Carvalho Date: Fri, 10 Jul 2026 22:17:40 +0200 Subject: [PATCH] Detach OAuth continuation replay from invoke --- .../hosting/core/app/agent_application.py | 50 +++++++++- .../hosting/core/app/oauth/authorization.py | 9 ++ .../app/test_agent_application.py | 96 ++++++++++++++++++- 3 files changed, 149 insertions(+), 6 deletions(-) diff --git a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/agent_application.py b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/agent_application.py index 748060ee..b0662a91 100644 --- a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/agent_application.py +++ b/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 diff --git a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/oauth/authorization.py b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/oauth/authorization.py index 542b2866..113d361e 100644 --- a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/oauth/authorization.py +++ b/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) diff --git a/tests/hosting_core/app/test_agent_application.py b/tests/hosting_core/app/test_agent_application.py index 30c56f73..763bff00 100644 --- a/tests/hosting_core/app/test_agent_application.py +++ b/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(),