lib/newrelic-grape/instrument.rb in newrelic-grape-1.1.0 vs lib/newrelic-grape/instrument.rb in newrelic-grape-1.2.0

@@ -28,11 +28,11 @@
       end
     end
   end
 end
 
-DependencyDetection.defer do 
+DependencyDetection.defer do
   @name = :grape
 
   depends_on do
     defined?(::Grape) && ! ::NewRelic::Control.instance['disable_grape'] && ! ENV['DISABLE_NEW_RELIC_GRAPE']
   end
@@ -40,16 +40,17 @@
   executes do
     NewRelic::Agent.logger.debug 'Installing Grape instrumentation'
   end
 
   executes do
-    ::Grape::Endpoint.class_eval do
-      alias_method :grape_build_middleware, :build_middleware
+    ::Rack::Builder.class_eval do
+      alias_method :origin_use, :use
 
-      def build_middleware
-        builder = grape_build_middleware
-        builder.use ::NewRelic::Agent::Instrumentation::Grape
-        builder
+      def use(middleware, *args, &block)
+        if middleware == Grape::Middleware::Error
+          use ::NewRelic::Agent::Instrumentation::Grape
+        end
+        origin_use(middleware, *args, &block)
       end
     end
   end
 end