diff --git a/coreblocks/core.py b/coreblocks/core.py
index 6c06d96c8..070500582 100644
--- a/coreblocks/core.py
+++ b/coreblocks/core.py
@@ -1,17 +1,15 @@
 from amaranth import *
 from amaranth.lib.wiring import Component, flipped, connect, In, Out
 from transactron.lib.allocators import PriorityEncoderAllocator
-from transactron.utils.amaranth_ext.elaboratables import ModuleConnector
 
 from transactron.utils.dependencies import DependencyContext
 from coreblocks.priv.traps.instr_counter import CoreInstructionCounter
 from coreblocks.func_blocks.interface.func_blocks_unifier import FuncBlocksUnifier
 from coreblocks.priv.traps.interrupt_controller import ISA_RESERVED_INTERRUPTS, InternalInterruptController
 from transactron.core import TModule
-from transactron.lib import ConnectTrans, MethodProduct
+from transactron.lib import MethodProduct
 from coreblocks.interface.layouts import *
 from coreblocks.interface.keys import (
-    FetchResumeKey,
     CSRInstancesKey,
     CommonBusDataKey,
 )
@@ -140,13 +138,6 @@ def elaborate(self, platform):
 
         m.submodules.exception_information_register = self.exception_information_register
 
-        fetch_resume = self.connections.get_optional_dependency(FetchResumeKey())
-        if fetch_resume is not None:
-            fetch_resume_fb, fetch_resume_unifiers = fetch_resume
-            m.submodules.fetch_resume_unifiers = ModuleConnector(**fetch_resume_unifiers)
-
-            m.submodules.fetch_resume_connector = ConnectTrans(fetch_resume_fb, self.frontend.resume_from_unsafe)
-
         m.submodules.announcement = self.announcement
         m.submodules.func_blocks_unifier = self.func_blocks_unifier
         m.submodules.retirement = Retirement(
diff --git a/coreblocks/frontend/fetch/fetch.py b/coreblocks/frontend/fetch/fetch.py
index 497ba81ee..eac7f8f67 100644
--- a/coreblocks/frontend/fetch/fetch.py
+++ b/coreblocks/frontend/fetch/fetch.py
@@ -1,11 +1,9 @@
 from amaranth import *
 from amaranth.lib.data import ArrayLayout
-from coreblocks.interface.keys import FetchResumeKey
 from transactron.lib import BasicFifo, WideFifo, Semaphore, logging, Pipe
 from transactron.lib.metrics import *
 from transactron.lib.simultaneous import condition
 from transactron.utils import popcount, assign, StableSelectingNetwork
-from transactron.utils.dependencies import DependencyContext
 from transactron.utils.transactron_helpers import make_layout
 from transactron.utils.amaranth_ext.coding import PriorityEncoder
 from transactron import *
@@ -32,9 +30,18 @@ class FetchUnit(Elaboratable):
 
     The unit also deals with expanding compressed instructions and managing instructions that aren't aligned to
     4-byte boundaries.
+
+    Attributes
+    ----------
+    redirect : Method
+        Redirects the fetch unit to the specified PC
+    flush : Method
+        Flushes the fetch unit from the currently processed fetch blocks, so it can be redirected or/and stalled.
     """
 
-    def __init__(self, gen_params: GenParams, icache: CacheInterface, cont: Method) -> None:
+    def __init__(
+        self, gen_params: GenParams, icache: CacheInterface, cont: Method, stall_lock: Method, stall_unsafe: Method
+    ) -> None:
         """
         Parameters
         ----------
@@ -46,16 +53,21 @@ def __init__(self, gen_params: GenParams, icache: CacheInterface, cont: Method)
         cont : Method
             Method which should be invoked to send fetched instruction to the next step.
             It has layout as described by `FetchLayout`.
+        stall_lock : Method
+            Method whose readiness determines if the fetch unit is stalled
+        stall_unsafe : Method
+            Method that is called when an unsafe instruction is fetched
         """
         self.gen_params = gen_params
         self.icache = icache
         self.cont = cont
+        self.stall_lock = stall_lock
+        self.stall_unsafe = stall_unsafe
 
         self.layouts = self.gen_params.get(FetchLayouts)
 
-        self.resume_from_unsafe = Method(i=self.layouts.resume)
-        self.resume_from_exception = Method(i=self.layouts.resume)
-        self.stall_exception = Method()
+        self.redirect = Method(i=self.layouts.redirect)
+        self.flush = Method()
 
         self.perf_fetch_utilization = TaggedCounter(
             "frontend.fetch.fetch_block_util",
@@ -83,7 +95,9 @@ def elaborate(self, platform):
         )
 
         with Transaction(name="cont").body(m):
-            self.cont(m, serializer.read(m, count=1).data[0])
+            raw_instr = serializer.read(m, count=1).data[0]
+            log.info(m, True, "Sending an instr to the backend pc=0x{:x} instr=0x{:x}", raw_instr.pc, raw_instr.instr)
+            self.cont(m, raw_instr)
 
         m.submodules.cache_requests = cache_requests = BasicFifo(layout=[("addr", self.gen_params.isa.xlen)], depth=2)
 
@@ -100,18 +114,13 @@ def flush():
 
         current_pc = Signal(self.gen_params.isa.xlen, init=self.gen_params.start_pc)
 
-        stalled_unsafe = Signal()
-        stalled_exception = Signal()
-
-        stalled = Signal()
-        m.d.av_comb += stalled.eq(stalled_unsafe | stalled_exception)
-
         #
         # Fetch - stage 0
         # ================
         # - send a request to the instruction cache
         #
-        with Transaction(name="Fetch_Stage0").body(m, request=~stalled):
+        with Transaction(name="Fetch_Stage0").body(m):
+            self.stall_lock(m)
             req_counter.acquire(m)
             self.icache.issue_req(m, addr=current_pc)
             cache_requests.write(m, addr=current_pc)
@@ -351,7 +360,7 @@ def flush():
                         # TODO: Raise different code for page fault when supported
                         # could be passed in 3rd bit of access_fault
                         flush()
-                        m.d.sync += stalled_unsafe.eq(1)
+                        self.stall_unsafe(m)
                     with m.Elif(redirect):
                         self.perf_fetch_redirects.incr(m)
                         new_pc = Signal.like(current_pc)
@@ -372,86 +381,14 @@ def flush():
         with m.If(flush_now):
             m.d.sync += flushing_counter.eq(req_counter.count_next)
 
-        @def_method(m, self.resume_from_unsafe, ready=(stalled & (flushing_counter == 0)))
-        def _(pc: Value):
-            log.info(m, ~stalled_exception, "Resuming from unsafe instruction new_pc=0x{:x}", pc)
-            m.d.sync += current_pc.eq(pc)
-            # If core is stalled because of exception, effect of this call will be ignored, as
-            # `stalled_exception` is not changed
-            m.d.sync += stalled_unsafe.eq(0)
-
-        @def_method(m, self.resume_from_exception, ready=(stalled_exception & (flushing_counter == 0)))
-        def _(pc: Value):
-            log.info(m, True, "Resuming from exception new_pc=0x{:x}", pc)
-            # Resume from exception has implicit priority to resume from unsafe instructions call.
-            # Both could happen at the same time due to resume methods being blocked.
-            # `resume_from_unsafe` will never overwrite `resume_from_exception` event, because there is at most one
-            # unsafe instruction in the core that will call resume_from_unsafe before or at the same time as
-            # `resume_from_exception`.
-            # `current_pc` is set to correct entry at a complete unstall due to method declaration order
-            # See https://github.com/kuznia-rdzeni/coreblocks/pull/654#issuecomment-2057478960
-            m.d.sync += current_pc.eq(pc)
-            m.d.sync += stalled_unsafe.eq(0)
-            m.d.sync += stalled_exception.eq(0)
-
-        # Fetch can be resumed to unstall from 'unsafe' instructions, and stalled because
-        # of exception report, both can happen at any time during normal excecution.
-        # In case of simultaneous call, fetch will be correctly stalled, becasue separate signal is used
-        @def_method(m, self.stall_exception)
+        @def_method(m, self.flush)
         def _():
-            log.info(m, True, "Stalling the fetch unit because of an exception")
-            serializer.clear(m)
-            m.d.sync += stalled_exception.eq(1)
             flush()
+            serializer.clear(m)
 
-        # Fetch resume verification
-        if self.gen_params.extra_verification:
-            expect_unstall_unsafe = Signal()
-            prev_stalled_unsafe = Signal()
-            dependencies = DependencyContext.get()
-            fetch_resume = dependencies.get_optional_dependency(FetchResumeKey())
-            if fetch_resume is not None:
-                unifier_ready = fetch_resume[0].ready
-            else:
-                unifier_ready = C(0)
-
-            m.d.sync += prev_stalled_unsafe.eq(stalled_unsafe)
-            with m.FSM("running"):
-                with m.State("running"):
-                    log.error(m, stalled_exception | prev_stalled_unsafe, "fetch was expected to be running")
-                    log.error(
-                        m,
-                        unifier_ready,
-                        "resume_from_unsafe unifier is ready before stall",
-                    )
-                    with m.If(stalled_unsafe):
-                        m.next = "stalled_unsafe"
-                    with m.If(self.stall_exception.run):
-                        m.next = "stalled_exception"
-                with m.State("stalled_unsafe"):
-                    m.d.sync += expect_unstall_unsafe.eq(1)
-                    with m.If(self.resume_from_unsafe.run):
-                        m.d.sync += expect_unstall_unsafe.eq(0)
-                        m.d.sync += prev_stalled_unsafe.eq(0)  # it is fine to be stalled now
-                        m.next = "running"
-                    with m.If(self.stall_exception.run):
-                        m.next = "stalled_exception"
-                    log.error(
-                        m,
-                        self.resume_from_exception.run & ~self.stall_exception.run,
-                        "unexpected resume_from_exception",
-                    )
-                with m.State("stalled_exception"):
-                    with m.If(self.resume_from_unsafe.run):
-                        log.error(m, ~expect_unstall_unsafe, "unexpected resume_from_unsafe")
-                        m.d.sync += expect_unstall_unsafe.eq(0)
-                    with m.If(self.resume_from_exception.run):
-                        # unstall_form_unsafe may be skipped if excpetion was reported on unsafe instruction,
-                        # invalid cases are verified by readiness check in running state
-                        m.d.sync += expect_unstall_unsafe.eq(0)
-                        m.d.sync += prev_stalled_unsafe.eq(0)  # it is fine to be stalled now
-                        with m.If(~self.stall_exception.run):
-                            m.next = "running"
+        @def_method(m, self.redirect)
+        def _(pc):
+            m.d.sync += current_pc.eq(pc)
 
         return m
 
diff --git a/coreblocks/frontend/frontend.py b/coreblocks/frontend/frontend.py
index f4212d342..d6b847059 100644
--- a/coreblocks/frontend/frontend.py
+++ b/coreblocks/frontend/frontend.py
@@ -7,6 +7,7 @@
 from coreblocks.params.genparams import GenParams
 from coreblocks.frontend.decoder.decode_stage import DecodeStage
 from coreblocks.frontend.fetch.fetch import FetchUnit
+from coreblocks.frontend.stall_controller import StallController
 from coreblocks.cache.icache import ICache, ICacheBypass
 from coreblocks.cache.refiller import SimpleCommonBusCacheRefiller
 from coreblocks.interface.layouts import *
@@ -23,8 +24,6 @@ class CoreFrontend(Elaboratable):
         Consume a single decoded instruction.
     resume_from_exception: Method
         Resume the frontend from the given PC after an exception.
-    resume_from_unsafe: Method
-        Resume the frontend from the given PC after an unsafe instruction.
     stall: Method
         Stall and flush the frontend.
     """
@@ -44,7 +43,15 @@ def __init__(self, *, gen_params: GenParams, instr_bus: BusMasterInterface):
 
         self.connections.add_dependency(FlushICacheKey(), self.icache.flush)
 
-        self.fetch = FetchUnit(self.gen_params, self.icache, self.instr_buffer.write)
+        self.stall_ctrl = StallController(self.gen_params)
+
+        self.fetch = FetchUnit(
+            self.gen_params,
+            self.icache,
+            self.instr_buffer.write,
+            self.stall_ctrl.stall_guard,
+            self.stall_ctrl.stall_unsafe,
+        )
 
         self.decode_pipe = Pipe(self.gen_params.get(DecodeLayouts).decoded_instr)
 
@@ -55,8 +62,7 @@ def __init__(self, *, gen_params: GenParams, instr_bus: BusMasterInterface):
         DependencyContext.get().add_dependency(PredictedJumpTargetKey(), (self.target_pred_req, self.target_pred_resp))
 
         self.consume_instr = self.decode_pipe.read
-        self.resume_from_exception = self.fetch.resume_from_exception
-        self.resume_from_unsafe = self.fetch.resume_from_unsafe
+        self.resume_from_exception = self.stall_ctrl.resume_from_exception
         self.stall = Method()
 
     def elaborate(self, platform):
@@ -74,6 +80,10 @@ def elaborate(self, platform):
             gen_params=self.gen_params, get_raw=self.instr_buffer.read, push_decoded=self.decode_pipe.write
         )
 
