diff --git a/lib/datadog/ci.rb b/lib/datadog/ci.rb
index 7b90cfec..6332e387 100644
--- a/lib/datadog/ci.rb
+++ b/lib/datadog/ci.rb
@@ -12,13 +12,11 @@
 module Datadog
   # Public API for Datadog CI visibility
   module CI
-    module_function
-
-    def trace_test(test_name, service_name: nil, operation_name: nil, tags: {}, &block)
+    def self.trace_test(test_name, service_name: nil, operation_name: "test", tags: {}, &block)
       Recorder.trace_test(test_name, service_name: service_name, operation_name: operation_name, tags: tags, &block)
     end
 
-    def trace(span_type, span_name, &block)
+    def self.trace(span_type, span_name, &block)
       Recorder.trace(span_type, span_name, &block)
     end
   end
diff --git a/lib/datadog/ci/contrib/minitest/hooks.rb b/lib/datadog/ci/contrib/minitest/hooks.rb
index 76c14b02..10cdf993 100644
--- a/lib/datadog/ci/contrib/minitest/hooks.rb
+++ b/lib/datadog/ci/contrib/minitest/hooks.rb
@@ -35,20 +35,19 @@ def before_setup
           end
 
           def after_teardown
-            return super unless @current_test_span
-
-            Thread.current[:_datadog_test_span] = nil
+            test_span = @current_test_span
+            return super unless test_span
 
             case result_code
             when "."
-              @current_test_span.passed!
+              test_span.passed!
             when "E", "F"
-              @current_test_span.failed!(failure)
+              test_span.failed!(failure)
             when "S"
-              @current_test_span.skipped!(nil, failure.message)
+              test_span.skipped!(nil, failure.message)
             end
 
-            @current_test_span.finish
+            test_span.finish
             @current_test_span = nil
 
             super
diff --git a/lib/datadog/ci/recorder.rb b/lib/datadog/ci/recorder.rb
index 28bda982..73b972c9 100644
--- a/lib/datadog/ci/recorder.rb
+++ b/lib/datadog/ci/recorder.rb
@@ -14,7 +14,7 @@ module CI
     # Common behavior for CI tests
     module Recorder
       # Creates a new span for a CI test
