diff --git a/pyproject.toml b/pyproject.toml
index 0833849..0a8e1d4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -5,7 +5,7 @@ requires = [
 build-backend = "hatchling.build"
 [project]
 name = "itp_interface"
-version = "1.1.10"
+version = "1.1.12"
 authors = [
   { name="Amitayush Thakur", email="amitayush@utexas.edu" },
 ]
diff --git a/src/itp_interface/tools/lean4_sync_executor.py b/src/itp_interface/tools/lean4_sync_executor.py
index 9d8311f..1e7ba6c 100644
--- a/src/itp_interface/tools/lean4_sync_executor.py
+++ b/src/itp_interface/tools/lean4_sync_executor.py
@@ -30,8 +30,8 @@ class Lean4SyncExecutor:
     # We ONLY support proofs which are written in tactic mode i.e. with := syntax
     theorem_endings = r"(:=|(\|[\S|\s]*=>))"
     theorem_end_regex = r"(theorem|lemma|example)([\s|\S]*?)(:=|=>)"
-    theorem_regex = r"((((theorem|lemma)[\s]+([\S]*))|example)([\S|\s]*?)(:=|=>)[\s]*?)[\s]+"
-    theorem_name_regex = r"(((theorem|lemma)[\s]+([\S]*))|example)"
+    theorem_regex = r"((((theorem|lemma)[\s]+([^\s:]*))|example)([\S|\s]*?)(:=|=>)[\s]*?)[\s]+"
+    theorem_name_regex = r"(((theorem|lemma)[\s]+([^\s:]*))|example)"
     remove_proof_regex = r"([\s|\S]*(:=|\|))[\s|\S]*?"
     proof_context_separator = "⊢"
     proof_context_regex = r"((\d+) goals)*([\s|\S]*?)\n\n"
@@ -771,7 +771,7 @@ def _parse_response(self, idx, response, relevant_messages = [], dont_update_err
             messages = response['messages']
             # Go over all sev after the line number and check if there is an error
             for msg_idx, msg in enumerate(messages):
-                full_error_msg = self._get_error_msg(msg_idx, msg)
+                full_error_msg = f"Line: {idx} " + self._get_error_msg(msg_idx, msg)
                 unsolved_goal_never_seen_before = not (full_error_msg in self._error_messages_since_last_thm.values())
                 if msg['severity'] == 'error' and 'pos' in msg and 'endPos' in msg and \
                 ((msg['endPos'] is not None and 'line' in msg['endPos']) or \
diff --git a/src/test/simple_env_test.py b/src/test/simple_env_test.py
index d6f3b1d..62b9fea 100644
--- a/src/test/simple_env_test.py
+++ b/src/test/simple_env_test.py
@@ -80,13 +80,17 @@ def test_simple_lean4(self):
         retrieval_strategy = ProofEnvReRankStrategy.NO_RE_RANK
         env = ProofEnv("test_lean4", proof_exec_callback, theorem_name, retrieval_strategy=retrieval_strategy, max_proof_depth=10, always_retrieve_thms=always_retrieve_thms)
         proof_steps = [
+            '-- TODO',
             'apply And.intro',
             'exact hp',
             'apply And.intro',
+            '--TODO',
+            '-- This is just some extra comment',
             'exact hq',
             'exact hp'
         ]
         with env:
+            proof_was_finished = False
             for proof_step in proof_steps:
                 state, _, next_state, _, done, info = env.step(ProofAction(
                     ProofAction.ActionType.RUN_TACTIC, 
@@ -99,6 +103,7 @@ def test_simple_lean4(self):
                 print('-'*30)
                 if done:
                     print("Proof Finished!!")
+                    proof_was_finished = True
                 else:
                     s1 : ProofState = state
                     s2 : ProofState = next_state
@@ -120,6 +125,7 @@ def test_simple_lean4(self):
                         print('|- ', end='')
                         print(goal.goal)
                     print(f"="*30)
+            assert proof_was_finished, "Proof was not finished"
 
     def test_lean4_backtracking(self):
         from itp_interface.rl.proof_state import ProofState
@@ -156,6 +162,7 @@ def test_lean4_backtracking(self):
         ]
         with env:
             prev_state = env.state
+            proof_was_finished = False
             for idx, proof_step in enumerate(proof_steps):
                 if idx > 0 and random.random() <= 0.5:
                     print(f"Backtracking at step {idx + 1} i.e. {proof_step}")
@@ -179,6 +186,8 @@ def test_lean4_backtracking(self):
                 prev_state = state
                 if done:
                     print("Proof Finished!!")
+                    proof_was_finished = True
+            assert proof_was_finished, "Proof was not finished"
 
     def test_simple_coq(self):
         from itp_interface.rl.proof_state import ProofState
@@ -286,6 +295,7 @@ def test_simple_lean_calc(self):
 "_ = (n + 1)*(n + 1) := by \n   rw [Nat.right_distrib n 1 (n + 1)]"
 ]
         with env:
+            proof_was_finished = False
             for proof_step in proof_steps:
                 state, _, next_state, _, done, info = env.step(ProofAction(
                     ProofAction.ActionType.RUN_TACTIC, 
@@ -310,6 +320,7 @@ def test_simple_lean_calc(self):
                     print(f"Action: {proof_step}")
                     print(f"="*30)
                     print("Proof Finished!!")
+                    proof_was_finished = True
                 else:
                     s1 : ProofState = state
                     s2 : ProofState = next_state
@@ -331,6 +342,7 @@ def test_simple_lean_calc(self):
                         print('|- ', end='')
                         print(goal.goal)
                     print(f"="*30)
+            assert proof_was_finished, "Proof was not finished"
 
     def test_simple_lean_enforce_done_test(self):
         from itp_interface.rl.proof_state import ProofState
@@ -372,6 +384,7 @@ def test_simple_lean_enforce_done_test(self):
 "done"
 ]
         with env:
+            proof_finished = False
             for proof_step in proof_steps:
                 state, _, next_state, _, done, info = env.step(ProofAction(
                     ProofAction.ActionType.RUN_TACTIC, 
@@ -397,6 +410,7 @@ def test_simple_lean_enforce_done_test(self):
                     print(f"Action: {proof_step}")
                     print(f"="*30)
                     print("Proof Finished!!")
+                    proof_finished = True
                 else:
                     s1 : ProofState = state
                     s2 : ProofState = next_state
@@ -418,6 +432,7 @@ def test_simple_lean_enforce_done_test(self):
                         print('|- ', end='')
                         print(goal.goal)
                     print(f"="*30)
+            assert proof_finished, "Proof was not finished"
 
     def test_simple_lean4_done_test(self):
         from itp_interface.rl.proof_state import ProofState