diff --git a/src/app.rs b/src/app.rs
index f440c8b..74b9d83 100644
--- a/src/app.rs
+++ b/src/app.rs
@@ -28,14 +28,17 @@ impl App {
     }
 
     fn render(&mut self) {
-        self.renderer.render_fitch(&self.model, &self.info_text());
-        match self.state {
-            State::Noraml => (),
-            State::AddAssumption => self
-                .renderer
-                .render_expression_box("Assumption expression", &self.expression_buffer),
-            _ => (),
-        }
+        let (title, render_box) = match self.state {
+            State::AddAssumption => ("Assumption expression", true),
+            _ => ("", false),
+        };
+        self.renderer.render_fitch(
+            &self.model,
+            &self.info_text(),
+            title,
+            &self.expression_buffer,
+            render_box,
+        );
     }
 
     pub fn listen(&mut self) {
diff --git a/src/ui.rs b/src/ui.rs
index 5078912..bedd9ba 100644
--- a/src/ui.rs
+++ b/src/ui.rs
@@ -30,7 +30,14 @@ impl Renderer {
         Ok(renderer)
     }
 
-    pub fn render_fitch(&mut self, model: &Fitch, info: &str) {
+    pub fn render_fitch(
+        &mut self,
+        model: &Fitch,
+        info: &str,
+        title: &str,
+        buffer: &str,
+        render_box: bool,
+    ) {
         self.terminal
             .draw(|frame| {
                 let (f_a, i_a) = base_area(frame.size());
@@ -43,13 +50,12 @@ impl Renderer {
 
                 frame.render_widget(fitch_widget, f_a);
                 frame.render_widget(info_widget, i_a);
-            })
-            .unwrap();
-    }
 
-    pub fn render_expression_box(&mut self, title: &str, buffer: &str) {
-        self.terminal
-            .draw(|frame| {
+                if !render_box {
+                    return;
+                }
+
+                // Render expression BOX
                 let area = expression_box_area(frame.size());
                 let expression_widget = Paragraph::new(buffer).block(
                     Block::default()