diff --git a/src/ptrace.rs b/src/ptrace.rs
index 51f4e78464d09b2e826e9924d6c7942059f65779..5ae54f456747c229dc008158e659de5b61bef650 100644
--- a/src/ptrace.rs
+++ b/src/ptrace.rs
@@ -205,8 +205,8 @@ pub fn clear_breakpoint(pid: ContextId) {
     data.breakpoint = None;
 }
 
-// TODO: All these small functions should be moved to be on the session instance
-pub fn notify(pid: ContextId) {
+/// Notify the tracee of the current session. Returns None if session does not exist
+pub fn notify_tracee(pid: ContextId) {
     let sessions = sessions();
     let session = match sessions.get(&pid) {
         Some(session) => session,
@@ -218,7 +218,7 @@ pub fn notify(pid: ContextId) {
 
 /// Create a new breakpoint for the specified tracee, optionally with
 /// a sysemu flag. Panics if the session is invalid.
-pub fn set_breakpoint(pid: ContextId, flags: PtraceFlags, should_continue: bool) {
+pub fn set_breakpoint(pid: ContextId, flags: PtraceFlags) {
     let sessions = sessions_mut();
     let session = sessions.get(&pid).expect("proc (set_breakpoint): invalid session");
     let mut data = session.data.lock();
diff --git a/src/scheme/proc.rs b/src/scheme/proc.rs
index b198c69bb66aa203fef31a86efad3b7079d42fef..5b538517f8138d57f0caf7a4624d3deaa9fb39b5 100644
--- a/src/scheme/proc.rs
+++ b/src/scheme/proc.rs
@@ -468,6 +468,13 @@ impl Scheme for ProcScheme {
 
                 let should_continue = !op.contains(PTRACE_FLAG_WAIT) || op.intersects(PTRACE_STOP_MASK);
 
+                // Set next breakpoint, or clear it if no stop condition was set and we should continue
+                if op.intersects(PTRACE_STOP_MASK) {
+                    ptrace::set_breakpoint(info.pid, op);
+                } else if should_continue {
+                    ptrace::clear_breakpoint(info.pid);
+                }
+
                 if op.contains(PTRACE_STOP_SINGLESTEP) {
                     try_stop_context(info.pid, |context| {
                         match unsafe { ptrace::regs_for_mut(context) } {
@@ -483,13 +490,6 @@ impl Scheme for ProcScheme {
                     })?;
                 }
 
-                // Set next breakpoint, and potentially restart tracee
-                if op.intersects(PTRACE_STOP_MASK) {
-                    ptrace::set_breakpoint(info.pid, op, should_continue);
-                } else if should_continue {
-                    ptrace::clear_breakpoint(info.pid);
-                }
-
                 if should_continue {
                     // disable the ptrace_stop flag, which is used in some cases
                     with_context_mut(info.pid, |context| {
@@ -498,7 +498,7 @@ impl Scheme for ProcScheme {
                     })?;
 
                     // and notify the tracee's WaitCondition, which is used in other cases
-                    ptrace::notify(info.pid);
+                    ptrace::notify_tracee(info.pid);
                 }
 
                 // And await the tracee, if requested to