feat: add opt-in shared OAuth token cache via object store

Add oauth_token_store to FastApiFrontEndConfig (optional, unset by default).
When set, the OAuth token obtained during the authorization-code flow is
persisted in the specified object store (e.g. Redis) so that a WebSocket
reconnecting to a different pod does not repeat authentication. When unset,
the existing in-memory dict is used unchanged.

Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>

Author: thepatrickchin

packages/nvidia_nat_core/src/nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import asyncio
+import json
 import logging
 import secrets
 import time
@@ -32,7 +33,10 @@
 from nat.data_models.authentication import AuthFlowType
 from nat.data_models.authentication import AuthProviderBaseConfig
 from nat.data_models.interactive import _HumanPromptOAuthConsent
+from nat.data_models.object_store import NoSuchKeyError
 from nat.front_ends.fastapi.message_handler import WebSocketMessageHandler
+from nat.object_store.interfaces import ObjectStore
+from nat.object_store.models import ObjectStoreItem
 
 logger = logging.getLogger(__name__)
 
@@ -55,15 +59,15 @@ def __init__(self,
                  web_socket_message_handler: WebSocketMessageHandler,
                  auth_timeout_seconds: float = 300.0,
                  return_url: str | None = None,
-                 token_store: dict | None = None,
+                 token_store: dict | ObjectStore | None = None,
                  session_id: str | None = None):
 
         self._add_flow_cb: Callable[[str, FlowState], Awaitable[None]] = add_flow_cb
         self._remove_flow_cb: Callable[[str], Awaitable[None]] = remove_flow_cb
         self._web_socket_message_handler: WebSocketMessageHandler = web_socket_message_handler
         self._auth_timeout_seconds: float = auth_timeout_seconds
         self._return_url: str | None = return_url
