Refresh upstream tokens transparently instead of forcing re-auth

The upstreamswap middleware returned 401 when upstream access tokens
expired, forcing users through full re-authentication even though
valid refresh tokens existed in storage. This happened because:

1. Redis/memory storage TTL was set to access token expiry, deleting
   the entry (and refresh token) when the access token expired
2. Storage returned nil on ErrExpired, discarding the refresh token
3. The middleware had no refresh path — only 401

Fix all three layers:

- Add DefaultRefreshTokenTTL (30 days) to storage entry TTL so
  refresh tokens survive past access token expiry
- Return token data alongside ErrExpired from storage so callers
  can use the refresh token
- Add UpstreamTokenRefresher interface and implementation that wraps
  the upstream OAuth2Provider and storage
- Plumb the refresher through Server → EmbeddedAuthServer → Runner →
  MiddlewareRunner
- Update upstreamswap middleware to attempt refresh before returning
  401, only requiring re-auth when the refresh token itself fails

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: aron-muon

pkg/auth/upstreamswap/middleware.go

@@ -6,6 +6,7 @@
 package upstreamswap
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -48,6 +49,10 @@ type MiddlewareParams struct {
 // This allows lazy access to the storage, which may not be available at middleware creation time.
 type StorageGetter func() storage.UpstreamTokenStorage
 
+// RefresherGetter is a function that returns an upstream token refresher.
+// This allows lazy access to the refresher, which may not be available at middleware creation time.
+type RefresherGetter func() storage.UpstreamTokenRefresher
+
 // Middleware wraps the upstream swap middleware functionality.
 type Middleware struct {
 	middleware types.MiddlewareFunction
@@ -81,12 +86,13 @@ func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRun
 		return fmt.Errorf("invalid upstream swap configuration: %w", err)
 	}
 
-	// Get storage getter from runner.
-	// The storage getter is a lazy accessor that checks storage availability at request time,
-	// so it's always non-nil. Actual storage availability is verified when processing requests.
+	// Get storage getter and refresher getter from runner.
+	// These are lazy accessors that check availability at request time,
+	// so they're always non-nil. Actual availability is verified when processing requests.
 	storageGetter := runner.GetUpstreamTokenStorage()
+	refresherGetter := runner.GetUpstreamTokenRefresher()
 
-	middleware := createMiddlewareFunc(cfg, storageGetter)
+	middleware := createMiddlewareFunc(cfg, storageGetter, refresherGetter)
 
 	upstreamSwapMw := &Middleware{
 		middleware: middleware,
@@ -141,7 +147,7 @@ func createCustomInjector(headerName string) injectionFunc {
 }
 
 // createMiddlewareFunc creates the actual middleware function.
-func createMiddlewareFunc(cfg *Config, storageGetter StorageGetter) types.MiddlewareFunction {
+func createMiddlewareFunc(cfg *Config, storageGetter StorageGetter, refresherGetter RefresherGetter) types.MiddlewareFunction {
 	// Determine injection strategy at startup time
 	strategy := cfg.HeaderStrategy
 	if strategy == "" {
@@ -188,35 +194,21 @@ func createMiddlewareFunc(cfg *Config, storageGetter StorageGetter) types.Middle
 				return
 			}
 
-			// 4. Lookup upstream tokens
-			tokens, err := stor.GetUpstreamTokens(r.Context(), tsid)
+			// 4. Lookup upstream tokens, refreshing if expired
+			tokens, err := getOrRefreshUpstreamTokens(r.Context(), stor, tsid, refresherGetter)
 			if err != nil {
 				slog.Warn("Failed to get upstream tokens",
 					"middleware", "upstreamswap", "error", err)
-				// Token is expired, was not found, or failed binding validation
-				// (e.g., subject/client mismatch). All three are client-attributable
-				// errors that require the caller to re-authenticate with the upstream IdP.
 				if errors.Is(err, storage.ErrExpired) ||
 					errors.Is(err, storage.ErrNotFound) ||
 					errors.Is(err, storage.ErrInvalidBinding) {
 					writeUpstreamAuthRequired(w)
 					return
 				}
-				// Other storage errors: fail closed to avoid bypassing the token swap
 				http.Error(w, "authentication service temporarily unavailable", http.StatusServiceUnavailable)
 				return
 			}
 
-			// 5. Check if expired
-			// Defense in depth: some storage implementations may return tokens
-			// without checking expiry (the interface does not require it).
-			if tokens.IsExpired(time.Now()) {
-				slog.Warn("Upstream tokens expired",
-					"middleware", "upstreamswap")
-				writeUpstreamAuthRequired(w)
-				return
-			}
-
 			// 6. Inject access token
 			if tokens.AccessToken == "" {
 				slog.Warn("Access token is empty",
@@ -233,3 +225,71 @@ func createMiddlewareFunc(cfg *Config, storageGetter StorageGetter) types.Middle
 		})
 	}
 }
+
+// getOrRefreshUpstreamTokens retrieves upstream tokens from storage, automatically
+// refreshing them if expired and a refresh token is available.
+func getOrRefreshUpstreamTokens(
+	ctx context.Context,
+	stor storage.UpstreamTokenStorage,
+	sessionID string,
+	refresherGetter RefresherGetter,
+) (*storage.UpstreamTokens, error) {
+	tokens, err := stor.GetUpstreamTokens(ctx, sessionID)
+	if err != nil {
+		// ErrExpired returns tokens (including refresh token) alongside the error.
+		// Attempt a refresh before giving up.
+		if errors.Is(err, storage.ErrExpired) && tokens != nil {
+			if refreshed := tryRefreshUpstreamTokens(ctx, sessionID, tokens, refresherGetter); refreshed != nil {
+				return refreshed, nil
+			}
+		}
+		return nil, err
+	}
+
+	// Defense in depth: some storage implementations may return tokens
+	// without checking expiry (the interface does not require it).
+	if tokens.IsExpired(time.Now()) {
+		if refreshed := tryRefreshUpstreamTokens(ctx, sessionID, tokens, refresherGetter); refreshed != nil {
+			return refreshed, nil
+		}
+		return nil, storage.ErrExpired
+	}
+
+	return tokens, nil
+}
+
+// tryRefreshUpstreamTokens attempts to refresh expired upstream tokens using the
+// configured refresher. Returns the refreshed tokens on success, or nil on failure.
+func tryRefreshUpstreamTokens(
+	ctx context.Context,
+	sessionID string,
+	expired *storage.UpstreamTokens,
+	refresherGetter RefresherGetter,
+) *storage.UpstreamTokens {
+	if expired.RefreshToken == "" {
+		slog.Debug("No refresh token available, cannot refresh upstream tokens",
+			"middleware", "upstreamswap")
+		return nil
+	}
+
+	if refresherGetter == nil {
+		return nil
+	}
+	refresher := refresherGetter()
+	if refresher == nil {
+		slog.Debug("Token refresher unavailable, cannot refresh upstream tokens",
+			"middleware", "upstreamswap")
+		return nil
+	}
+
+	refreshed, err := refresher.RefreshAndStore(ctx, sessionID, expired)
+	if err != nil {
+		slog.Warn("Upstream token refresh failed",
+			"middleware", "upstreamswap", "error", err)
+		return nil
+	}
+
+	slog.Debug("Successfully refreshed upstream tokens",
+		"middleware", "upstreamswap")
+	return refreshed
+}

pkg/auth/upstreamswap/middleware_test.go

@@ -99,7 +99,7 @@ func TestMiddleware_NoIdentity(t *testing.T) {
 	}
 
 	cfg := &Config{}
-	middleware := createMiddlewareFunc(cfg, storageGetter)
+	middleware := createMiddlewareFunc(cfg, storageGetter, nil)
 
 	var nextCalled bool
 	nextHandler := http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
@@ -131,7 +131,7 @@ func TestMiddleware_NoTsidClaim(t *testing.T) {
 	}
 
 	cfg := &Config{}
-	middleware := createMiddlewareFunc(cfg, storageGetter)
+	middleware := createMiddlewareFunc(cfg, storageGetter, nil)
 
 	var nextCalled bool
 	nextHandler := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
@@ -166,7 +166,7 @@ func TestMiddleware_StorageUnavailable(t *testing.T) {
 	}
 
 	cfg := &Config{}
-	middleware := createMiddlewareFunc(cfg, storageGetter)
+	middleware := createMiddlewareFunc(cfg, storageGetter, nil)
 
 	var nextCalled bool
 	nextHandler := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
@@ -222,7 +222,7 @@ func TestMiddleware_ClientAttributableStorageErrors_Returns401(t *testing.T) {
 			}
 
 			cfg := &Config{}
-			middleware := createMiddlewareFunc(cfg, storageGetter)
+			middleware := createMiddlewareFunc(cfg, storageGetter, nil)
 
 			var nextCalled bool
 			nextHandler := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
@@ -268,7 +268,7 @@ func TestMiddleware_StorageError(t *testing.T) {
 	}
 
 	cfg := &Config{}
-	middleware := createMiddlewareFunc(cfg, storageGetter)
+	middleware := createMiddlewareFunc(cfg, storageGetter, nil)
 
 	var nextCalled bool
 	nextHandler := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
@@ -319,7 +319,7 @@ func TestMiddleware_SuccessfulSwap_AccessToken(t *testing.T) {
 	cfg := &Config{
 		HeaderStrategy: HeaderStrategyReplace,
 	}
-	middleware := createMiddlewareFunc(cfg, storageGetter)
+	middleware := createMiddlewareFunc(cfg, storageGetter, nil)
 
 	var capturedAuthHeader string
 	nextHandler := http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
@@ -369,7 +369,7 @@ func TestMiddleware_CustomHeader(t *testing.T) {
 		HeaderStrategy:   HeaderStrategyCustom,
 		CustomHeaderName: "X-Upstream-Token",
 	}
-	middleware := createMiddlewareFunc(cfg, storageGetter)
+	middleware := createMiddlewareFunc(cfg, storageGetter, nil)
 
 	var capturedCustomHeader string
 	var capturedAuthHeader string
@@ -422,7 +422,7 @@ func TestMiddleware_ExpiredTokens_Returns401(t *testing.T) {
 	}
 
 	cfg := &Config{}
-	middleware := createMiddlewareFunc(cfg, storageGetter)
+	middleware := createMiddlewareFunc(cfg, storageGetter, nil)
 
 	var nextCalled bool
 	nextHandler := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
@@ -473,7 +473,7 @@ func TestMiddleware_EmptySelectedToken(t *testing.T) {
 	}
 
 	cfg := &Config{}
-	middleware := createMiddlewareFunc(cfg, storageGetter)
+	middleware := createMiddlewareFunc(cfg, storageGetter, nil)
 
 	var nextCalled bool
 	var capturedAuthHeader string
@@ -563,7 +563,7 @@ func TestMiddlewareWithContext(t *testing.T) {
 	}
 
 	cfg := &Config{}
-	middleware := createMiddlewareFunc(cfg, storageGetter)
+	middleware := createMiddlewareFunc(cfg, storageGetter, nil)
 
 	// Test that context is properly passed through
 	var receivedCtx context.Context
@@ -677,6 +677,9 @@ func TestCreateMiddleware(t *testing.T) {
 				mockRunner.EXPECT().GetUpstreamTokenStorage().Return(func() storage.UpstreamTokenStorage {
 					return nil // Storage availability is checked at request time
 				})
+				mockRunner.EXPECT().GetUpstreamTokenRefresher().Return(func() storage.UpstreamTokenRefresher {
+					return nil // Refresher availability is checked at request time
+				})
 				mockRunner.EXPECT().AddMiddleware(gomock.Any(), gomock.Any()).Do(func(_ string, mw types.Middleware) {
 					_, ok := mw.(*Middleware)
 					assert.True(t, ok, "Expected middleware to be of type *upstreamswap.Middleware")
@@ -738,7 +741,7 @@ func TestMiddleware_TsidClaimWrongType(t *testing.T) {
 	}
 
 	cfg := &Config{}
-	middleware := createMiddlewareFunc(cfg, storageGetter)
+	middleware := createMiddlewareFunc(cfg, storageGetter, nil)
 
 	var nextCalled bool
 	nextHandler := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {

pkg/authserver/refresher.go

@@ -0,0 +1,83 @@
+// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package authserver
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"log/slog"
+
+	"github.com/stacklok/toolhive/pkg/authserver/storage"
+	"github.com/stacklok/toolhive/pkg/authserver/upstream"
+)
+
+// upstreamTokenRefresher implements storage.UpstreamTokenRefresher by wrapping
+// an upstream OAuth2Provider (for token refresh) and UpstreamTokenStorage (for
+// persisting the refreshed tokens).
+type upstreamTokenRefresher struct {
+	provider upstream.OAuth2Provider
+	storage  storage.UpstreamTokenStorage
+}
+
+// RefreshAndStore refreshes expired upstream tokens using the stored refresh token,
+// persists the new tokens, and returns them.
+func (r *upstreamTokenRefresher) RefreshAndStore(
+	ctx context.Context,
+	sessionID string,
+	expired *storage.UpstreamTokens,
+) (*storage.UpstreamTokens, error) {
+	if expired == nil {
+		return nil, errors.New("expired tokens are required")
+	}
+	if expired.RefreshToken == "" {
+		return nil, errors.New("no refresh token available for upstream token refresh")
+	}
+
+	slog.Debug("attempting upstream token refresh",
+		"session_id", sessionID,
+		"provider_id", expired.ProviderID,
+	)
+
+	// Refresh tokens via the upstream provider
+	newTokens, err := r.provider.RefreshTokens(ctx, expired.RefreshToken, expired.UpstreamSubject)
+	if err != nil {
+		return nil, fmt.Errorf("upstream token refresh failed: %w", err)
+	}
+
+	// Build updated storage tokens preserving binding fields from the original
+	updated := &storage.UpstreamTokens{
+		ProviderID:      expired.ProviderID,
+		AccessToken:     newTokens.AccessToken,
+		RefreshToken:    newTokens.RefreshToken,
+		IDToken:         newTokens.IDToken,
+		ExpiresAt:       newTokens.ExpiresAt,
+		UserID:          expired.UserID,
+		UpstreamSubject: expired.UpstreamSubject,
+		ClientID:        expired.ClientID,
+	}
+
+	// If the provider didn't rotate the refresh token, keep the original
+	if updated.RefreshToken == "" {
+		updated.RefreshToken = expired.RefreshToken
+	}
+
+	// Store the refreshed tokens
+	if err := r.storage.StoreUpstreamTokens(ctx, sessionID, updated); err != nil {
+		// Log but still return the refreshed tokens — the current request can
+		// proceed even if storage fails. The next request will retry the refresh.
+		slog.Warn("failed to store refreshed upstream tokens",
+			"session_id", sessionID,
+			"error", err,
+		)
+		return updated, nil
+	}
+
+	slog.Debug("upstream tokens refreshed successfully",
+		"session_id", sessionID,
+		"provider_id", expired.ProviderID,
+	)
+
+	return updated, nil
+}

pkg/authserver/runner/embeddedauthserver.go

@@ -135,6 +135,12 @@ func (e *EmbeddedAuthServer) IDPTokenStorage() storage.UpstreamTokenStorage {
 	return e.server.IDPTokenStorage()
 }
 
+// UpstreamTokenRefresher returns a refresher that can refresh expired upstream
+// tokens using the upstream provider's refresh token grant.
+func (e *EmbeddedAuthServer) UpstreamTokenRefresher() storage.UpstreamTokenRefresher {
+	return e.server.UpstreamTokenRefresher()
+}
+
 // createKeyProvider creates a KeyProvider from SigningKeyRunConfig.
 // Returns a GeneratingProvider if config is nil or empty (development mode).
 func createKeyProvider(cfg *authserver.SigningKeyRunConfig) (keys.KeyProvider, error) {

pkg/authserver/server.go

@@ -31,6 +31,11 @@ type Server interface {
 	// Returns nil if no upstream IDP is configured.
 	IDPTokenStorage() storage.UpstreamTokenStorage
 
+	// UpstreamTokenRefresher returns a refresher that can refresh expired upstream
+	// tokens using the upstream provider's refresh token grant.
+	// Returns nil if no upstream IDP is configured.
+	UpstreamTokenRefresher() storage.UpstreamTokenRefresher
+
 	// Close releases resources held by the server.
 	Close() error
 }

pkg/authserver/server_impl.go

@@ -161,6 +161,18 @@ func (s *server) IDPTokenStorage() storage.UpstreamTokenStorage {
 	return s.storage
 }
 
+// UpstreamTokenRefresher returns a refresher that wraps the upstream provider
+// and storage to transparently refresh expired upstream tokens.
+func (s *server) UpstreamTokenRefresher() storage.UpstreamTokenRefresher {
+	if s.upstreamIDP == nil {
+		return nil
+	}
+	return &upstreamTokenRefresher{
+		provider: s.upstreamIDP,
+		storage:  s.storage,
+	}
+}
+
 // Close releases resources held by the server.
 func (s *server) Close() error {
 	slog.Debug("closing OAuth authorization server")