+        m.submodules.stall_ctrl = self.stall_ctrl
+
+        self.stall_ctrl.redirect_frontend.proxy(m, self.fetch.redirect)
+
         # TODO: Remove when Branch Predictor implemented
         with Transaction(name="DiscardBranchVerify").body(m):
             read = self.connections.get_dependency(BranchVerifyKey())
@@ -89,7 +99,8 @@ def _(arg):
 
         @def_method(m, self.stall)
         def _():
-            self.fetch.stall_exception(m)
+            self.stall_ctrl.stall_exception(m)
+            self.fetch.flush(m)
             self.instr_buffer.clear(m)
             self.decode_pipe.clean(m)
 
diff --git a/coreblocks/frontend/stall_controller.py b/coreblocks/frontend/stall_controller.py
new file mode 100644
index 000000000..c3cdb7c3c
--- /dev/null
+++ b/coreblocks/frontend/stall_controller.py
@@ -0,0 +1,107 @@
+from typing import Sequence
+import operator
+from functools import reduce
+
+from amaranth import *
+
+
+from transactron.lib import logging
+from transactron.lib.metrics import *
+from transactron.lib.simultaneous import condition
+from transactron.utils import popcount, DependencyContext, MethodStruct
+from transactron.utils.assign import AssignArg
+from transactron import *
+
+from coreblocks.params import *
+from coreblocks.interface.layouts import *
+from coreblocks.interface.keys import UnsafeInstructionResolvedKey
+
+log = logging.HardwareLogger("frontend.stall_ctrl")
+
+
+class StallController(Elaboratable):
+    """
+    The stall controller is responsible for managing all stall/unstall signals
+    that may be triggered in the core and based on them deciding if and where
+    the frontend should be resumed.
+
+
+    Attributes
+    ----------
+    stall_unsafe : Method
+        Signals that the frontend should be stalled because an unsafe (i.e. causing
+        difficult to handle side effects) was just fetched.
+    stall_exception : Method
+        Signals that the frontend should be stalled because of an exception.
+    stall_guard : Method
+        A non-exclusive method whose readiness denotes if the frontend is currently stalled.
+    resume_from_exception: Method
+        Signals that the backend handled the exception and the frontend can be resumed.
+    redirect_frontend : Method (bodyless)
+        A method that will be called when the frontend needs to be redirected.
+    """
+
+    def __init__(self, gen_params: GenParams):
+        self.gen_params = gen_params
+
+        layouts = self.gen_params.get(FetchLayouts)
+
+        self.stall_unsafe = Method()
+        self.stall_exception = Method()
+        self.stall_guard = Method()
+        self.resume_from_exception = Method(i=layouts.resume)
+        self._resume_from_unsafe = Method(i=layouts.resume)
+
+        self.redirect_frontend = Method(i=layouts.redirect)
+
+        DependencyContext.get().add_dependency(UnsafeInstructionResolvedKey(), self._resume_from_unsafe)
+
+    def elaborate(self, platform):
+        m = TModule()
+
+        stalled_unsafe = Signal()
+        stalled_exception = Signal()
+
+        @def_method(m, self.stall_guard, ready=~(stalled_unsafe | stalled_exception), nonexclusive=True)
+        def _():
+            pass
+
+        def resume_combiner(m: Module, args: Sequence[MethodStruct], runs: Value) -> AssignArg:
+            # Make sure that there is at most one caller - there can be only one unsafe instruction
+            log.assertion(m, popcount(runs) <= 1)
+            pcs = [Mux(runs[i], v.pc, 0) for i, v in enumerate(args)]
+
+            return {"pc": reduce(operator.or_, pcs, 0)}
+
+        @def_method(m, self._resume_from_unsafe, nonexclusive=True, combiner=resume_combiner)
+        def _(pc):
+            log.assertion(m, stalled_unsafe)
+            m.d.sync += stalled_unsafe.eq(0)
+
+            with condition(m, nonblocking=True) as branch:
+                with branch(~stalled_exception):
+                    log.info(m, True, "Resuming from unsafe instruction new_pc=0x{:x}", pc)
+                    self.redirect_frontend(m, pc=pc)
+
+        @def_method(m, self.resume_from_exception)
+        def _(pc):
+            m.d.sync += stalled_unsafe.eq(0)
+            m.d.sync += stalled_exception.eq(0)
+
+            log.info(m, True, "Resuming from exception new_pc=0x{:x}", pc)
+            self.redirect_frontend(m, pc=pc)
+
+        @def_method(m, self.stall_unsafe)
+        def _():
+            log.assertion(m, ~stalled_unsafe, "Can't be stalled twice because of an unsafe instruction")
+            log.info(m, True, "Stalling the frontend because of an unsafe instruction")
+            m.d.sync += stalled_unsafe.eq(1)
+
+        # Fetch can be resumed to unstall from 'unsafe' instructions, and stalled because
+        # of exception report, both can happen at any time during normal excecution.
+        @def_method(m, self.stall_exception)
+        def _():
+            log.info(m, ~stalled_exception, "Stalling the frontend because of an exception")
+            m.d.sync += stalled_exception.eq(1)
+
+        return m
diff --git a/coreblocks/func_blocks/csr/csr.py b/coreblocks/func_blocks/csr/csr.py
index 75cd23d19..e5fbb0f37 100644
--- a/coreblocks/func_blocks/csr/csr.py
+++ b/coreblocks/func_blocks/csr/csr.py
@@ -12,10 +12,10 @@
 from coreblocks.params import GenParams
 from coreblocks.params.fu_params import BlockComponentParams
 from coreblocks.func_blocks.interface.func_protocols import FuncBlock
