diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 5696edca12..747be2b7ce 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -29,7 +29,7 @@ steps:
             transport: "8.2"
     env:
       RUBY_VERSION: "{{ matrix.ruby }}"
-      STACK_VERSION: 9.0.0-SNAPSHOT
+      STACK_VERSION: 9.1.0-SNAPSHOT
       ES_YAML_TESTS_BRANCH: main
       TRANSPORT_VERSION: "{{ matrix.transport }}"
       RUBY_SOURCE: "{{ matrix.ruby_source }}"
diff --git a/rake_tasks/automation.rake b/rake_tasks/automation.rake
index 1f7dfe77f9..7ccca18e40 100644
--- a/rake_tasks/automation.rake
+++ b/rake_tasks/automation.rake
@@ -107,11 +107,16 @@ namespace :automation do
       if file == '.buildkite/pipeline.yml'
         require 'yaml'
         yaml = YAML.safe_load(content)
-        branch = version.match(/([0-9]+\.[0-9]+)\.[0-9]+.*/)[1]
         yaml_tests_branch = yaml['steps'][0]['env']['ES_YAML_TESTS_BRANCH']
-        next if yaml_tests_branch == 'main'
 
-        content.gsub!(yaml_tests_branch, branch)
+        if yaml_tests_branch == 'main'
+          old = content.match(/STACK_VERSION: (.*)/)[1]
+          new = "STACK_VERSION: #{version}"
+          content.gsub!(new, old)
+        else
+          branch = version.match(/([0-9]+\.[0-9]+)\.[0-9]+.*/)[1]
+          content.gsub!(yaml_tests_branch, branch)
+        end
         puts "[#{yaml_tests_branch}] -> [#{branch}] in #{file.gsub('./', '')}"
       end
       match = content.match(regexp)