From d5214e6cd2031fb1c8cf51be79cea45eb31c2d08 Mon Sep 17 00:00:00 2001
From: you-n-g <you-n-g@users.noreply.github.com>
Date: Wed, 17 Jul 2024 18:07:30 +0800
Subject: [PATCH] Fix Logger context bug. (#83)

---
 rdagent/log/logger.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/rdagent/log/logger.py b/rdagent/log/logger.py
index 8a17768d..2c7c81ea 100644
--- a/rdagent/log/logger.py
+++ b/rdagent/log/logger.py
@@ -70,8 +70,10 @@ class RDAgentLog(SingletonBaseClass):
 
         # TODO: It may result in error in mutithreading or co-routine
         self._tag = self._tag + tag
-        yield
-        self._tag = self._tag[: -len(tag)]
+        try:
+            yield
+        finally:
+            self._tag = self._tag[: -len(tag)]
 
     def get_pids(self) -> str:
         """