-from coreblocks.interface.layouts import FetchLayouts, FuncUnitLayouts, CSRUnitLayouts
+from coreblocks.interface.layouts import FuncUnitLayouts, CSRUnitLayouts
 from coreblocks.interface.keys import (
     CSRListKey,
-    FetchResumeKey,
+    UnsafeInstructionResolvedKey,
     CSRInstancesKey,
     InstructionPrecommitKey,
     ExceptionReportKey,
@@ -71,8 +71,6 @@ def __init__(self, gen_params: GenParams):
         self.gen_params = gen_params
         self.dependency_manager = DependencyContext.get()
 
-        self.fetch_resume = Method(o=gen_params.get(FetchLayouts).resume)
-
         # Standard RS interface
         self.csr_layouts = gen_params.get(CSRUnitLayouts)
         self.fu_layouts = gen_params.get(FuncUnitLayouts)
@@ -99,7 +97,6 @@ def elaborate(self, platform):
         reserved = Signal()
         ready_to_process = Signal()
         done = Signal()
-        call_resume = Signal()
         exception = Signal()
 
         current_result = Signal(self.gen_params.isa.xlen)
@@ -210,6 +207,7 @@ def _():
 
             report = self.dependency_manager.get_dependency(ExceptionReportKey())
             interrupt = self.dependency_manager.get_dependency(AsyncInterruptInsertSignalKey())
+            resume_core = self.dependency_manager.get_dependency(UnsafeInstructionResolvedKey())
 
             with m.If(exception):
                 mtval = Signal(self.gen_params.isa.xlen)
@@ -243,7 +241,9 @@ def _():
 
             m.d.sync += exception.eq(0)
 
-            m.d.comb += call_resume.eq(exe_side_fx & ~exception & ~interrupt)
+            with m.If(exe_side_fx & ~exception & ~interrupt):
+                # CSR instructions are never compressed, PC+4 is always next instruction
+                resume_core(m, pc=instr.pc + self.gen_params.isa.ilen_bytes)
 
             return {
                 "rob_id": instr.rob_id,
@@ -252,23 +252,13 @@ def _():
                 "exception": exception | interrupt,
             }
 
-        @def_method(m, self.fetch_resume, call_resume)
-        def _():
-            # This call will always execute, because there is at most one unsafe instruction in the core, and it can be
-            # stored in unifer's Forwarder unitl resume becomes ready.
-            # CSR instructions are never compressed, PC+4 is always next instruction
-            return {"pc": instr.pc + self.gen_params.isa.ilen_bytes}
-
         return m
 
 
 @dataclass(frozen=True)
 class CSRBlockComponent(BlockComponentParams):
     def get_module(self, gen_params: GenParams) -> FuncBlock:
-        connections = DependencyContext.get()
-        unit = CSRUnit(gen_params)
-        connections.add_dependency(FetchResumeKey(), unit.fetch_resume)
-        return unit
+        return CSRUnit(gen_params)
 
     def get_optypes(self) -> set[OpType]:
         return {OpType.CSR_REG, OpType.CSR_IMM}
diff --git a/coreblocks/func_blocks/fu/priv.py b/coreblocks/func_blocks/fu/priv.py
index 755a73887..8ad116439 100644
--- a/coreblocks/func_blocks/fu/priv.py
+++ b/coreblocks/func_blocks/fu/priv.py
@@ -7,7 +7,7 @@
 
 
 from transactron import *
-from transactron.lib import BasicFifo, logging
+from transactron.lib import logging
 from transactron.lib.metrics import TaggedCounter
 from transactron.lib.simultaneous import condition
 from transactron.utils import DependencyContext, OneHotSwitch
@@ -15,14 +15,14 @@
 from coreblocks.params import *
 from coreblocks.params import GenParams, FunctionalComponentParams
 from coreblocks.arch import OpType, ExceptionCause
-from coreblocks.interface.layouts import FuncUnitLayouts, FetchLayouts
+from coreblocks.interface.layouts import FuncUnitLayouts
 from coreblocks.interface.keys import (
     MretKey,
     AsyncInterruptInsertSignalKey,
     ExceptionReportKey,
     CSRInstancesKey,
     InstructionPrecommitKey,
-    FetchResumeKey,
+    UnsafeInstructionResolvedKey,
     FlushICacheKey,
     WaitForInterruptResumeKey,
 )
@@ -56,8 +56,6 @@ def __init__(self, gen_params: GenParams, priv_fn=PrivilegedFn()):
         self.issue = Method(i=layouts.issue)
         self.push_result = Method(i=layouts.push_result)
 
-        self.fetch_resume_fifo = BasicFifo(self.gen_params.get(FetchLayouts).resume, 2)
-
         self.perf_instr = TaggedCounter(
             "backend.fu.priv.instr",
             "Number of instructions precommited with side effects by the priviledge unit",
@@ -86,8 +84,7 @@ def elaborate(self, platform):
         csr = self.dm.get_dependency(CSRInstancesKey())
         priv_mode = csr.m_mode.priv_mode
         flush_icache = self.dm.get_dependency(FlushICacheKey())
-
-        m.submodules.fetch_resume_fifo = self.fetch_resume_fifo
+        resume_core = self.dm.get_dependency(UnsafeInstructionResolvedKey())
 
         @def_method(m, self.issue, ready=~instr_valid)
         def _(arg):
@@ -177,7 +174,7 @@ def _(arg):
             with m.Else():
                 log.info(m, True, "Unstalling fetch from the priv unit new_pc=0x{:x}", ret_pc)
                 # Unstall the fetch
-                self.fetch_resume_fifo.write(m, pc=ret_pc)
+                resume_core(m, pc=ret_pc)
 
             self.push_result(
                 m,
@@ -194,8 +191,5 @@ def _(arg):
 class PrivilegedUnitComponent(FunctionalComponentParams):
     decoder_manager: PrivilegedFn = PrivilegedFn()
 
-    def get_module(self, gp: GenParams) -> FuncUnit:
-        unit = PrivilegedFuncUnit(gp, self.decoder_manager)
-        connections = DependencyContext.get()
-        connections.add_dependency(FetchResumeKey(), unit.fetch_resume_fifo.read)
-        return unit
+    def get_module(self, gen_params: GenParams) -> FuncUnit:
+        return PrivilegedFuncUnit(gen_params, self.decoder_manager)
diff --git a/coreblocks/interface/keys.py b/coreblocks/interface/keys.py
index 6c3dc661c..da42a84ea 100644
--- a/coreblocks/interface/keys.py
+++ b/coreblocks/interface/keys.py
@@ -1,9 +1,8 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
-from transactron.lib.dependencies import SimpleKey, UnifierKey, ListKey
+from transactron.lib.dependencies import SimpleKey, ListKey
 from transactron import Method
-from transactron.lib import Collector
 from coreblocks.peripherals.bus_adapter import BusMasterInterface
 from amaranth import Signal
 
@@ -16,7 +15,7 @@
     "InstructionPrecommitKey",
     "BranchVerifyKey",
     "PredictedJumpTargetKey",
-    "FetchResumeKey",
+    "UnsafeInstructionResolvedKey",
     "ExceptionReportKey",
     "CSRInstancesKey",
     "AsyncInterruptInsertSignalKey",
@@ -48,7 +47,12 @@ class PredictedJumpTargetKey(SimpleKey[tuple[Method, Method]]):
 
 
 @dataclass(frozen=True)
-class FetchResumeKey(UnifierKey, unifier=Collector):
+class UnsafeInstructionResolvedKey(SimpleKey[Method]):
+    """
+    Represents a method that is called by functional units when
+    an unsafe instruction is executed and the core should be resumed.
+    """
+
     pass
 
 
diff --git a/coreblocks/interface/layouts.py b/coreblocks/interface/layouts.py
index bb5cf217b..6579dde23 100644
--- a/coreblocks/interface/layouts.py
+++ b/coreblocks/interface/layouts.py
@@ -467,6 +467,7 @@ def __init__(self, gen_params: GenParams):
             fields.predicted_taken,
         )
 
+        self.redirect = make_layout(fields.pc)
         self.resume = make_layout(fields.pc)
 
         self.predecoded_instr = make_layout(fields.cfi_type, ("cfi_offset", signed(21)), ("unsafe", 1))
diff --git a/scripts/core_graph.py b/scripts/core_graph.py
index 5024a88f8..67e91e93d 100755
--- a/scripts/core_graph.py
+++ b/scripts/core_graph.py
@@ -19,20 +19,22 @@
 from coreblocks.params.configurations import basic_core_config  # noqa: E402
 from transactron.core import TransactionModule  # noqa: E402
 from transactron.core.keys import TransactionManagerKey  # noqa: E402
+from transactron.utils import DependencyManager, DependencyContext # noqa: E402
 
-gp = GenParams(basic_core_config)
-elaboratable = CoreTestElaboratable(gp)
-tm = TransactionModule(elaboratable)
-fragment = TracingFragment.get(tm, platform=None).prepare()
+with DependencyContext(DependencyManager()):
+    gp = GenParams(basic_core_config)
+    elaboratable = CoreTestElaboratable(gp)
+    tm = TransactionModule(elaboratable)
+    fragment = TracingFragment.get(tm, platform=None).prepare()
 
-core = fragment
-while not hasattr(core, "manager"):
-    core = core._tracing_original  # type: ignore
+    core = fragment
+    while not hasattr(core, "manager"):
+        core = core._tracing_original  # type: ignore
 
-mgr = core.manager.get_dependency(TransactionManagerKey())  # type: ignore
+    mgr = core.manager.get_dependency(TransactionManagerKey())  # type: ignore
 
-with arg.ofile as fp:
-    graph = mgr.visual_graph(fragment)
-    if arg.prune:
-        graph.prune()
-    graph.dump(fp, arg.format)
+    with arg.ofile as fp:
+        graph = mgr.visual_graph(fragment)
+        if arg.prune:
+            graph.prune()
+        graph.dump(fp, arg.format)
diff --git a/test/frontend/test_fetch.py b/test/frontend/test_fetch.py
index 6310f0d14..e60a15c4d 100644
--- a/test/frontend/test_fetch.py
+++ b/test/frontend/test_fetch.py
@@ -6,7 +6,6 @@
 import random
 
 from amaranth import Elaboratable, Module
-from coreblocks.interface.keys import FetchResumeKey
 
 from transactron.core import Method
 from transactron.lib import AdapterTrans, Adapter, BasicFifo
@@ -27,7 +26,6 @@
 from coreblocks.params import *
 from coreblocks.params.configurations import test_core_config
 from coreblocks.interface.layouts import ICacheLayouts, FetchLayouts
-from transactron.utils.dependencies import DependencyContext
 
 
 class MockedICache(Elaboratable, CacheInterface):
@@ -77,19 +75,32 @@ def setup(self, fixture_initialize_testing_env):
         self.icache = MockedICache(self.gen_params)
         fifo = BasicFifo(self.gen_params.get(FetchLayouts).raw_instr, depth=2)
         self.io_out = TestbenchIO(AdapterTrans(fifo.read))
-        self.clean_fifo = TestbenchIO(AdapterTrans(fifo.clear))
+        self.clear_fifo = TestbenchIO(AdapterTrans(fifo.clear))
         self.fetch_resume_mock = TestbenchIO(Adapter.create())
-        DependencyContext.get().add_dependency(FetchResumeKey(), self.fetch_resume_mock.adapter.iface)
 
-        self.fetch = SimpleTestCircuit(FetchUnit(self.gen_params, self.icache, fifo.write))
+        self.mock_stall_lock = TestbenchIO(Adapter.create())
+        self.mock_stall_unsafe = TestbenchIO(Adapter.create())
 
-        self.m = ModuleConnector(self.icache, fifo, self.io_out, self.clean_fifo, self.fetch)
+        self.fetch = SimpleTestCircuit(
+            FetchUnit(
+                self.gen_params,
+                self.icache,
+                fifo.write,
+                self.mock_stall_lock.adapter.iface,
+                self.mock_stall_unsafe.adapter.iface,
+            )
+        )
+
+        self.m = ModuleConnector(
+            self.icache, fifo, self.io_out, self.clear_fifo, self.fetch, self.mock_stall_lock, self.mock_stall_unsafe
+        )
 
         self.instr_queue = deque()
         self.mem = {}
         self.memerr = set()
         self.input_q = deque()
         self.output_q = deque()
+        self.stalled = False
 
         random.seed(41)
 
@@ -183,6 +194,14 @@ def eff():
         if self.output_q:
             return self.output_q[0]
 
+    @def_method_mock(lambda self: self.mock_stall_lock, enable=lambda self: not self.stalled)
+    def stall_lock_mock(self):
+        pass
+
+    @def_method_mock(lambda self: self.mock_stall_unsafe)
+    def stall_lock_unsafe(self):
+        pass
+
     async def fetch_out_check(self, sim: TestbenchContext):
         while self.instr_queue:
             instr = self.instr_queue.popleft()
@@ -202,11 +221,12 @@ async def fetch_out_check(self, sim: TestbenchContext):
 
             if (instr["jumps"] and (instr["branch_taken"] != v["predicted_taken"])) or access_fault:
                 await self.random_wait(sim, 5)
-                await self.fetch.stall_exception.call(sim)
+                self.stalled = True
+                await self.fetch.flush.call(sim)
                 await self.random_wait(sim, 5)
 
                 # Empty the pipeline
-                await self.clean_fifo.call_try(sim)
+                await self.clear_fifo.call_try(sim)
                 await sim.tick()
 
                 resume_pc = instr["next_pc"]
@@ -217,9 +237,11 @@ async def fetch_out_check(self, sim: TestbenchContext):
                     ) + self.gen_params.fetch_block_bytes
 
                 # Resume the fetch unit
-                while await self.fetch.resume_from_exception.call_try(sim, pc=resume_pc) is None:
+                while await self.fetch.redirect.call_try(sim, pc=resume_pc) is None:
                     pass
 
+                self.stalled = False
+
     def run_sim(self):
         with self.run_simulation(self.m) as sim:
             sim.add_process(self.cache_process)
diff --git a/test/func_blocks/csr/test_csr.py b/test/func_blocks/csr/test_csr.py
index 7fc847f69..a42a9b30f 100644
--- a/test/func_blocks/csr/test_csr.py
+++ b/test/func_blocks/csr/test_csr.py
@@ -9,9 +9,10 @@
 from coreblocks.params import GenParams
 from coreblocks.arch import Funct3, ExceptionCause, OpType
 from coreblocks.params.configurations import test_core_config
-from coreblocks.interface.layouts import ExceptionRegisterLayouts, RetirementLayouts
+from coreblocks.interface.layouts import ExceptionRegisterLayouts, RetirementLayouts, FetchLayouts
 from coreblocks.interface.keys import (
     AsyncInterruptInsertSignalKey,
+    UnsafeInstructionResolvedKey,
     ExceptionReportKey,
     InstructionPrecommitKey,
     CSRInstancesKey,
@@ -48,6 +49,9 @@ def elaborate(self, platform):
         m.submodules.insert = self.insert = TestbenchIO(AdapterTrans(self.dut.insert))
         m.submodules.update = self.update = TestbenchIO(AdapterTrans(self.dut.update))
         m.submodules.accept = self.accept = TestbenchIO(AdapterTrans(self.dut.get_result))
+        m.submodules.fetch_resume = self.fetch_resume = TestbenchIO(
+            Adapter.create(i=self.gen_params.get(FetchLayouts).resume)
+        )
         m.submodules.exception_report = self.exception_report = TestbenchIO(
             Adapter.create(i=self.gen_params.get(ExceptionRegisterLayouts).report)
         )
@@ -56,8 +60,7 @@ def elaborate(self, platform):
         DependencyContext.get().add_dependency(ExceptionReportKey(), self.exception_report.adapter.iface)
         DependencyContext.get().add_dependency(AsyncInterruptInsertSignalKey(), Signal())
         DependencyContext.get().add_dependency(CSRInstancesKey(), self.csr_instances)
-
-        m.submodules.fetch_resume = self.fetch_resume = TestbenchIO(AdapterTrans(self.dut.fetch_resume))
+        DependencyContext.get().add_dependency(UnsafeInstructionResolvedKey(), self.fetch_resume.adapter.iface)
 
         self.csr = {}