-      def self.trace_test(test_name, service_name: nil, operation_name: nil, tags: {})
+      def self.trace_test(test_name, service_name: nil, operation_name: "test", tags: {}, &block)
         span_options = {
           resource: test_name,
           service: service_name,
@@ -23,10 +23,11 @@ def self.trace_test(test_name, service_name: nil, operation_name: nil, tags: {})
 
         tags[:test_name] = test_name
 
-        if block_given?
+        if block
           ::Datadog::Tracing.trace(operation_name, **span_options) do |span, trace|
             set_tags!(trace, span, tags)
-            yield(Test.new(span))
+
+            block.call(Test.new(span))
           end
         else
           span = ::Datadog::Tracing.trace(operation_name, **span_options)
@@ -36,15 +37,15 @@ def self.trace_test(test_name, service_name: nil, operation_name: nil, tags: {})
         end
       end
 
-      def self.trace(span_type, span_name)
+      def self.trace(span_type, span_name, &block)
         span_options = {
           resource: span_name,
           span_type: span_type
         }
 
-        if block_given?
+        if block
           ::Datadog::Tracing.trace(span_name, **span_options) do |tracer_span|
-            yield Span.new(tracer_span)
+            block.call(Span.new(tracer_span))
           end
         else
           tracer_span = Datadog::Tracing.trace(span_name, **span_options)
diff --git a/sig/datadog/ci.rbs b/sig/datadog/ci.rbs
index 75d8c788..0815ae6b 100644
--- a/sig/datadog/ci.rbs
+++ b/sig/datadog/ci.rbs
@@ -1,7 +1,7 @@
 module Datadog
   module CI
-    # See the writing guide of rbs: https://github.com/ruby/rbs#guides
-    class Error < StandardError
-    end
+    def self.trace_test: (String span_name, ?service_name: String?, ?operation_name: String, ?tags: Hash[untyped, untyped]) ?{ (Datadog::CI::Test span) -> untyped } -> untyped
+
+    def self.trace: (String span_type, String span_name) ?{ (Datadog::CI::Span span) -> untyped } -> untyped
   end
 end
diff --git a/sig/datadog/ci/contrib/minitest/hooks.rbs b/sig/datadog/ci/contrib/minitest/hooks.rbs
index 94e1fe62..d376a62a 100644
--- a/sig/datadog/ci/contrib/minitest/hooks.rbs
+++ b/sig/datadog/ci/contrib/minitest/hooks.rbs
@@ -5,6 +5,8 @@ module Datadog
         module Hooks : ::Minitest::Test
           include ::Minitest::Test::LifecycleHooks
 
+          @current_test_span: Datadog::CI::Test?
+
           def before_setup: () -> (nil | untyped)
 
           def after_teardown: () -> untyped
diff --git a/sig/datadog/ci/recorder.rbs b/sig/datadog/ci/recorder.rbs
index 248cfe9b..fe8c0c12 100644
--- a/sig/datadog/ci/recorder.rbs
+++ b/sig/datadog/ci/recorder.rbs
@@ -3,16 +3,13 @@ module Datadog
     module Recorder
       self.@environment_tags: Hash[String, String]
 
-      def self.trace: (untyped span_name, ?::Hash[untyped, untyped] options) ?{ (untyped, untyped) -> untyped } -> untyped
-      def self.set_tags!: (untyped trace, untyped span, ?::Hash[untyped, untyped] tags) -> untyped
+      def self.trace_test: (String span_name, ?service_name: String?, ?operation_name: String, ?tags: Hash[untyped, untyped]) ?{ (Datadog::CI::Test span) -> untyped } -> untyped
 
-      def self.passed!: (untyped span) -> untyped
+      def self.trace: (String span_type, String span_name) ?{ (Datadog::CI::Span span) -> untyped } -> untyped
 
-      def self.failed!: (untyped span, ?untyped? exception) -> untyped
+      def self.set_tags!: (Datadog::Tracing::TraceOperation trace, Datadog::Tracing::SpanOperation span, ?::Hash[untyped, untyped] tags)  -> untyped
 
-      def self.skipped!: (untyped span, ?untyped? exception) -> untyped
-
-      def self.set_environment_runtime_tags!: (untyped span) -> untyped
+      def self.set_environment_runtime_tags!: (Datadog::Tracing::SpanOperation span) -> untyped
     end
   end
 end
diff --git a/sig/datadog/ci/span.rbs b/sig/datadog/ci/span.rbs
new file mode 100644
index 00000000..8fab668f
--- /dev/null
+++ b/sig/datadog/ci/span.rbs
@@ -0,0 +1,19 @@
+module Datadog
+  module CI
+    class Span
+      @tracer_span: Datadog::Tracing::SpanOperation
+
+      attr_reader tracer_span: Datadog::Tracing::SpanOperation
+
+      def initialize: (Datadog::Tracing::SpanOperation tracer_span) -> void
+
+      def passed!: () -> void
+
+      def failed!: (?untyped? exception) -> void
+
+      def skipped!: (?untyped? exception, ?String? reason) -> void
+
+      def finish: () -> untyped
+    end
+  end
+end
diff --git a/sig/datadog/ci/test.rbs b/sig/datadog/ci/test.rbs
new file mode 100644
index 00000000..178981dd
--- /dev/null
+++ b/sig/datadog/ci/test.rbs
@@ -0,0 +1,6 @@
+module Datadog
+  module CI
+    class Test < Span
+    end
+  end
+end
diff --git a/vendor/rbs/ddtrace/0/datadog/tracing.rbs b/vendor/rbs/ddtrace/0/datadog/tracing.rbs
index 778750c1..a195f44b 100644
--- a/vendor/rbs/ddtrace/0/datadog/tracing.rbs
+++ b/vendor/rbs/ddtrace/0/datadog/tracing.rbs
@@ -1,6 +1,6 @@
 module Datadog
   module Tracing
-    def self.active_trace: () -> Datadog::Tracing::TraceSegment
+    def self.active_trace: () -> Datadog::Tracing::TraceOperation
     def self.trace: (String span_name, Hash[untyped, untyped] options) ?{ (untyped span, untyped trace) -> untyped } -> Datadog::Tracing::SpanOperation
   end
 end