fix: reuse persisted auth tokens across server restarts

src/okta_mcp_server/server.py

@@ -33,15 +33,21 @@ async def okta_authorisation_flow(server: FastMCP) -> AsyncIterator[OktaAppConte
     """
     logger.info("Starting Okta authorization flow")
     manager = OktaAuthManager()
-    await manager.authenticate()
+    # Reuse a cached or refreshable token when one is available; only fall back
+    # to an interactive device grant when no valid token can be obtained.
+    # is_valid_token() silently refreshes via the stored refresh token, so a
+    # restart or reconnect no longer re-prompts for device authorization now
+    # that tokens persist across shutdowns.
+    if not await manager.is_valid_token():
+        await manager.authenticate()
     logger.info("Okta authentication completed successfully")
     prune_tools_by_scope(server, manager)
 
-    try:
-        yield OktaAppContext(okta_auth_manager=manager)
-    finally:
-        logger.debug("Clearing Okta tokens")
-        manager.clear_tokens()
+    # Tokens are intentionally left in the keyring on shutdown so a restart can
+    # reuse the cached access token (and silently refresh it via the refresh
+    # token) without forcing another device-authorization prompt. Clearing only
+    # happens on an explicit logout via OktaAuthManager.clear_tokens().
+    yield OktaAppContext(okta_auth_manager=manager)
 
 
 mcp = FastMCP("Okta IDaaS MCP Server", lifespan=okta_authorisation_flow)

tests/test_lifespan.py

@@ -0,0 +1,79 @@
+# The Okta software accompanied by this notice is provided pursuant to the following terms:
+# Copyright © 2026-Present, Okta, Inc.
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0.
+# Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and limitations under the License.
+
+"""Tests for the server lifespan — token reuse and persistence across shutdowns."""
+
+from __future__ import annotations
+
+from contextlib import suppress
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from okta_mcp_server.server import OktaAppContext, okta_authorisation_flow
+
+
+def _fake_manager(token_valid: bool = True):
+    manager = MagicMock()
+    manager.is_valid_token = AsyncMock(return_value=token_valid)
+    manager.authenticate = AsyncMock()
+    return manager
+
+
+class TestLifespanStartupAuth:
+    """Startup should reuse a cached/refreshable token instead of always re-authenticating."""
+
+    @pytest.mark.asyncio
+    async def test_reuses_token_without_device_prompt_when_valid(self):
+        # A valid (or silently refreshable) token must NOT trigger another device
+        # grant on startup — this is the reconnect-without-re-prompt regression.
+        manager = _fake_manager(token_valid=True)
+        with patch("okta_mcp_server.server.OktaAuthManager", return_value=manager), \
+                patch("okta_mcp_server.server.prune_tools_by_scope") as prune:
+            async with okta_authorisation_flow(MagicMock()) as ctx:
+                assert isinstance(ctx, OktaAppContext)
+                assert ctx.okta_auth_manager is manager
+
+        manager.is_valid_token.assert_awaited_once()
+        manager.authenticate.assert_not_awaited()
+        prune.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_authenticates_when_no_valid_token(self):
+        manager = _fake_manager(token_valid=False)
+        with patch("okta_mcp_server.server.OktaAuthManager", return_value=manager), \
+                patch("okta_mcp_server.server.prune_tools_by_scope"):
+            async with okta_authorisation_flow(MagicMock()):
+                pass
+
+        manager.is_valid_token.assert_awaited_once()
+        manager.authenticate.assert_awaited_once()
+
+
+class TestLifespanTokenPersistence:
+    """The lifespan must not wipe stored tokens on shutdown (only explicit logout should)."""
+
+    @pytest.mark.asyncio
+    async def test_tokens_not_cleared_on_normal_exit(self):
+        manager = _fake_manager(token_valid=True)
+        with patch("okta_mcp_server.server.OktaAuthManager", return_value=manager), \
+                patch("okta_mcp_server.server.prune_tools_by_scope"):
+            async with okta_authorisation_flow(MagicMock()):
+                pass
+
+        manager.clear_tokens.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_tokens_not_cleared_on_exception(self):
+        manager = _fake_manager(token_valid=True)
+        with patch("okta_mcp_server.server.OktaAuthManager", return_value=manager), \
+                patch("okta_mcp_server.server.prune_tools_by_scope"):
+            with suppress(RuntimeError):
+                async with okta_authorisation_flow(MagicMock()):
+                    raise RuntimeError("boom")
+
+        manager.clear_tokens.assert_not_called()