diff --git a/docsgpt/agents/tools/mcp_tool.py b/docsgpt/agents/tools/mcp_tool.py index 6454c7d4..4fee49d6 100644 --- a/docsgpt/agents/tools/mcp_tool.py +++ b/docsgpt/agents/tools/mcp_tool.py @@ -37,6 +37,9 @@ logger = logging.getLogger(__name__) _mcp_clients_cache = {} +# A token expiry long past: a stored token of unknown age is renewed before use. +_EXPIRED = 1.0 + _ANNOTATION_HINTS = ("readOnlyHint", "destructiveHint", "idempotentHint", "openWorldHint") @@ -720,6 +723,10 @@ class MCPTool(Tool): } +class MCPReauthorizationRequired(Exception): + """An MCP sign-in expired and could not be renewed: its owner must sign in again.""" + + class DocsGPTOAuth(OAuthClientProvider): """ Custom OAuth handler for DocsGPT that uses frontend redirect instead of browser. @@ -813,6 +820,25 @@ class DocsGPTOAuth(OAuthClientProvider): return prepare_token_auth + async def _initialize(self) -> None: + """Load the stored sign-in along with when its access token expires. + + The SDK learns a token's expiry only from a token response in this + process, so a token read back from the connection gets the expiry + saved with it, and an expired one is renewed with its refresh token + before it is sent. One saved before expiries were kept is of unknown + age and is renewed once. + """ + await super()._initialize() + tokens = self.context.current_tokens + if tokens is None: + return + expires_at = getattr(self.context.storage, "expires_at", None) + if expires_at: + self.context.token_expiry_time = float(expires_at) + elif tokens.refresh_token and tokens.expires_in: + self.context.token_expiry_time = _EXPIRED + def _process_auth_url(self, authorization_url: str) -> tuple[str, str]: """Process authorization URL to extract state""" try: @@ -920,12 +946,12 @@ class NonInteractiveOAuth(DocsGPTOAuth): super().__init__(**kwargs) async def redirect_handler(self, authorization_url: str) -> None: - raise Exception( + raise MCPReauthorizationRequired( "OAuth session expired — please re-authorize this MCP server in tool settings." ) async def callback_handler(self) -> AuthorizationCodeResult: - raise Exception( + raise MCPReauthorizationRequired( "OAuth session expired — please re-authorize this MCP server in tool settings." ) @@ -951,6 +977,8 @@ class DBTokenStorage(TokenStorage): self.user_id = user_id self.expected_redirect_uri = expected_redirect_uri self.connection_id = connection_id + # When the stored access token expires (epoch seconds), once read. + self.expires_at: Optional[float] = None @staticmethod def get_base_url(url: str) -> str: @@ -972,6 +1000,7 @@ class DBTokenStorage(TokenStorage): data = await asyncio.to_thread(self._fetch_session_data) if not data or "tokens" not in data: return None + self.expires_at = data.get("tokens_expires_at") try: return OAuthToken.model_validate(data["tokens"]) except ValidationError as e: @@ -1007,7 +1036,10 @@ class DBTokenStorage(TokenStorage): async def set_tokens(self, tokens: OAuthToken) -> None: base_url = self.get_base_url(self.server_url) token_dump = tokens.model_dump() - await asyncio.to_thread(self._merge, {"tokens": token_dump}) + # Kept beside the token: a later process cannot tell from + # ``expires_in`` alone whether it has expired. + self.expires_at = time.time() + tokens.expires_in if tokens.expires_in else None + await asyncio.to_thread(self._merge, {"tokens": token_dump, "tokens_expires_at": self.expires_at}) logger.info("Saved tokens for %s", base_url) async def get_client_info(self) -> OAuthClientInformationFull | None: diff --git a/tests/agents/tools/test_mcp_tool.py b/tests/agents/tools/test_mcp_tool.py index 3dac7d1f..fbef348c 100644 --- a/tests/agents/tools/test_mcp_tool.py +++ b/tests/agents/tools/test_mcp_tool.py @@ -1203,6 +1203,128 @@ class TestDBTokenStorage: is None ) + def test_stored_tokens_remember_when_they_expire(self, monkeypatch, pg_conn): + import time + + from mcp.shared.auth import OAuthToken + + from docsgpt.agents.tools.mcp_tool import DBTokenStorage + + self._patch_db(monkeypatch, pg_conn) + storage = DBTokenStorage(server_url="https://mcp.expiry.example.com/mcp", user_id="user-expiry") + loop = asyncio.new_event_loop() + try: + before = time.time() + loop.run_until_complete(storage.set_tokens(OAuthToken( + access_token="at", token_type="Bearer", expires_in=3600, refresh_token="rt", + ))) + # A later process reads the tokens back with when they expire. + reader = DBTokenStorage(server_url="https://mcp.expiry.example.com/mcp", user_id="user-expiry") + tokens = loop.run_until_complete(reader.get_tokens()) + finally: + loop.close() + assert tokens.access_token == "at" + assert before + 3600 <= reader.expires_at <= time.time() + 3600 + + def test_tokens_without_a_lifetime_have_no_expiry(self, monkeypatch, pg_conn): + from mcp.shared.auth import OAuthToken + + from docsgpt.agents.tools.mcp_tool import DBTokenStorage + + self._patch_db(monkeypatch, pg_conn) + storage = DBTokenStorage(server_url="https://mcp.forever.example.com/mcp", user_id="user-forever") + loop = asyncio.new_event_loop() + try: + loop.run_until_complete(storage.set_tokens(OAuthToken(access_token="at", token_type="Bearer"))) + loop.run_until_complete(storage.get_tokens()) + finally: + loop.close() + assert storage.expires_at is None + + +# ===================================================================== +# Renewing a stored sign-in +# ===================================================================== + + +class _StoredTokens: + """A token storage holding one sign-in, with or without its expiry.""" + + def __init__(self, tokens, expires_at=None): + self.tokens = tokens + self.expires_at = None + self._expires_at = expires_at + + async def get_tokens(self): + self.expires_at = self._expires_at + return self.tokens + + async def get_client_info(self): + return None + + +@pytest.mark.unit +class TestStoredSignInRenewal: + """A stored MCP sign-in is renewed with its refresh token once it expires. + + The SDK only knows when a token expires if it obtained it in this + process; a worker that loads it from the connection must be told, or it + sends the expired token, gets a 401 and asks the user to sign in again. + """ + + @staticmethod + def _oauth(storage): + from docsgpt.agents.tools.mcp_tool import NonInteractiveOAuth + + oauth = NonInteractiveOAuth( + mcp_url="https://mcp.example.com/mcp", + redirect_uri="https://docsgpt.example.com/api/mcp_server/callback", + user_id="alice", + connection_id="c1", + ) + oauth.context.storage = storage + return oauth + + @staticmethod + def _token(**extra): + from mcp.shared.auth import OAuthToken + + return OAuthToken(access_token="at", token_type="Bearer", **extra) + + def test_expired_stored_token_is_due_for_renewal(self): + import time + + oauth = self._oauth(_StoredTokens(self._token(refresh_token="rt", expires_in=3600), time.time() - 5)) + asyncio.run(oauth._initialize()) + assert oauth.context.is_token_valid() is False + + def test_unexpired_stored_token_is_used_as_it_is(self): + import time + + expires_at = time.time() + 600 + oauth = self._oauth(_StoredTokens(self._token(refresh_token="rt", expires_in=3600), expires_at)) + asyncio.run(oauth._initialize()) + assert oauth.context.token_expiry_time == expires_at + assert oauth.context.is_token_valid() is True + + def test_token_saved_before_expiries_were_kept_is_renewed_once(self): + # Its age is unknown, so it is treated as expired while it can be renewed. + oauth = self._oauth(_StoredTokens(self._token(refresh_token="rt", expires_in=3600))) + asyncio.run(oauth._initialize()) + assert oauth.context.is_token_valid() is False + + def test_token_that_cannot_be_renewed_is_still_tried(self): + oauth = self._oauth(_StoredTokens(self._token(expires_in=3600))) + asyncio.run(oauth._initialize()) + assert oauth.context.is_token_valid() is True + + def test_a_lost_sign_in_raises_a_typed_error(self): + from docsgpt.agents.tools.mcp_tool import MCPReauthorizationRequired + + oauth = self._oauth(_StoredTokens(None)) + with pytest.raises(MCPReauthorizationRequired, match="OAuth session expired"): + asyncio.run(oauth.redirect_handler("https://mcp.example.com/authorize?state=x")) + # ===================================================================== # NonInteractiveOAuth