Renew a stored MCP sign-in once its access token expires

The MCP SDK knows a token's expiry only when it obtained the token in the
same process. A worker or a restarted API loading tokens from the
connection sent an expired one, got a 401 and asked the owner to sign in
again. The expiry is now saved with the tokens and handed back to the SDK,
which then renews the token with its refresh token first; a token saved
before expiries were kept is renewed once. A sign-in that cannot be
renewed raises MCPReauthorizationRequired.

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:

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