--- a/src/topic/__init__.py	Sat Mar 12 18:19:27 2016 +0000
+++ b/src/topic/__init__.py	Sat Mar 12 18:42:16 2016 +0000
@@ -14,11 +14,13 @@
 
 from mercurial.i18n import _
 from mercurial import branchmap
+from mercurial import bundle2
 from mercurial import cmdutil
 from mercurial import commands
 from mercurial import context
 from mercurial import discovery as discoverymod
 from mercurial import error
+from mercurial import exchange
 from mercurial import extensions
 from mercurial import localrepo
 from mercurial import lock
@@ -295,6 +297,11 @@
 extensions.wrapfunction(merge, 'update', mergeupdatewrap)
 extensions.wrapfunction(discoverymod, '_headssummary', discovery._headssummary)
 extensions.wrapfunction(wireproto, 'branchmap', discovery.wireprotobranchmap)
+extensions.wrapfunction(bundle2, 'handlecheckheads', discovery.handlecheckheads)
+bundle2.handlecheckheads.params = frozenset() # we need a proper wrape b2 part stuff
+bundle2.parthandlermapping['check:heads'] = bundle2.handlecheckheads
+extensions.wrapfunction(exchange, '_pushb2phases', discovery._pushb2phases)
+exchange.b2partsgenmapping['phase'] = exchange._pushb2phases
 topicrevset.modsetup()
 cmdutil.summaryhooks.add('topic', summaryhook)