diff --git a/CHANGELOG.md b/CHANGELOG.md index f41a770..967775a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,14 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/), and this ## [Unreleased] +## [0.9.4] - 2026-10-21 + +### Fixed + +- OAuth (MCP and A2A): an expired sign-in no longer blocks the request for 5 minutes and ends in `TimeoutError`. When the IdP refuses the stored refresh token (`invalid_grant`/`invalid_token`, e.g. Keycloak "Token is not active" once an offline session reaches its maximum lifespan), the dead token is deleted, the user is sent a fresh sign-in link straight away, and the call returns a `sign_in_expired` result telling the user to sign in and ask again. The sign-in completes in the background, so the next request uses the new token. Background pollers get the same `ReauthRequired` handling as before, without a prompt. + +- MCP over HTTP (`slack_agents.tools.mcp_http`): a closed or dead session is now reopened instead of failing every later call. Previously the session was opened once at startup and never replaced, so after the connection dropped each tool call failed at once with `anyio.ClosedResourceError` until the agent restarted. A call whose request could not be written to the closed transport (`ClosedResourceError`/`BrokenResourceError`, so it never reached the server) is retried once on a new session. When the connection closes after the request was sent, or the server no longer knows the session (`Session terminated`, e.g. after a server restart), the call fails without replay and the next call reconnects. The connection now runs in its own task, so the transport is always entered and exited by the same task. + ## [0.9.3] - 2026-07-08 ### Fixed diff --git a/docs/oauth.md b/docs/oauth.md index 63fbe33..d3657f4 100644 --- a/docs/oauth.md +++ b/docs/oauth.md @@ -97,6 +97,14 @@ If they don't click within the configured `auth_timeout` (default 5 minutes), the agent surfaces a "timed out — please try again" error and the tool call ends. They can re-ask whenever they're ready and a fresh prompt appears. +Tokens are refreshed silently. If the identity provider refuses the refresh +token (`invalid_grant`/`invalid_token`, for example once a Keycloak offline +session reaches its maximum lifespan), the agent deletes the dead token, sends +a fresh "Authenticate" prompt straight away, and ends that tool call telling +the user their sign-in expired and to sign in and ask again. The request does +not wait for them; the sign-in completes in the background and their next +request uses the new token. + If the upstream tool later requires additional permissions (e.g. they had read access but are now trying to write), the same flow re-runs requesting the broader scope. Most identity providers auto-collapse the consent screen if the diff --git a/llms-full.txt b/llms-full.txt index 93b7880..20dd1d3 100644 --- a/llms-full.txt +++ b/llms-full.txt @@ -1458,6 +1458,14 @@ If they don't click within the configured `auth_timeout` (default 5 minutes), the agent surfaces a "timed out — please try again" error and the tool call ends. They can re-ask whenever they're ready and a fresh prompt appears. +Tokens are refreshed silently. If the identity provider refuses the refresh +token (`invalid_grant`/`invalid_token`, for example once a Keycloak offline +session reaches its maximum lifespan), the agent deletes the dead token, sends +a fresh "Authenticate" prompt straight away, and ends that tool call telling +the user their sign-in expired and to sign in and ask again. The request does +not wait for them; the sign-in completes in the background and their next +request uses the new token. + If the upstream tool later requires additional permissions (e.g. they had read access but are now trying to write), the same flow re-runs requesting the broader scope. Most identity providers auto-collapse the consent screen if the diff --git a/pyproject.toml b/pyproject.toml index c969355..2eb4e1e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "python-slack-agents" -version = "0.9.3" +version = "0.9.4" description = "A Python framework for deploying AI agents as Slack bots" authors = [{name = "Eric Pichon", email = "epichon@comparenetworks.com"}] readme = "README.md" diff --git a/src/slack_agents/a2a/agent.py b/src/slack_agents/a2a/agent.py index c317885..613f684 100644 --- a/src/slack_agents/a2a/agent.py +++ b/src/slack_agents/a2a/agent.py @@ -307,6 +307,17 @@ async def call_tool( ) except Exception as e: if self._oauth_mode: + from slack_agents.oauth.errors import ( # noqa: PLC0415 + find_sign_in_expired, + sign_in_expired_tool_result, + ) + + expired = find_sign_in_expired(e) + if expired is not None: + # The cached client's provider holds the dead token; rebuild + # from storage once the user has signed in again. + self._user_clients.pop(user_conversation_context["user_id"], None) + return sign_in_expired_tool_result(expired, tool=name) handled = await self._handle_oauth_error(e, name, user_conversation_context) if handled is not None: return handled diff --git a/src/slack_agents/oauth/errors.py b/src/slack_agents/oauth/errors.py index 605b336..c7e2752 100644 --- a/src/slack_agents/oauth/errors.py +++ b/src/slack_agents/oauth/errors.py @@ -8,6 +8,7 @@ from slack_agents.tools.base import ( ERROR_PERMISSION_DENIED, ERROR_SYSTEM_ERROR, + RECOVERY_ABORT, RECOVERY_CONTACT_ADMIN, RECOVERY_CONTACT_SUPPORT, RECOVERY_RETRY, @@ -85,6 +86,17 @@ class ReauthRequired(Exception): Slack auth prompt when a non-interactive OAuth flow hits a full re-auth.""" +class SignInExpired(ReauthRequired): + """The IdP refused the stored refresh token (e.g. the offline session reached + its maximum lifespan). The dead token has been deleted; in an interactive + context a fresh sign-in link has already been posted to the user.""" + + def __init__(self, server_id: str, *, prompted: bool) -> None: + self.server_id = server_id + self.prompted = prompted + super().__init__(f"sign-in to {server_id} expired (refresh token refused)") + + class UserAuthorizationDenied(Exception): """Raised when the upstream IdP (or the resource server) rejects access because the user's account lacks permission for the operation. Distinct @@ -248,3 +260,30 @@ def message_from_tool_error(tr: ToolResult) -> str: except (json.JSONDecodeError, TypeError): return tr["content"] return payload.get("message") or tr["content"] + + +def find_sign_in_expired(exc: BaseException) -> SignInExpired | None: + for x in flatten_exceptions(exc): + if isinstance(x, SignInExpired): + return x + return None + + +def sign_in_expired_tool_result(expired: SignInExpired, *, tool: str | None = None) -> ToolResult: + """Tell the user their sign-in expired and to sign in and ask again. + + The LLM must not retry in this turn: the user has to sign in first. + """ + link = "I've just sent you a sign-in link (only you can see it). " if expired.prompted else "" + return make_tool_error( + error=ERROR_PERMISSION_DENIED, + code="sign_in_expired", + server=expired.server_id, + tool=tool, + recovery=RECOVERY_ABORT, + message=( + f"Your sign-in to {expired.server_id} has expired. {link}" + "Please sign in, then ask again. Tell the user exactly this and stop; " + "do not retry this tool in this turn." + ), + ) diff --git a/src/slack_agents/oauth/flow.py b/src/slack_agents/oauth/flow.py index 1d13739..b8ff91d 100644 --- a/src/slack_agents/oauth/flow.py +++ b/src/slack_agents/oauth/flow.py @@ -20,7 +20,7 @@ from slack_agents import OAuthCallbackResult, PendingFlowsRegistry from slack_agents.oauth.discovery import DiscoveryResult, DiscoveryStrategy -from slack_agents.oauth.errors import ReauthRequired, UserAuthorizationDenied +from slack_agents.oauth.errors import ReauthRequired, SignInExpired, UserAuthorizationDenied from slack_agents.oauth.prompts import AuthPromptDeliveryError, send_auth_prompt from slack_agents.oauth.scopes import ( OIDC_BASELINE_SCOPES, @@ -33,6 +33,45 @@ logger = logging.getLogger(__name__) +# Token-endpoint errors meaning the refresh token itself is dead (expired, revoked, +# or its session gone): retrying can never succeed, only a new sign-in can. +_DEAD_REFRESH_ERRORS = frozenset({"invalid_grant", "invalid_token"}) + + +class _PerUserProvider(OAuthClientProvider): + """OAuthClientProvider that fails fast when the IdP refuses the refresh token. + + The SDK's default clears the token in memory and falls through to a full + authorization inside the same request, blocking it until the user signs in + or the callback wait times out. Instead the dead token is deleted at once, + `on_refresh_refused` starts a fresh sign-in, and `SignInExpired` ends the call. + """ + + def __init__(self, *args, on_refresh_refused, server_id: str, interactive: bool, **kwargs): + super().__init__(*args, **kwargs) + self._on_refresh_refused = on_refresh_refused + self._server_id = server_id + self._interactive = interactive + + async def _handle_refresh_response(self, response: httpx.Response) -> bool: + if response.status_code in (400, 401): + await response.aread() + try: + error = response.json().get("error") + except ValueError: + error = None + if error in _DEAD_REFRESH_ERRORS: + logger.warning( + "oauth: %s refused the refresh token (HTTP %s %s) — sign-in expired", + self._server_id, + response.status_code, + error, + ) + self.context.clear_tokens() + await self._on_refresh_refused() + raise SignInExpired(self._server_id, prompted=self._interactive) + return await super()._handle_refresh_response(response) + @dataclass class PerUserHandle: @@ -138,6 +177,7 @@ def __init__( self._auth_timeout = auth_timeout self._discovery = discovery self._cached: DiscoveryResult | None = None + self._sign_ins: set[asyncio.Task] = set() def make_handle(self, user_id, channel_id, thread_id, *, interactive=True) -> PerUserHandle: return PerUserHandle( @@ -200,7 +240,17 @@ async def build_provider( response_types=["code"], scope=seed_scope, ) - oauth = OAuthClientProvider( + + async def on_refresh_refused() -> None: + scopes = await self.cached_scope_set(user_id) + await self._storage.delete_oauth_token(user_id, self._server_id) + if interactive: + self.start_sign_in(user_id, channel_id, thread_id, scopes) + + oauth = _PerUserProvider( + on_refresh_refused=on_refresh_refused, + server_id=self._server_id, + interactive=interactive, server_url=self._server_url, client_metadata=client_metadata, storage=token_storage, @@ -218,6 +268,42 @@ async def build_provider( logger.warning("oauth: pre-initialize for %s failed", self._server_id, exc_info=True) return oauth + def start_sign_in( + self, user_id: str, channel_id: str, thread_id: str | None, scopes: set[str] + ) -> None: + """Post the user a sign-in link now and finish the sign-in in the background. + + The turn that hit the expired sign-in returns immediately; when the user + completes sign-in the new token is stored and their next request uses it. + """ + task = asyncio.create_task(self._sign_in(user_id, channel_id, thread_id, scopes)) + self._sign_ins.add(task) + task.add_done_callback(self._sign_ins.discard) + + async def _sign_in( + self, user_id: str, channel_id: str, thread_id: str | None, scopes: set[str] + ) -> None: + try: + oauth = await self.build_provider(user_id, channel_id, thread_id) + # Ask for at least what the expired token had, so the new token keeps + # offline_access and the resource scopes. + requested = set((oauth.context.client_metadata.scope or "").split()) + oauth.context.client_metadata.scope = " ".join( + sorted(OIDC_BASELINE_SCOPES | requested | scopes) + ) + token_request = await oauth._perform_authorization() + async with httpx.AsyncClient(timeout=httpx.Timeout(30.0)) as http: + response = await http.send(token_request) + await oauth._handle_token_response(response) + logger.info("oauth: %s user %s signed in again", self._server_id, user_id) + except Exception: + logger.warning( + "oauth: %s background sign-in for %s did not complete", + self._server_id, + user_id, + exc_info=True, + ) + def _apply_cached_metadata(self, oauth: OAuthClientProvider) -> None: if self._cached is None: return diff --git a/src/slack_agents/tools/mcp_http.py b/src/slack_agents/tools/mcp_http.py index 2910a9b..465de9c 100644 --- a/src/slack_agents/tools/mcp_http.py +++ b/src/slack_agents/tools/mcp_http.py @@ -7,10 +7,12 @@ import logging from urllib.parse import unquote, urlparse +import anyio import httpx import mcp from mcp.client.streamable_http import streamable_http_client -from mcp.types import BlobResourceContents, EmbeddedResource, ImageContent +from mcp.shared.exceptions import McpError +from mcp.types import CONNECTION_CLOSED, BlobResourceContents, EmbeddedResource, ImageContent from slack_agents import UserConversationContext from slack_agents.llm import CHARS_PER_TOKEN @@ -36,10 +38,30 @@ def _uri_to_filename(uri: str) -> str: return name or "file" +def _session_is_gone(exc: Exception) -> bool: + """True when an MCP error means the session can never be used again. + + CONNECTION_CLOSED: the transport ended while the request was pending. + "Session terminated" (code 32600): the server answered 404 for our session id, + e.g. after the server restarted. + """ + if not isinstance(exc, McpError): + return False + code, message = exc.error.code, exc.error.message + return code == CONNECTION_CLOSED or (code == 32600 and message == "Session terminated") + + class Provider(BaseToolProvider): - """MCP over HTTP tool provider. Connects to a single MCP server.""" + """MCP over HTTP tool provider. Connects to a single MCP server. + + The session lives in a dedicated task that owns the transport (see + `_run_connection`). When the connection dies the session is dropped and the + next call reconnects; a call whose request could not even be written to the + closed transport is retried once on the new session. + """ DEFAULT_INIT_RETRIES = [5, 10, 30] + DISCONNECT_TIMEOUT = 5.0 def __init__( self, @@ -52,34 +74,94 @@ def __init__( self._url = url self._headers = headers or {} self._init_retries = init_retries if init_retries is not None else self.DEFAULT_INIT_RETRIES - self._tool_map: dict[str, mcp.ClientSession] = {} + self._tool_names: set[str] = set() self._all_tools: list[dict] = [] self._session: mcp.ClientSession | None = None - self._exit_stack: contextlib.AsyncExitStack | None = None + self._conn_task: asyncio.Task | None = None + self._conn_stop: asyncio.Event | None = None + self._reconnect_lock = asyncio.Lock() def _get_all_tools(self) -> list[dict]: return self._all_tools - async def _connect(self) -> None: - """Establish connection to the MCP server.""" - http_client = httpx.AsyncClient( - headers=self._headers, - timeout=httpx.Timeout(30.0, read=300.0), - follow_redirects=True, - ) + async def _connect(self) -> mcp.ClientSession: + """Open a session in a new connection task and wait until it is initialized. - stack = contextlib.AsyncExitStack() - self._exit_stack = stack + The transport's anyio task groups must be entered and exited by the same + task, so the connection runs in its own task for its whole life instead + of in whichever task happens to call initialize, close or reconnect. + """ + ready: asyncio.Future[mcp.ClientSession] = asyncio.get_running_loop().create_future() + stop = asyncio.Event() + task = asyncio.create_task(self._run_connection(ready, stop)) + try: + session = await ready + except BaseException: + task.cancel() + with contextlib.suppress(BaseException): + await task + raise + self._session, self._conn_task, self._conn_stop = session, task, stop + return session - await stack.enter_async_context(http_client) - read_stream, write_stream, _get_session_id = await stack.enter_async_context( - streamable_http_client(url=self._url, http_client=http_client) - ) + async def _run_connection( + self, ready: "asyncio.Future[mcp.ClientSession]", stop: asyncio.Event + ) -> None: + try: + async with contextlib.AsyncExitStack() as stack: + http_client = await stack.enter_async_context( + httpx.AsyncClient( + headers=self._headers, + timeout=httpx.Timeout(30.0, read=300.0), + follow_redirects=True, + ) + ) + read_stream, write_stream, _get_session_id = await stack.enter_async_context( + streamable_http_client(url=self._url, http_client=http_client) + ) + session = await stack.enter_async_context( + mcp.ClientSession(read_stream, write_stream) + ) + await session.initialize() + # Populate the session's tool output-schema cache now, so call_tool + # never sends a list_tools of its own after the tool has run. + await session.list_tools() + ready.set_result(session) + await stop.wait() + except BaseException as exc: + if not ready.done(): + if isinstance(exc, asyncio.CancelledError): + ready.cancel() + else: + ready.set_exception(exc) + elif not stop.is_set(): + logger.warning("MCP %s: connection closed: %r", self._url, exc) + if isinstance(exc, asyncio.CancelledError): + raise + finally: + # A connection that ended on its own leaves no usable session. + if self._conn_task is asyncio.current_task(): + self._session = None - session = mcp.ClientSession(read_stream, write_stream) - await stack.enter_async_context(session) - await session.initialize() - self._session = session + async def _disconnect(self) -> None: + task, stop = self._conn_task, self._conn_stop + self._session = self._conn_task = self._conn_stop = None + if task is None: + return + stop.set() + try: + await asyncio.wait_for(task, self.DISCONNECT_TIMEOUT) + except BaseException: + logger.debug("MCP %s: error closing connection", self._url, exc_info=True) + + async def _reconnect(self, stale: mcp.ClientSession | None) -> mcp.ClientSession: + """Replace `stale` with a fresh session. Concurrent callers share one reconnect.""" + async with self._reconnect_lock: + if self._session is not None and self._session is not stale: + return self._session + logger.warning("MCP %s: reconnecting", self._url) + await self._disconnect() + return await self._connect() async def initialize(self) -> None: """Connect to the MCP server and discover tools, retrying on connection errors.""" @@ -108,11 +190,6 @@ async def initialize(self) -> None: exc, backoff, ) - # Clean up partial state before retrying - if self._exit_stack: - with contextlib.suppress(BaseException): - await self._exit_stack.aclose() - self._exit_stack = None await asyncio.sleep(backoff) tools_result = await self._session.list_tools() @@ -126,7 +203,7 @@ async def initialize(self) -> None: } tool_tokens = len(json.dumps(tool_def)) // CHARS_PER_TOKEN server_tokens += tool_tokens - self._tool_map[tool.name] = self._session + self._tool_names.add(tool.name) self._all_tools.append(tool_def) # Log filtered tools @@ -156,8 +233,7 @@ async def call_tool( storage: BaseStorageProvider, ) -> ToolResult: """Execute a tool call and return the result.""" - session = self._tool_map.get(tool_name) - if not session: + if tool_name not in self._tool_names: return make_tool_error( error=ERROR_INPUT_ERROR, code="unknown_tool", @@ -170,7 +246,17 @@ async def call_tool( logger.info("Calling MCP tool %s", tool_name) try: - result = await session.call_tool(name=tool_name, arguments=arguments) + session = self._session or await self._reconnect(None) + try: + result = await session.call_tool(name=tool_name, arguments=arguments) + except (anyio.ClosedResourceError, anyio.BrokenResourceError): + # Raised by the write to the session's closed transport stream, so + # the request never reached the server and sending it once more on + # a new session cannot run the tool twice. Any later failure may + # follow execution and is never replayed. + logger.warning("MCP %s: session closed, reconnecting for %s", self._url, tool_name) + session = await self._reconnect(session) + result = await session.call_tool(name=tool_name, arguments=arguments) text_parts = [] files = [] @@ -203,6 +289,10 @@ async def call_tool( "files": files, } except Exception as e: + if _session_is_gone(e): + # The request reached the server or may have run, so it is not + # retried; the next call reconnects. + self._session = None logger.exception("MCP tool call failed: %s", tool_name) return make_tool_error( error=ERROR_SYSTEM_ERROR, @@ -214,12 +304,6 @@ async def call_tool( ) async def close(self) -> None: - if self._exit_stack: - try: - await self._exit_stack.aclose() - except Exception: - logger.exception("Error closing MCP connection") - self._exit_stack = None - self._session = None - self._tool_map.clear() - self._all_tools.clear() + await self._disconnect() + self._tool_names.clear() + self._all_tools.clear() diff --git a/src/slack_agents/tools/mcp_http_oauth.py b/src/slack_agents/tools/mcp_http_oauth.py index 61fbfce..8f2104d 100644 --- a/src/slack_agents/tools/mcp_http_oauth.py +++ b/src/slack_agents/tools/mcp_http_oauth.py @@ -35,10 +35,12 @@ from slack_agents.oauth.errors import ( AuthSetupError, UserAuthorizationDenied, + find_sign_in_expired, find_user_authorization_denied, is_redirect_uri_mismatch, message_from_tool_error, record_error, + sign_in_expired_tool_result, user_denied_tool_result, ) from slack_agents.oauth.flow import PerUserOAuth @@ -254,6 +256,9 @@ async def call_tool( try: return await self._call_mcp_with_token(user_conversation_context, tool_name, arguments) except Exception as e: + expired = find_sign_in_expired(e) + if expired is not None: + return sign_in_expired_tool_result(expired, tool=tool_name) # Specific recovery: IdP rejected our authorize request because the # client's registered redirect_uri doesn't match what we sent. # Means the cached `oauth_clients` row's redirect_uri is stale — diff --git a/tests/a2a/test_oauth.py b/tests/a2a/test_oauth.py index 5c279d1..6528cc3 100644 --- a/tests/a2a/test_oauth.py +++ b/tests/a2a/test_oauth.py @@ -456,3 +456,19 @@ def handler(request: httpx.Request) -> httpx.Response: assert row is not None assert row.client_id == "fresh-cid" await storage.close() + + +async def test_call_tool_sign_in_expired_returns_clear_message(monkeypatch, store): + import json + + from slack_agents.oauth.errors import SignInExpired + + p, result = await _build_oauth_provider( + monkeypatch, store, SignInExpired("secure-agent", prompted=True) + ) + payload = json.loads(result["content"]) + assert result["is_error"] is True + assert payload["code"] == "sign_in_expired" + assert "sign in, then ask again" in payload["message"] + # The cached client holds the dead token; it is rebuilt after sign-in. + assert "U1" not in p._user_clients diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index 81657ab..e950b4f 100644 --- a/tests/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -1,11 +1,24 @@ """Tests for MCP HTTP tool provider binary result extraction.""" +import asyncio import base64 +import contextlib from unittest.mock import AsyncMock, MagicMock +import anyio +import httpx import pytest -from mcp.types import BlobResourceContents, EmbeddedResource, ImageContent, TextContent +from mcp.shared.exceptions import McpError +from mcp.types import ( + CONNECTION_CLOSED, + BlobResourceContents, + EmbeddedResource, + ErrorData, + ImageContent, + TextContent, +) +from slack_agents.tools import mcp_http from slack_agents.tools.mcp_http import Provider, _uri_to_filename @@ -31,7 +44,8 @@ class TestCallToolBinaryResults: def provider(self): prov = Provider(url="https://example.com/mcp", allowed_functions=[".*"]) session = AsyncMock() - prov._tool_map["test_tool"] = session + prov._tool_names.add("test_tool") + prov._session = session return prov, session async def test_text_content(self, provider): @@ -115,3 +129,190 @@ async def test_unknown_tool(self): out = await prov.call_tool("nonexistent", {}, None, None) assert out["is_error"] is True assert out["files"] == [] + + +def _ok_result(text="ok"): + result = MagicMock() + result.content = [TextContent(type="text", text=text)] + result.isError = False + return result + + +class TestReconnect: + """A closed MCP session is reopened; a call is replayed only if it never reached the server.""" + + @pytest.fixture + def provider(self): + prov = Provider(url="https://example.com/mcp", allowed_functions=[".*"]) + prov._tool_names.add("scrape") + prov._session = AsyncMock(name="stale") + fresh = [] + + async def connect(): + session = AsyncMock(name=f"fresh{len(fresh)}") + session.call_tool.return_value = _ok_result() + fresh.append(session) + prov._session = session + return session + + prov._connect = AsyncMock(side_effect=connect) + return prov, fresh + + @pytest.mark.parametrize("exc", [anyio.ClosedResourceError, anyio.BrokenResourceError]) + async def test_send_on_closed_stream_reconnects_and_retries_once(self, provider, exc): + prov, fresh = provider + stale = prov._session + stale.call_tool.side_effect = exc() + + out = await prov.call_tool("scrape", {"url": "u"}, None, None) + + assert out["is_error"] is False + assert out["content"] == "ok" + assert stale.call_tool.await_count == 1 + assert prov._connect.await_count == 1 + fresh[0].call_tool.assert_awaited_once_with(name="scrape", arguments={"url": "u"}) + + async def test_retries_at_most_once(self, provider): + prov, fresh = provider + prov._session.call_tool.side_effect = anyio.ClosedResourceError() + + async def connect_dead(): + session = AsyncMock() + session.call_tool.side_effect = anyio.ClosedResourceError() + fresh.append(session) + prov._session = session + return session + + prov._connect.side_effect = connect_dead + + out = await prov.call_tool("scrape", {}, None, None) + + assert out["is_error"] is True + assert prov._connect.await_count == 1 + assert fresh[0].call_tool.await_count == 1 + + @pytest.mark.parametrize( + "error", + [ + ErrorData(code=CONNECTION_CLOSED, message="Connection closed"), + # The SDK's error for a 404 on our session id (server restarted). + ErrorData(code=32600, message="Session terminated"), + ], + ) + async def test_dead_session_after_send_is_not_replayed(self, provider, error): + prov, fresh = provider + stale = prov._session + stale.call_tool.side_effect = McpError(error) + + out = await prov.call_tool("scrape", {}, None, None) + + assert out["is_error"] is True + assert stale.call_tool.await_count == 1 + assert prov._connect.await_count == 0 + # ...but the next call gets a fresh session. + out = await prov.call_tool("scrape", {}, None, None) + assert out["is_error"] is False + assert prov._connect.await_count == 1 + + async def test_other_errors_keep_the_session(self, provider): + prov, _ = provider + stale = prov._session + stale.call_tool.side_effect = McpError(ErrorData(code=-32603, message="boom")) + + out = await prov.call_tool("scrape", {}, None, None) + + assert out["is_error"] is True + assert prov._session is stale + assert prov._connect.await_count == 0 + + async def test_concurrent_failures_reconnect_once(self, provider): + prov, fresh = provider + prov._session.call_tool.side_effect = anyio.ClosedResourceError() + + outs = await asyncio.gather(*(prov.call_tool("scrape", {}, None, None) for _ in range(3))) + + assert all(o["is_error"] is False for o in outs) + assert prov._connect.await_count == 1 + assert fresh[0].call_tool.await_count == 3 + + async def test_reconnect_failure_is_a_tool_error_and_next_call_tries_again(self, provider): + prov, fresh = provider + prov._session.call_tool.side_effect = anyio.ClosedResourceError() + real_connect = prov._connect.side_effect + prov._connect.side_effect = httpx.ConnectError("down") + + out = await prov.call_tool("scrape", {}, None, None) + assert out["is_error"] is True + assert prov._session is None + + prov._connect.side_effect = real_connect + out = await prov.call_tool("scrape", {}, None, None) + assert out["is_error"] is False + + +class _FakeCM: + """Async context manager recording which task entered and exited it.""" + + def __init__(self, log, name, value=None): + self.log, self.name, self.value = log, name, value + + async def __aenter__(self): + self.log.append(("enter", self.name, asyncio.current_task())) + return self.value + + async def __aexit__(self, *exc): + self.log.append(("exit", self.name, asyncio.current_task())) + return False + + +class TestConnectionTask: + @pytest.fixture + def fakes(self, monkeypatch): + log = [] + sessions = [] + + def transport(url, http_client): + return _FakeCM(log, "transport", (MagicMock(), MagicMock(), MagicMock())) + + def client_session(read, write): + session = AsyncMock() + session.list_tools.return_value = MagicMock(tools=[]) + sessions.append(session) + return _FakeCM(log, "session", session) + + monkeypatch.setattr(mcp_http, "streamable_http_client", transport) + monkeypatch.setattr(mcp_http.mcp, "ClientSession", client_session) + return log, sessions + + async def test_transport_is_entered_and_exited_by_one_task(self, fakes): + log, sessions = fakes + prov = Provider(url="https://example.com/mcp", allowed_functions=[".*"]) + + await prov.initialize() + assert prov._session is sessions[0] + sessions[0].list_tools.assert_awaited() # output-schema cache filled + + # Reconnect and close from other tasks; each connection still exits in + # the task that opened it. + await asyncio.create_task(prov._reconnect(prov._session)) + assert prov._session is sessions[1] + await asyncio.create_task(prov.close()) + assert prov._session is None + + enters = [t for k, n, t in log if k == "enter" and n == "transport"] + exits = [t for k, n, t in log if k == "exit" and n == "transport"] + assert len(enters) == len(exits) == 2 + assert enters == exits + assert asyncio.current_task() not in enters + + async def test_connection_that_dies_drops_the_session(self, fakes): + _, sessions = fakes + prov = Provider(url="https://example.com/mcp", allowed_functions=[".*"]) + await prov.initialize() + + prov._conn_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await prov._conn_task + + assert prov._session is None + await prov.close() diff --git a/tests/test_oauth_sign_in_expired.py b/tests/test_oauth_sign_in_expired.py new file mode 100644 index 0000000..669bcbb --- /dev/null +++ b/tests/test_oauth_sign_in_expired.py @@ -0,0 +1,183 @@ +"""A refresh token the IdP refuses ends the call at once with a fresh sign-in link.""" + +import asyncio +import json + +import httpx +import pytest +from mcp.shared.auth import OAuthClientInformationFull, OAuthMetadata, OAuthToken + +from slack_agents import PendingFlowsRegistry +from slack_agents.oauth import flow +from slack_agents.oauth.crypto import derive_subkeys +from slack_agents.oauth.discovery import DiscoveryResult +from slack_agents.oauth.errors import ( + ReauthRequired, + SignInExpired, + sign_in_expired_tool_result, +) +from slack_agents.oauth.flow import PerUserOAuth +from slack_agents.oauth.storage import DBTokenStorage +from slack_agents.storage.sqlite import Provider as Sqlite + +TOKEN_URL = "https://idp.example.com/token" +RS_URL = "https://rs.example.com/a2a" + + +class _StubDiscovery: + async def discover(self): + raise AssertionError("discovery is pre-cached") + + +@pytest.fixture +async def store(): + s = Sqlite(path=":memory:") + await s.initialize() + yield s + await s.close() + + +@pytest.fixture +async def puo(store): + sk, tk = derive_subkeys(b"0" * 32) + p = PerUserOAuth( + server_url=RS_URL, + server_id="agent-geo", + agent_name="t", + storage=store, + pending_flows=PendingFlowsRegistry(), + slack_client=None, + state_key=sk, + token_key=tk, + redirect_uri="https://a.example.com/oauth/callback", + public_url="https://a.example.com", + auth_timeout=300, + discovery=_StubDiscovery(), + ) + p._cached = DiscoveryResult( + resource_metadata=None, + authorization_server_metadata=OAuthMetadata( + issuer="https://idp.example.com", + authorization_endpoint="https://idp.example.com/auth", + token_endpoint=TOKEN_URL, + ), + scope_catalog=[], + required_scopes=[], + ) + storage = DBTokenStorage( + backend=store, + user_id="__system__", + server_id="agent-geo", + redirect_uri=p._redirect_uri, + token_key=tk, + ) + await storage.set_client_info( + OAuthClientInformationFull(client_id="c1", redirect_uris=[p._redirect_uri]) + ) + # An expired access token with a refresh token, as stored after 60 days. + await p._token_storage("U1").set_tokens( + OAuthToken( + access_token="old", + token_type="Bearer", + expires_in=-60, + refresh_token="dead-refresh", + scope="openid offline_access agent:agent-geo:read", + ) + ) + return p + + +def _client(provider, token_response): + seen = [] + + def handler(request): + seen.append(str(request.url)) + if str(request.url) == TOKEN_URL: + return token_response + return httpx.Response(200, json={"ok": True}) + + return httpx.AsyncClient(auth=provider, transport=httpx.MockTransport(handler)), seen + + +async def test_refused_refresh_deletes_token_prompts_and_fails_fast(puo, store, monkeypatch): + started = [] + monkeypatch.setattr(puo, "start_sign_in", lambda *a: started.append(a)) + provider = await puo.build_provider("U1", "C1", "T1") + client, seen = _client( + provider, + httpx.Response( + 400, json={"error": "invalid_grant", "error_description": "Token is not active"} + ), + ) + + with pytest.raises(SignInExpired) as info: + async with client: + await asyncio.wait_for(client.post(RS_URL, json={}), timeout=5) + + assert info.value.prompted is True + assert isinstance(info.value, ReauthRequired) # the poller's existing handling still applies + assert seen == [TOKEN_URL] # never sent to the resource, never waited for a callback + assert await store.get_oauth_token("U1", "agent-geo") is None + assert started == [("U1", "C1", "T1", {"openid", "offline_access", "agent:agent-geo:read"})] + + +async def test_noninteractive_refused_refresh_does_not_prompt(puo, store, monkeypatch): + started = [] + monkeypatch.setattr(puo, "start_sign_in", lambda *a: started.append(a)) + provider = await puo.build_provider("U1", "C1", "T1", interactive=False) + client, _ = _client(provider, httpx.Response(400, json={"error": "invalid_token"})) + + with pytest.raises(SignInExpired) as info: + async with client: + await client.post(RS_URL, json={}) + + assert info.value.prompted is False + assert started == [] + assert await store.get_oauth_token("U1", "agent-geo") is None + + +async def test_other_refresh_failures_keep_sdk_behaviour(puo, store, monkeypatch): + monkeypatch.setattr(puo, "start_sign_in", lambda *a: pytest.fail("must not prompt")) + provider = await puo.build_provider("U1", "C1", "T1", interactive=False) + client, _ = _client(provider, httpx.Response(503, text="unavailable")) + + # A transient IdP failure is not a dead sign-in: no SignInExpired, token kept. + async with client: + response = await client.post(RS_URL, json={}) + assert response.status_code == 200 + assert await store.get_oauth_token("U1", "agent-geo") is not None + + +async def test_background_sign_in_stores_new_token_with_previous_scopes(puo, store, monkeypatch): + requested_scope = {} + + async def perform_authorization(self): + requested_scope["scope"] = self.context.client_metadata.scope + return httpx.Request("POST", TOKEN_URL, data={"code": "abc"}) + + monkeypatch.setattr(flow._PerUserProvider, "_perform_authorization", perform_authorization) + real_client = httpx.AsyncClient + fresh = {"access_token": "new", "token_type": "Bearer", "expires_in": 600} + fresh["refresh_token"] = "r2" + transport = httpx.MockTransport(lambda r: httpx.Response(200, json=fresh)) + monkeypatch.setattr(flow.httpx, "AsyncClient", lambda **kw: real_client(transport=transport)) + await store.delete_oauth_token("U1", "agent-geo") + + puo.start_sign_in("U1", "C1", "T1", {"agent:agent-geo:read", "offline_access"}) + await asyncio.gather(*puo._sign_ins) + + expected = {"agent:agent-geo:read", "offline_access", "openid"} + assert set(requested_scope["scope"].split()) >= expected + tokens = await puo._token_storage("U1").get_tokens() + assert tokens.access_token == "new" + assert tokens.refresh_token == "r2" + + +def test_tool_result_tells_the_user_to_sign_in_and_stop(): + result = sign_in_expired_tool_result(SignInExpired("agent-geo", prompted=True), tool="geo") + payload = json.loads(result["content"]) + assert result["is_error"] is True + assert payload["code"] == "sign_in_expired" + assert payload["recovery"] == "abort" + assert "sign-in link" in payload["message"] + assert "sign in, then ask again" in payload["message"]