Refactor to eliminate runtime type assertions

Changed tokenStore from UserTokenStore to Storage interface to avoid
runtime type assertions. This makes the code cleaner and type-safe.

Author: dgellow

internal/server/handler.go

@@ -23,7 +23,7 @@ type Server struct {
 	mux            *http.ServeMux
 	config         *config.Config
 	oauthServer    *oauth.Server
-	tokenStore     storage.UserTokenStore
+	storage        storage.Storage
 	sessionManager *client.StdioSessionManager
 	sseServers     map[string]*server.SSEServer // serverName -> SSE server for stdio servers
 }
@@ -143,8 +143,7 @@ func NewServer(ctx context.Context, cfg *config.Config) (*Server, error) {
 			return nil, fmt.Errorf("failed to create OAuth server: %w", err)
 		}
 
-		// Use the storage directly as token store
-		s.tokenStore = store
+		s.storage = store
 
 		// Initialize admin users if admin is enabled
 		if cfg.Proxy.Admin != nil && cfg.Proxy.Admin.Enabled {
@@ -180,7 +179,7 @@ func NewServer(ctx context.Context, cfg *config.Config) (*Server, error) {
 		mux.Handle("/register", chainMiddleware(http.HandlerFunc(s.oauthServer.RegisterHandler), oauthMiddlewares...))
 
 		// Protected endpoints - require authentication
-		tokenHandlers := NewTokenHandlers(s.tokenStore, cfg.MCPServers, s.oauthServer != nil)
+		tokenHandlers := NewTokenHandlers(s.storage, cfg.MCPServers, s.oauthServer != nil)
 		tokenMiddlewares := []MiddlewareFunc{
 			corsMiddleware(allowedOrigins),
 			loggerMiddleware("tokens"),
@@ -280,9 +279,8 @@ func NewServer(ctx context.Context, cfg *config.Config) (*Server, error) {
 					}
 					s.sessionManager.RemoveSession(key)
 
-					// Remove session from storage
-					if store, ok := handler.h.tokenStore.(storage.Storage); ok {
-						if err := store.RevokeSession(sessionCtx, session.SessionID()); err != nil {
+					if handler.h.storage != nil {
+						if err := handler.h.storage.RevokeSession(sessionCtx, session.SessionID()); err != nil {
 							internal.LogWarnWithFields("server", "Failed to revoke session from storage", map[string]interface{}{
 								"error":     err.Error(),
 								"sessionID": session.SessionID(),
@@ -321,7 +319,7 @@ func NewServer(ctx context.Context, cfg *config.Config) (*Server, error) {
 		handler := NewMCPHandler(
 			serverName,
 			serverConfig,
-			s.tokenStore,
+			s.storage,
 			baseURL.String(),
 			info,
 			s.sessionManager,
@@ -371,12 +369,12 @@ func NewServer(ctx context.Context, cfg *config.Config) (*Server, error) {
 			encryptionKey = oauthAuth.EncryptionKey
 		}
 
-		adminHandlers := NewAdminHandlers(s.tokenStore.(storage.Storage), cfg, s.sessionManager, encryptionKey)
+		adminHandlers := NewAdminHandlers(s.storage, cfg, s.sessionManager, encryptionKey)
 		adminMiddlewares := []MiddlewareFunc{
 			corsMiddleware(allowedOrigins),
 			loggerMiddleware("admin"),
-			s.oauthServer.SSOMiddleware(),                                    // Browser SSO
-			adminMiddleware(cfg.Proxy.Admin, s.tokenStore.(storage.Storage)), // Admin check
+			s.oauthServer.SSOMiddleware(),               // Browser SSO
+			adminMiddleware(cfg.Proxy.Admin, s.storage), // Admin check
 		}
 
 		// Admin routes - all protected by admin middleware
@@ -478,7 +476,7 @@ func handleSessionRegistration(
 		handler.mcpServer,
 		handler.userEmail,
 		handler.config.RequiresUserToken,
-		handler.h.tokenStore,
+		handler.h.storage,
 		handler.h.serverName,
 		handler.h.setupBaseURL,
 		handler.config.TokenSetup,
@@ -495,15 +493,15 @@ func handleSessionRegistration(
 	}
 
 	if handler.userEmail != "" {
-		if store, ok := handler.h.tokenStore.(storage.Storage); ok {
+		if handler.h.storage != nil {
 			activeSession := storage.ActiveSession{
 				SessionID:  session.SessionID(),
 				UserEmail:  handler.userEmail,
 				ServerName: handler.h.serverName,
 				Created:    time.Now(),
 				LastActive: time.Now(),
 			}
-			if err := store.TrackSession(sessionCtx, activeSession); err != nil {
+			if err := handler.h.storage.TrackSession(sessionCtx, activeSession); err != nil {
 				internal.LogWarnWithFields("server", "Failed to track session", map[string]interface{}{
 					"error":     err.Error(),
 					"sessionID": session.SessionID(),

internal/server/mcp_handler.go

@@ -31,7 +31,7 @@ type SessionManager interface {
 type MCPHandler struct {
 	serverName      string
 	serverConfig    *config.MCPClientConfig
-	tokenStore      storage.UserTokenStore
+	storage         storage.Storage
 	setupBaseURL    string
 	info            mcp.Implementation
 	sessionManager  SessionManager
@@ -42,7 +42,7 @@ type MCPHandler struct {
 func NewMCPHandler(
 	serverName string,
 	serverConfig *config.MCPClientConfig,
-	tokenStore storage.UserTokenStore,
+	storage storage.Storage,
 	setupBaseURL string,
 	info mcp.Implementation,
 	sessionManager SessionManager,
@@ -51,7 +51,7 @@ func NewMCPHandler(
 	return &MCPHandler{
 		serverName:      serverName,
 		serverConfig:    serverConfig,
-		tokenStore:      tokenStore,
+		storage:         storage,
 		setupBaseURL:    setupBaseURL,
 		info:            info,
 		sessionManager:  sessionManager,
@@ -136,8 +136,8 @@ func (h *MCPHandler) isMessageRequest(r *http.Request) bool {
 // trackUserAccess tracks user access if user email is provided
 func (h *MCPHandler) trackUserAccess(ctx context.Context, userEmail string) {
 	if userEmail != "" {
-		if store, ok := h.tokenStore.(storage.Storage); ok {
-			if err := store.UpsertUser(ctx, userEmail); err != nil {
+		if h.storage != nil {
+			if err := h.storage.UpsertUser(ctx, userEmail); err != nil {
 				internal.LogWarnWithFields("mcp", "Failed to track user", map[string]any{
 					"error": err.Error(),
 					"user":  userEmail,
@@ -263,7 +263,11 @@ func (h *MCPHandler) getUserTokenIfAvailable(ctx context.Context, userEmail stri
 		return "", fmt.Errorf("authentication required")
 	}
 
-	token, err := h.tokenStore.GetUserToken(ctx, userEmail, h.serverName)
+	if h.storage == nil {
+		return "", fmt.Errorf("storage not configured")
+	}
+
+	token, err := h.storage.GetUserToken(ctx, userEmail, h.serverName)
 	if err != nil {
 		return "", err
 	}

internal/server/mcp_handler_test.go

@@ -15,7 +15,7 @@ import (
 	"github.com/dgellow/mcp-front/internal/client"
 	"github.com/dgellow/mcp-front/internal/config"
 	"github.com/dgellow/mcp-front/internal/jsonrpc"
-	"github.com/dgellow/mcp-front/internal/testutil"
+	"github.com/dgellow/mcp-front/internal/storage"
 	"github.com/mark3labs/mcp-go/mcp"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -26,6 +26,28 @@ type mockSessionManager struct {
 	mock.Mock
 }
 
+type mockStorage struct {
+	*storage.MemoryStorage
+	mock.Mock
+}
+
+// Override only the methods we want to mock
+func (m *mockStorage) GetUserToken(ctx context.Context, userEmail, service string) (string, error) {
+	if m.Mock.ExpectedCalls != nil {
+		args := m.Called(ctx, userEmail, service)
+		return args.String(0), args.Error(1)
+	}
+	return m.MemoryStorage.GetUserToken(ctx, userEmail, service)
+}
+
+func (m *mockStorage) UpsertUser(ctx context.Context, email string) error {
+	if m.Mock.ExpectedCalls != nil {
+		args := m.Called(ctx, email)
+		return args.Error(0)
+	}
+	return m.MemoryStorage.UpsertUser(ctx, email)
+}
+
 func (m *mockSessionManager) GetSession(key client.SessionKey) (*client.StdioSession, bool) {
 	args := m.Called(key)
 	if args.Get(0) == nil {
@@ -52,14 +74,16 @@ func (m *mockSessionManager) Shutdown() {
 
 // Test helper to create MCPHandler for SSE tests
 func createTestMCPHandler(serverName string, config *config.MCPClientConfig) *MCPHandler {
-	tokenStore := new(testutil.MockUserTokenStore)
+	mockStore := &mockStorage{
+		MemoryStorage: storage.NewMemoryStorage(),
+	}
 	sessionManager := new(mockSessionManager)
 	info := mcp.Implementation{Name: "test", Version: "1.0"}
 
 	return NewMCPHandler(
 		serverName,
 		config,
-		tokenStore,
+		mockStore,
 		"http://localhost:8080",
 		info,
 		sessionManager,