diff --git a/runbot/models/batch.py b/runbot/models/batch.py
index 5f022810..da6749a0 100644
--- a/runbot/models/batch.py
+++ b/runbot/models/batch.py
@@ -335,7 +335,8 @@ class Batch(models.Model):
                 'create_batch_id': self.id,
                 'used_custom_trigger': bool(trigger_custom),
             }
-            params_value['builds_reference_ids'] = trigger._reference_builds(bundle)
+
+            params_value['builds_reference_ids'] = trigger._reference_builds(self)
 
             params = self.env['runbot.build.params'].create(params_value)
 
diff --git a/runbot/models/build_config.py b/runbot/models/build_config.py
index 64ce96ea..10927d0c 100644
--- a/runbot/models/build_config.py
+++ b/runbot/models/build_config.py
@@ -777,9 +777,9 @@ class ConfigStep(models.Model):
 
         return dict(cmd=cmd, container_name=build._get_docker_name(), cpu_limit=self.cpu_limit)
 
-    def _reference_builds(self, bundle, trigger):
+    def _reference_builds(self, batch, trigger):
         upgrade_dumps_trigger_id = trigger.upgrade_dumps_trigger_id
-        refs_batches = self._reference_batches(bundle, trigger)
+        refs_batches = self._reference_batches(batch, trigger)
         refs_builds = refs_batches.mapped('slot_ids').filtered(
             lambda slot: slot.trigger_id == upgrade_dumps_trigger_id
             ).mapped('build_id')
@@ -791,13 +791,14 @@ class ConfigStep(models.Model):
     def _is_upgrade_step(self):
         return self.job_type in ('configure_upgrade', 'configure_upgrade_complement')
 
-    def _reference_batches(self, bundle, trigger):
+    def _reference_batches(self, batch, trigger):
         if self.job_type == 'configure_upgrade_complement':
-            return self._reference_batches_complement(bundle, trigger)
+            return self._reference_batches_complement(batch, trigger)
         else:
-            return self._reference_batches_upgrade(bundle, trigger.upgrade_dumps_trigger_id.category_id.id)
+            return self._reference_batches_upgrade(batch, trigger.upgrade_dumps_trigger_id.category_id.id)
 
-    def _reference_batches_complement(self, bundle, trigger):
+    def _reference_batches_complement(self, batch, trigger):
+        bundle = batch.bundle_id
         category_id = trigger.upgrade_dumps_trigger_id.category_id.id
         version = bundle.version_id
         next_versions = version.next_major_version_id | version.next_intermediate_version_ids  # TODO filter on trigger version
@@ -810,10 +811,11 @@ class ConfigStep(models.Model):
                 if bundle.version_id in upgrade_complement_step._get_upgrade_source_versions(next_version):
                     target_versions |= next_version
         return target_versions.with_context(
-            category_id=category_id, project_id=bundle.project_id.id
+            category_id=category_id, project_id=bundle.project_id.id,
             ).mapped('base_bundle_id').filtered('to_upgrade').mapped('last_done_batch')
 
-    def _reference_batches_upgrade(self, bundle, category_id):
+    def _reference_batches_upgrade(self, batch, category_id):
+        bundle = batch.bundle_id
         target_refs_bundles = self.env['runbot.bundle']
         upgrade_domain = [('to_upgrade', '=', True), ('project_id', '=', bundle.project_id.id)]
         if self.upgrade_to_version_ids:
diff --git a/runbot/models/repo.py b/runbot/models/repo.py
index df5dbd4a..d13d01cd 100644
--- a/runbot/models/repo.py
+++ b/runbot/models/repo.py
@@ -86,12 +86,12 @@ class Trigger(models.Model):
             raise UserError('Upgrade trigger should have a config with step of type Configure Upgrade')
         return upgrade_step
 
-    def _reference_builds(self, bundle):
+    def _reference_builds(self, batch):
         self.ensure_one()
         if self.upgrade_step_id:  # this is an upgrade trigger, add corresponding builds
-            custom_config = next((trigger_custom.config_id for trigger_custom in bundle.trigger_custom_ids if trigger_custom.trigger_id == self), False)
+            custom_config = next((trigger_custom.config_id for trigger_custom in batch.bundle_id.trigger_custom_ids if trigger_custom.trigger_id == self), False)
             step = self._upgrade_step_from_config(custom_config) if custom_config else self.upgrade_step_id
-            refs_builds = step._reference_builds(bundle, self)
+            refs_builds = step._reference_builds(batch, self)
             return [(4, b.id) for b in refs_builds]
         return []