   """
   @spec remove(String.t()) :: Map
   def remove(delivery_tag) do
-    messages = Agent.get(__MODULE__, fn messages -> messages end)
-    deleted_message = messages[delivery_tag]
-    messages = Map.delete(messages, delivery_tag)
-    
-    # workload = Enum.reduce Map.keys(messages), [], fn(delivery_tag, workload) ->
-    #   workload ++ [%{
-    #     description: "Request:  #{delivery_tag}"
-    #   }]
-    # end
-    #Heartbeat.set_workload(workload)
+    Logger.debug("#{@logprefix} Finished tracking message #{delivery_tag}...")
 
-    Agent.update(__MODULE__, fn _ -> messages end)
-    deleted_message
+    Agent.get_and_update(__MODULE__, fn messages ->
+      message = messages[delivery_tag]
+      {message, Map.delete(messages, delivery_tag)}
+    end)
+  end
   end
 end
  """
  @spec track(Map) :: term
  def track(%{subscription_handler: subscription_handler, delivery_tag: delivery_tag} = _async_info) do
    new_message = %{
      process: self(),
      subscription_handler: subscription_handler, 

  """
  @spec start_link() :: {:ok, pid} | {:error, String.t()} 
  def start_link() do
    Logger.debug("#{@logprefix} Starting...")      
    Agent.start_link(fn -> %{} end, name: __MODULE__)
  end