-        self._token_store: dict | None = token_store
+        self._token_store: dict | ObjectStore | None = token_store
         self._session_id: str | None = session_id
 
     async def authenticate(
@@ -137,7 +141,7 @@ async def _handle_oauth2_auth_code_flow(self, config: OAuth2AuthCodeFlowProvider
             raise ValueError("Redirect-based authentication (use_redirect_auth=True) requires a return URL, "
                              "but none was configured. Pass return_url when constructing the flow handler.")
 
-        cached = self._get_cached_token(config)
+        cached = await self._get_cached_token(config)
         if cached is not None:
             logger.debug("OAuth token cache hit for client_id=%s", config.client_id)
             return cached
@@ -174,7 +178,7 @@ async def _handle_oauth2_auth_code_flow(self, config: OAuth2AuthCodeFlowProvider
                                    metadata={
                                        "expires_at": token.get("expires_at"), "raw_token": token
                                    })
-        self._store_token(config, ctx)
+        await self._store_token(config, ctx)
         return ctx
 
     def _token_cache_key(self, config: OAuth2AuthCodeFlowProviderConfig) -> str | None:
@@ -183,24 +187,73 @@ def _token_cache_key(self, config: OAuth2AuthCodeFlowProviderConfig) -> str | No
             return None
         return f"{self._session_id}:{config.client_id}:{config.token_url}"
 
-    def _get_cached_token(self, config: OAuth2AuthCodeFlowProviderConfig) -> AuthenticatedContext | None:
+    async def _get_cached_token(self, config: OAuth2AuthCodeFlowProviderConfig) -> AuthenticatedContext | None:
         """Return a cached, non-expired token for *config*, or None."""
         key = self._token_cache_key(config)
         if key is None or self._token_store is None:
             return None
-        entry = self._token_store.get(key)
-        if entry is None:
+
+        if isinstance(self._token_store, dict):
+            entry = self._token_store.get(key)
+            if entry is None:
+                return None
+            ctx, expires_at = entry
+            if expires_at is not None and time.time() >= expires_at - 60:
+                del self._token_store[key]
+                return None
+            return ctx
+
+        # ObjectStore path
+        try:
+            item = await self._token_store.get_object(key)
+        except NoSuchKeyError:
+            return None
+        try:
+            payload = json.loads(item.data)
+        except Exception:
+            logger.warning("OAuth token cache entry for key %s is corrupt; evicting", key)
+            try:
+                await self._token_store.delete_object(key)
+            except NoSuchKeyError:
+                pass
             return None
-        ctx, expires_at = entry
+        expires_at = payload.get("expires_at")
         if expires_at is not None and time.time() >= expires_at - 60:
-            del self._token_store[key]
+            try:
+                await self._token_store.delete_object(key)
+            except NoSuchKeyError:
+                pass
             return None
-        return ctx
+        return AuthenticatedContext(
+            headers=payload.get("headers"),
+            query_params=payload.get("query_params"),
+            cookies=payload.get("cookies"),
+            body=payload.get("body"),
+            metadata=payload.get("metadata"),
+        )
 
-    def _store_token(self, config: OAuth2AuthCodeFlowProviderConfig, ctx: AuthenticatedContext) -> None:
+    async def _store_token(self, config: OAuth2AuthCodeFlowProviderConfig, ctx: AuthenticatedContext) -> None:
         """Cache *ctx* for *config* if caching is available."""
         key = self._token_cache_key(config)
         if key is None or self._token_store is None:
             return
         expires_at = ctx.metadata.get("expires_at") if ctx.metadata else None
-        self._token_store[key] = (ctx, expires_at)
+
+        if isinstance(self._token_store, dict):
+            self._token_store[key] = (ctx, expires_at)
+            return
+
+        # ObjectStore path
+        payload = {
+            "headers": dict(ctx.headers) if ctx.headers else None,
+            "query_params": dict(ctx.query_params) if ctx.query_params else None,
+            "cookies": dict(ctx.cookies) if ctx.cookies else None,
+            "body": ctx.body,
+            "metadata": ctx.metadata,
+            "expires_at": expires_at,
+        }
+        try:
+            item = ObjectStoreItem(data=json.dumps(payload).encode(), content_type="application/json")
+            await self._token_store.upsert_object(key, item)
+        except Exception:
+            logger.warning("Failed to store OAuth token in object store for key %s", key, exc_info=True)

packages/nvidia_nat_core/src/nat/front_ends/fastapi/fastapi_front_end_config.py

@@ -319,6 +319,15 @@ class CrossOriginResourceSharing(BaseModel):
             "request to '/static' and files will be served from the object store. The files will be served from the "
             "object store at '/static/{file_name}'."))
 
+    oauth_token_store: ObjectStoreRef | None = Field(
+        default=None,
+        description=(
+            "Object store reference for sharing the OAuth token cache across processes or replicas. "
+            "When set, the access token obtained during the OAuth authorization-code flow is persisted "
+            "in the specified object store (e.g. a Redis store) so that a WebSocket reconnecting to a "
+            "different pod does not need to repeat authentication. "
+            "If not set, tokens are cached in process memory, which is sufficient for single-replica deployments."))
+
     disable_legacy_routes: bool = Field(
         default=False,
         description="Disable the legacy routes for the FastAPI app. If True, the legacy routes are disabled.")

packages/nvidia_nat_core/src/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py

@@ -21,6 +21,10 @@
 from collections.abc import Awaitable
 from collections.abc import Callable
 from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from nat.object_store.interfaces import ObjectStore
 
 from fastapi import FastAPI
 from fastapi import Request
@@ -203,8 +207,9 @@ def __init__(self, config: Config):
         self._conversation_handlers: dict[str, WebSocketMessageHandler] = {}
 
         # OAuth token cache: maps "{session_id}:{client_id}:{token_url}" to
-        # (AuthenticatedContext, expires_at_unix_or_None)
-        self._oauth_token_store: dict[str, tuple] = {}
+        # (AuthenticatedContext, expires_at_unix_or_None) when using the default in-memory dict.
+        # Replaced with an ObjectStore-backed instance in add_routes() when oauth_token_store is set in config.
+        self._oauth_token_store: dict[str, tuple] | ObjectStore = {}
 
         # Track session managers for each route
         self._session_managers: list[SessionManager] = []
