diff --git a/python/ofagent/agent.py b/python/ofagent/agent.py
index aa45f67..e4d415d 100755
--- a/python/ofagent/agent.py
+++ b/python/ofagent/agent.py
@@ -35,10 +35,10 @@
     cached_generation_id = None
 
     def __init__(self,
+                 connection_manager,
                  controller_endpoint,
                  datapath_id,
                  device_id,
-                 rpc_stub,
                  enable_tls=False,
                  key_file=None,
                  cert_file=None,
@@ -47,7 +47,7 @@
         self.controller_endpoint = controller_endpoint
         self.datapath_id = datapath_id
         self.device_id = device_id
-        self.rpc_stub = rpc_stub
+        self.connection_manager = connection_manager
         self.enable_tls = enable_tls
         self.key_file = key_file
         self.cert_file = cert_file
@@ -142,7 +142,7 @@
     def protocol(self):
         cxn = OpenFlowConnection(self)  # Low level message handler
         self.proto_handler = OpenFlowProtocolHandler(
-            self.datapath_id, self.device_id, self, cxn, self.rpc_stub)
+            self.datapath_id, self.device_id, self, cxn)
         return cxn
 
     def clientConnectionFailed(self, connector, reason):
@@ -166,7 +166,6 @@
         else:
             log.error('unknown-change-event', change_event=event)
 
-
 if __name__ == '__main__':
     """Run this to test the agent for N concurrent sessions:
        python agent [<number-of-desired-instances>]