Allow cancellation of all capability contexts

capnp/helpers/capabilityHelper.h

@@ -75,10 +75,19 @@ public:
   kj::Own<PyRefCounter> py_server;
   kj::Own<PyRefCounter> kj_loop;
 
+#if (CAPNP_VERSION_MAJOR < 1)
   PythonInterfaceDynamicImpl(capnp::InterfaceSchema & schema,
                              kj::Own<PyRefCounter> _py_server,
                              kj::Own<PyRefCounter> kj_loop)
-    : capnp::DynamicCapability::Server(schema), py_server(kj::mv(_py_server)), kj_loop(kj::mv(kj_loop)) { }
+    : capnp::DynamicCapability::Server(schema),
+      py_server(kj::mv(_py_server)), kj_loop(kj::mv(kj_loop)) { }
+#else
+  PythonInterfaceDynamicImpl(capnp::InterfaceSchema & schema,
+                             kj::Own<PyRefCounter> _py_server,
+                             kj::Own<PyRefCounter> kj_loop)
+    : capnp::DynamicCapability::Server(schema, { true }),
+      py_server(kj::mv(_py_server)), kj_loop(kj::mv(kj_loop)) { }
+#endif
 
   ~PythonInterfaceDynamicImpl() {
   }
@@ -87,6 +96,12 @@ public:
                          capnp::CallContext< capnp::DynamicStruct, capnp::DynamicStruct> context);
 };
 
+inline void allowCancellation(capnp::CallContext<capnp::DynamicStruct, capnp::DynamicStruct> context) {
+#if (CAPNP_VERSION_MAJOR < 1)
+  context.allowCancellation();
+#endif
+}
+
 class PyAsyncIoStream: public kj::AsyncIoStream {
 public:
   kj::Own<PyRefCounter> protocol;

capnp/helpers/helpers.pxd

@@ -1,7 +1,7 @@
 from capnp.includes.capnp_cpp cimport (
     Maybe, PyPromise, VoidPromise, RemotePromise,
-    DynamicCapability, InterfaceSchema, EnumSchema, StructSchema, DynamicValue, Capability, 
-    RpcSystem, MessageBuilder, Own, PyRefCounter, Node, DynamicStruct
+    DynamicCapability, InterfaceSchema, EnumSchema, StructSchema, DynamicValue, Capability,
+    RpcSystem, MessageBuilder, Own, PyRefCounter, Node, DynamicStruct, CallContext
 )
 
 from capnp.includes.schema_cpp cimport ByteArray
@@ -22,6 +22,7 @@ cdef extern from "capnp/helpers/capabilityHelper.h":
     PyPromise convert_to_pypromise(RemotePromise)
     PyPromise convert_to_pypromise(VoidPromise)
     VoidPromise taskToPromise(Own[PyRefCounter] coroutine, PyObject* callback)
+    void allowCancellation(CallContext context) except +reraise_kj_exception nogil
     void init_capnp_api()
 
 cdef extern from "capnp/helpers/rpcHelper.h":

capnp/lib/capnp.pyx

@@ -1927,7 +1927,7 @@ async def kj_loop():
 
 async def run(coro):
     """Ensure that the coroutine runs while the KJ event loop is running
-    
+
     This is a shortcut for wrapping the coroutine in a :py:meth:`capnp.kj_loop` context manager.
 
     :param coro: Coroutine to run
@@ -1949,6 +1949,7 @@ cdef class _CallContext:
     cdef CallContext * thisptr
 
     cdef _init(self, CallContext other):
+        helpers.allowCancellation(other)
         self.thisptr = new CallContext(move(other))
         return self