@@ -331,6 +336,10 @@ async def configure(self, app: FastAPI, builder: WorkflowBuilder):
 
     async def add_routes(self, app: FastAPI, builder: WorkflowBuilder):
 
+        if self.front_end_config.oauth_token_store:
+            self._oauth_token_store = await builder.get_object_store_client(self.front_end_config.oauth_token_store)
+            logger.info("Using object-store-backed OAuth token cache: %s", self.front_end_config.oauth_token_store)
+
         session_manager = await self._create_session_manager(builder)
 
         await add_authorization_route(self, app)

packages/nvidia_nat_core/tests/nat/front_ends/auth_flow_handlers/test_websocket_flow_handler.py

@@ -30,6 +30,7 @@
 from nat.data_models.config import Config
 from nat.front_ends.fastapi.auth_flow_handlers.websocket_flow_handler import WebSocketAuthenticationFlowHandler
 from nat.front_ends.fastapi.fastapi_front_end_plugin_worker import FastApiFrontEndPluginWorker
+from nat.object_store.in_memory_object_store import InMemoryObjectStore
 from nat.test.functions import EchoFunctionConfig
 
 
@@ -396,11 +397,11 @@ def test_token_cache_key_format(noop_handler, minimal_oauth_config):
     assert key == f"sess-1:{minimal_oauth_config.client_id}:{minimal_oauth_config.token_url}"
 
 
-def test_get_cached_token_miss(noop_handler, minimal_oauth_config):
+async def test_get_cached_token_miss(noop_handler, minimal_oauth_config):
     """_get_cached_token returns None when the cache has no entry for the config."""
     noop_handler._token_store = {}
     noop_handler._session_id = "sess-1"
-    assert noop_handler._get_cached_token(minimal_oauth_config) is None
+    assert await noop_handler._get_cached_token(minimal_oauth_config) is None
 
 
 @pytest.mark.parametrize("expires_at,expect_hit",
@@ -410,37 +411,135 @@ def test_get_cached_token_miss(noop_handler, minimal_oauth_config):
                              pytest.param(time.time() - 1, False, id="past"),
                              pytest.param(time.time() + 30, False, id="within_buffer"),
                          ])
-def test_get_cached_token_expiry(noop_handler, minimal_oauth_config, expires_at, expect_hit):
+async def test_get_cached_token_expiry(noop_handler, minimal_oauth_config, expires_at, expect_hit):
     """_get_cached_token returns the context when valid and evicts it when expired or within the 60s buffer."""
     ctx = AuthenticatedContext(headers={"Authorization": "Bearer tok"}, metadata={})
     store: dict = {}
     noop_handler._token_store = store
     noop_handler._session_id = "sess-1"
     key = noop_handler._token_cache_key(minimal_oauth_config)
     store[key] = (ctx, expires_at)
-    result = noop_handler._get_cached_token(minimal_oauth_config)
+    result = await noop_handler._get_cached_token(minimal_oauth_config)
     if expect_hit:
         assert result is ctx
     else:
         assert result is None
         assert key not in store
 
 
-def test_store_token_writes_correctly(noop_handler, minimal_oauth_config):
+async def test_store_token_writes_correctly(noop_handler, minimal_oauth_config):
     """_store_token writes (ctx, expires_at) to the store under the expected key."""
     store: dict = {}
     noop_handler._token_store = store
     noop_handler._session_id = "sess-1"
     expires = 9999999999.0
     ctx = AuthenticatedContext(headers={"Authorization": "Bearer tok"}, metadata={"expires_at": expires})
-    noop_handler._store_token(minimal_oauth_config, ctx)
+    await noop_handler._store_token(minimal_oauth_config, ctx)
     key = noop_handler._token_cache_key(minimal_oauth_config)
     assert key in store
     stored_ctx, stored_expires = store[key]
     assert stored_ctx is ctx
     assert stored_expires == expires
 
 
