The head of the Iranian armed forces warned Islamabad on Monday that Tehran would hit bases inside Pakistan if the government does not confront militants who carry out cross-border attacks.
Ten Iranian border guards were killed by militants last month. Iran said Jaish-al-Adl, a militant group, had shot the guards with long-range guns, fired from inside Pakistan.
The border area has long been plagued by unrest from both drug smuggling gangs and separatist militants.
"We cannot accept the continuation of this situation," Major General Mohammad Baqeri, the head of the Iranian armed forces, was quoted as saying by state news agency IRNA.
"We expect the Pakistani officials to control the borders, arrest the terrorists and shut down their bases."
"If the terrorist attacks continue, we will hit their safe havens and cells, wherever they are," he said.
Iran's Foreign Minister Mohammad Javad Zarif visited Pakistan last week and asked Prime Minister Nawaz Sharif to improve border security. Pakistan assured Iran it would deploy additional troops along its border.
In 2014, Iran warned it would send troops to Pakistan to retrieve five Iranian border guards kidnapped by Jaish al Adl.
Pakistan said at the time that such action would be a violation of international law and warned Iranian forces not to cross the border.
Iran refrained from sending troops when a local cleric stepped in and resolved the situation.
Four of the guards were released a few months later, but one was killed by the militants.
Jaish al Adl is a militant group that has carried out several attacks against Iranian security forces with the aim of highlighting what they say is discrimination against minority groups in Iran.
The group claimed responsibility for attacks that killed eight border guards in April 2015 and 14 border guards in October 2013.