lib/rexec/server.rb in rexec-1.4.0 vs lib/rexec/server.rb in rexec-1.4.1

@@ -82,14 +82,14 @@
 
 		if block_given?
 			Task.open(command, options) do |task|
 				conn = Connection.build(task, options, &send_code)
 
-				yield conn, task
-
-				conn.stop
-				task.stop
-				task.wait
+				begin
+					yield conn, task
+				ensure
+					conn.stop
+				end
 			end
 		else
 			task = Task.open(command, options)
 			conn = Connection.build(task, options, &send_code)