+# --------------------------------------------------------------------------- #
+# ObjectStore-backed token cache tests                                        #
+# --------------------------------------------------------------------------- #
+
+
+async def test_get_cached_token_miss_object_store(noop_handler, minimal_oauth_config):
+    """_get_cached_token returns None when the ObjectStore has no entry for the config."""
+    noop_handler._token_store = InMemoryObjectStore()
+    noop_handler._session_id = "sess-1"
+    assert await noop_handler._get_cached_token(minimal_oauth_config) is None
+
+
+async def test_store_and_retrieve_token_object_store(noop_handler, minimal_oauth_config):
+    """_store_token persists to an ObjectStore and _get_cached_token retrieves it."""
+    noop_handler._token_store = InMemoryObjectStore()
+    noop_handler._session_id = "sess-1"
+    expires = time.time() + 3600
+    ctx = AuthenticatedContext(
+        headers={"Authorization": "Bearer obj-tok"},
+        metadata={
+            "expires_at": expires, "raw_token": {
+                "access_token": "obj-tok"
+            }
+        },
+    )
+    await noop_handler._store_token(minimal_oauth_config, ctx)
+    result = await noop_handler._get_cached_token(minimal_oauth_config)
+    assert result is not None
+    assert result.headers["Authorization"] == "Bearer obj-tok"
+
+
+async def test_get_cached_token_evicts_expired_object_store(noop_handler, minimal_oauth_config):
+    """_get_cached_token evicts and returns None for an expired entry in an ObjectStore."""
+    noop_handler._token_store = InMemoryObjectStore()
+    noop_handler._session_id = "sess-1"
+    expired_ctx = AuthenticatedContext(
+        headers={"Authorization": "Bearer old"},
+        metadata={"expires_at": time.time() - 1},
+    )
+    await noop_handler._store_token(minimal_oauth_config, expired_ctx)
+    result = await noop_handler._get_cached_token(minimal_oauth_config)
+    assert result is None
+
+
+@pytest.mark.usefixtures("set_nat_config_file_env_var")
+async def test_second_authenticate_uses_object_store_cache(monkeypatch, mock_server):
+    """After a successful flow the token is cached in an ObjectStore; a second call must not trigger OAuth again."""
+
+    redirect_port = _free_port()
+    mock_server.register_client(
+        client_id="cid",
+        client_secret="secret",
+        redirect_base=f"http://localhost:{redirect_port}",
+    )
+
+    cfg_nat = Config(workflow=EchoFunctionConfig())
+    worker = FastApiFrontEndPluginWorker(cfg_nat)
+    message_count = [0]
+
+    class _DummyWSHandler:
+
+        def set_flow_handler(self, _):
+            return
+
+        async def create_websocket_message(self, msg):
+            message_count[0] += 1
+            await _complete_oauth_redirect(msg.text, mock_server, worker._outstanding_flows)
+
+    object_store = InMemoryObjectStore()
+    ws_handler = _AuthHandler(
+        oauth_server=mock_server,
+        add_flow_cb=worker._add_flow,
+        remove_flow_cb=worker._remove_flow,
+        web_socket_message_handler=_DummyWSHandler(),
+        token_store=object_store,
+        session_id="test-session",
+    )
+
+    cfg_flow = OAuth2AuthCodeFlowProviderConfig(
+        client_id="cid",
+        client_secret="secret",
+        authorization_url="http://testserver/oauth/authorize",
+        token_url="http://testserver/oauth/token",
+        scopes=["read"],
+        use_pkce=True,
+        redirect_uri=f"http://localhost:{redirect_port}/auth/redirect",
+    )
+
+    monkeypatch.setattr("click.echo", lambda *_: None, raising=True)
+
+    ctx1 = await ws_handler.authenticate(cfg_flow, AuthFlowType.OAUTH2_AUTHORIZATION_CODE)
+    assert message_count[0] == 1, "OAuth flow should have run exactly once"
+
+    ctx2 = await ws_handler.authenticate(cfg_flow, AuthFlowType.OAUTH2_AUTHORIZATION_CODE)
+    assert message_count[0] == 1, "Second authenticate() must return from object store cache without triggering OAuth"
+    assert ctx2.headers["Authorization"] == ctx1.headers["Authorization"]
+
+
 # --------------------------------------------------------------------------- #
 # Token-cache integration: second authenticate() returns from cache          #
 # --------------------------------------------------------------------------- #