diff --git a/kern/include/thread.h b/kern/include/thread.h
index 42d56b4..04a46d3 100644
--- a/kern/include/thread.h
+++ b/kern/include/thread.h
@@ -178,5 +178,7 @@ void schedule(void);
  */
 void thread_consider_migration(void);
 
+extern unsigned thread_count;
+void thread_wait_for_count(unsigned);
 
 #endif /* _THREAD_H_ */
diff --git a/kern/main/menu.c b/kern/main/menu.c
index 057438a..2451b13 100644
--- a/kern/main/menu.c
+++ b/kern/main/menu.c
@@ -117,6 +117,7 @@ common_prog(int nargs, char **args)
 {
 	struct proc *proc;
 	int result;
+	unsigned tc;
 
 	/* Create a process for the new program to run in. */
 	proc = proc_create_runprogram(args[0] /* name */);
@@ -124,6 +125,8 @@ common_prog(int nargs, char **args)
 		return ENOMEM;
 	}
 
+	tc = thread_count;
+
 	result = thread_fork(args[0] /* thread name */,
 			proc /* new process */,
 			cmd_progthread /* thread function */,
@@ -139,6 +142,10 @@ common_prog(int nargs, char **args)
 	 * once you write the code for handling that.
 	 */
 
+	// Wait for all threads to finish cleanup, otherwise khu be a bit behind,
+	// especially once swapping is enabled.
+	thread_wait_for_count(tc);
+
 	return 0;
 }
 
diff --git a/kern/thread/thread.c b/kern/thread/thread.c
index 7a06d07..8ceeaf5 100644
--- a/kern/thread/thread.c
+++ b/kern/thread/thread.c
@@ -70,6 +70,11 @@ unsigned num_cpus;
 /* Used to wait for secondary CPUs to come online. */
 static struct semaphore *cpu_startup_sem;
 
+/* Used to synchronize exit cleanup. */
+unsigned thread_count = 0;
+static struct spinlock thread_count_lock = SPINLOCK_INITIALIZER;
+static struct wchan *thread_count_wchan;
+
 ////////////////////////////////////////////////////////////
 
 /*
@@ -425,6 +430,7 @@ thread_start_cpus(void)
 	kprintf("cpu0: %s\n", buf);
 
 	cpu_startup_sem = sem_create("cpu_hatch", 0);
+	thread_count_wchan = wchan_create("thread_count");
 	mainbus_start_cpus();
 
 	num_cpus = cpuarray_num(&allcpus);
@@ -438,6 +444,12 @@ thread_start_cpus(void)
 		kprintf("%d CPUs online\n", i + 1);
 	}
 	cpu_startup_sem = NULL;
+
+	// Gross hack to deal with os/161 "idle" threads. Hardcode the thread count
+	// to 1 so the inc/dec properly works in thread_[fork/exit]. The one thread
+	// is the cpu0 boot thread (menu), which is the only thread that hasn't
+	// exited yet.
+	thread_count = 1;
 }
 
 /*
@@ -536,6 +548,11 @@ thread_fork(const char *name,
 	 */
 	newthread->t_iplhigh_count++;
 
+	spinlock_acquire(&thread_count_lock);
+	++thread_count;
+	wchan_wakeall(thread_count_wchan, &thread_count_lock);
+	spinlock_release(&thread_count_lock);
+
 	/* Set up the switchframe so entrypoint() gets called */
 	switchframe_init(newthread, entrypoint, data1, data2);
 
@@ -799,6 +816,14 @@ thread_exit(void)
 	/* Check the stack guard band. */
 	thread_checkstack(cur);
 
+	// Decrement the thread count and notify anyone interested.
+	if (thread_count) {
+		spinlock_acquire(&thread_count_lock);
+		--thread_count;
+		wchan_wakeall(thread_count_wchan, &thread_count_lock);
+		spinlock_release(&thread_count_lock);
+	}
+
 	/* Interrupts off on this processor */
 	splhigh();
 	thread_switch(S_ZOMBIE, NULL, NULL);
@@ -1196,3 +1221,15 @@ interprocessor_interrupt(void)
 	curcpu->c_ipi_pending = 0;
 	spinlock_release(&curcpu->c_ipi_lock);
 }
+
+/*
+ * Wait for the thread count to equal tc.
+ */
+void thread_wait_for_count(unsigned tc)
+{
+	spinlock_acquire(&thread_count_lock);
+	while (thread_count != tc) {
+		wchan_sleep(thread_count_wchan, &thread_count_lock);
+	}
+	spinlock_release(&thread_count_lock);
+}