--- a/states.py	Mon May 23 15:02:16 2011 +0200
+++ b/states.py	Mon May 23 15:41:20 2011 +0200
@@ -29,6 +29,7 @@
 from mercurial.node import nullid
 from mercurial import discovery
 from mercurial import extensions
+from mercurial import wireproto
 
 
 _NOPULLPUSH=2
@@ -106,6 +107,19 @@
     extensions.wrapfunction(discovery, 'findcommonoutgoing', filterprivateout)
     extensions.wrapfunction(discovery, 'findcommonincoming', filterprivatein)
 
+    # Write protocols
+    ####################
+    def heads(repo, proto):
+        h = repo._publicheads
+        return wireproto.encodelist(h) + "\n"
+
+    def _reducehead(wirerepo, heads):
+        """heads filtering is done repo side"""
+        return heads
+
+    wireproto.wirerepository._reducehead = _reducehead
+    wireproto.commands['heads'] = (heads, '')
+
 def reposetup(ui, repo):
 
     if not repo.local():