--- a/mercurial/wireproto.py	Fri Jul 16 22:20:19 2010 +0200
+++ b/mercurial/wireproto.py	Tue Jul 20 20:53:33 2010 +0200
@@ -133,12 +133,18 @@
 
 # server side
 
+class streamres(object):
+    def __init__(self, gen):
+        self.gen = gen
+
+class pushres(object):
+    def __init__(self, res):
+        self.res = res
+
 def dispatch(repo, proto, command):
     func, spec = commands[command]
     args = proto.getargs(spec)
-    r = func(repo, proto, *args)
-    if r != None:
-        proto.sendresponse(r)
+    return func(repo, proto, *args)
 
 def between(repo, proto, pairs):
     pairs = [decodelist(p, '-') for p in pairs.split(" ")]
@@ -173,13 +179,13 @@
 def changegroup(repo, proto, roots):
     nodes = decodelist(roots)
     cg = repo.changegroup(nodes, 'serve')
-    proto.sendstream(proto.groupchunks(cg))
+    return streamres(proto.groupchunks(cg))
 
 def changegroupsubset(repo, proto, bases, heads):
     bases = decodelist(bases)
     heads = decodelist(heads)
     cg = repo.changegroupsubset(bases, heads, 'serve')
-    proto.sendstream(proto.groupchunks(cg))
+    return streamres(proto.groupchunks(cg))
 
 def heads(repo, proto):
     h = repo.heads()
@@ -215,7 +221,7 @@
     return '%s\n' % int(r)
 
 def stream(repo, proto):
-    proto.sendstream(streamclone.stream_out(repo))
+    return streamres(streamclone.stream_out(repo))
 
 def unbundle(repo, proto, heads):
     their_heads = decodelist(heads)
@@ -259,7 +265,7 @@
                 sys.stderr.write("abort: %s\n" % inst)
         finally:
             lock.release()
-            proto.sendpushresponse(r)
+            return pushres(r)
 
     finally:
         fp.close()