Refactor OAuth handling and error management in client connection

- Updated ShouldUseOAuth to prioritize header authentication before falling back to OAuth.
- Enhanced getServerToolCount with shorter timeout and improved error classification for UI status updates.
- Implemented detailed error handling in connect methods to provide clearer context for authentication failures.
- Introduced helper functions for error classification to streamline connection logic and improve maintainability.

Author: Dumbris

internal/oauth/config.go

@@ -305,7 +305,7 @@ func parseBaseURL(fullURL string) (string, error) {
 }
 
 // ShouldUseOAuth determines if OAuth should be attempted for a given server
-// This is now based on protocol type only - OAuth detection happens when server returns 401
+// Headers are tried first if configured, then OAuth as fallback on auth errors
 func ShouldUseOAuth(serverConfig *config.ServerConfig) bool {
 	logger := zap.L().Named("oauth")
 
@@ -321,9 +321,16 @@ func ShouldUseOAuth(serverConfig *config.ServerConfig) bool {
 		return false
 	}
 
-	// For HTTP/SSE servers, we should always try OAuth-enabled clients
-	// The mcp-go library will handle the OAuth flow only if the server requires it (returns 401)
-	logger.Debug("OAuth-enabled client will be used",
+	// If headers are configured, try headers first, not OAuth
+	if len(serverConfig.Headers) > 0 {
+		logger.Debug("Headers configured - will try headers first, OAuth as fallback if needed",
+			zap.String("server", serverConfig.Name),
+			zap.Int("header_count", len(serverConfig.Headers)))
+		return false
+	}
+
+	// For HTTP/SSE servers without headers, try OAuth-enabled clients
+	logger.Debug("No headers configured - OAuth-enabled client will be used",
 		zap.String("server", serverConfig.Name),
 		zap.String("protocol", serverConfig.Protocol))
 

internal/server/server.go

@@ -761,25 +761,60 @@ func (s *Server) QuarantineServer(serverName string, quarantined bool) error {
 }
 
 // getServerToolCount returns the number of tools for a specific server
+// Uses shorter timeout and better error handling for UI status display
 func (s *Server) getServerToolCount(serverID string) int {
 	client, exists := s.upstreamManager.GetClient(serverID)
 	if !exists || !client.IsConnected() {
 		return 0
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	// Use shorter timeout for UI status updates (5 seconds instead of 30)
+	// This reduces waiting time for unresponsive servers
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
+
 	tools, err := client.ListTools(ctx)
 	if err != nil {
-		s.logger.Warn("Failed to get tool count for server",
-			zap.String("server_id", serverID),
-			zap.Error(err))
+		// Classify errors to reduce noise from expected failures
+		if isTimeoutError(err) {
+			// Timeout errors are common for servers that don't support tool listing
+			// Log at debug level to reduce noise
+			s.logger.Debug("Tool listing timeout for server (server may not support tools)",
+				zap.String("server_id", serverID),
+				zap.String("error_type", "timeout"))
+		} else if isConnectionError(err) {
+			// Connection errors suggest the server is actually disconnected
+			s.logger.Debug("Connection error during tool listing",
+				zap.String("server_id", serverID),
+				zap.String("error_type", "connection"))
+		} else {
+			// Other errors might be more significant
+			s.logger.Warn("Failed to get tool count for server",
+				zap.String("server_id", serverID),
+				zap.Error(err))
+		}
 		return 0
 	}
 
 	return len(tools)
 }
 
+// Helper functions for error classification
+func isTimeoutError(err error) bool {
+	errStr := err.Error()
+	return strings.Contains(errStr, "timeout") ||
+		strings.Contains(errStr, "deadline exceeded") ||
+		strings.Contains(errStr, "context canceled")
+}
+
+func isConnectionError(err error) bool {
+	errStr := err.Error()
+	return strings.Contains(errStr, "connection refused") ||
+		strings.Contains(errStr, "no such host") ||
+		strings.Contains(errStr, "connection reset") ||
+		strings.Contains(errStr, "broken pipe")
+}
+
 // StartServer starts the server if it's not already running
 func (s *Server) StartServer(ctx context.Context) error {
 	s.mu.Lock()