diff --git a/firmware/2lib/2ui_screens.c b/firmware/2lib/2ui_screens.c
index 9571ab1..4aeeda1 100644
--- a/firmware/2lib/2ui_screens.c
+++ b/firmware/2lib/2ui_screens.c
@@ -60,6 +60,7 @@
  * current_page is valid in prev and next actions, and the back_item is assigned
  * to a correct menu item index.
  */
+/* TODO(b/174127808): Split out enabling/disabling buttons. */
 
 static vb2_error_t log_page_init(struct vb2_ui_context *ui)
 {
@@ -1153,6 +1154,7 @@
 		ui->error_code = VB2_UI_ERROR_DIAGNOSTICS;
 		return vb2_ui_screen_back(ui);
 	}
+	/* TODO(b/174127808): Integrate this into a new log_page_* function. */
 	if (ui->state->current_page >= ui->state->page_count)
 		ui->state->current_page = ui->state->page_count - 1;
 
@@ -1182,14 +1184,24 @@
 
 static vb2_error_t diagnostics_memory_init_quick(struct vb2_ui_context *ui)
 {
-	return diagnostics_memory_update_screen(
+	vb2_error_t rv;
+	rv = diagnostics_memory_update_screen(
 		ui, &vb2ex_diag_memory_quick_test, 1);
+
+	if (rv != VB2_REQUEST_UI_CONTINUE)
+		return rv;
+	return log_page_init(ui);
 }
 
 static vb2_error_t diagnostics_memory_init_full(struct vb2_ui_context *ui)
 {
-	return diagnostics_memory_update_screen(
+	vb2_error_t rv;
+	rv = diagnostics_memory_update_screen(
 		ui, &vb2ex_diag_memory_full_test, 1);
+
+	if (rv != VB2_REQUEST_UI_CONTINUE)
+		return rv;
+	return log_page_init(ui);
 }
 
 static vb2_error_t diagnostics_memory_update_quick(struct vb2_ui_context *ui)
@@ -1221,6 +1233,9 @@
 	.init = diagnostics_memory_init_quick,
 	.action = diagnostics_memory_update_quick,
 	.menu = MENU_ITEMS(diagnostics_memory_items),
+	.page_up_item = DIAGNOSTICS_MEMORY_ITEM_PAGE_UP,
+	.page_down_item = DIAGNOSTICS_MEMORY_ITEM_PAGE_DOWN,
+	.back_item = DIAGNOSTICS_MEMORY_ITEM_CANCEL,
 };
 
 static const struct vb2_screen_info diagnostics_memory_full_screen = {
@@ -1229,6 +1244,9 @@
 	.init = diagnostics_memory_init_full,
 	.action = diagnostics_memory_update_full,
 	.menu = MENU_ITEMS(diagnostics_memory_items),
+	.page_up_item = DIAGNOSTICS_MEMORY_ITEM_PAGE_UP,
+	.page_down_item = DIAGNOSTICS_MEMORY_ITEM_PAGE_DOWN,
+	.back_item = DIAGNOSTICS_MEMORY_ITEM_CANCEL,
 };
 
 /******************************************************************************/