diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
index 7d9641a..4de0c51 100644
--- a/.devcontainer/devcontainer.json
+++ b/.devcontainer/devcontainer.json
@@ -2,7 +2,7 @@
 // README at: https://github.com/devcontainers/templates/tree/main/src/docker-existing-dockerfile
 {
     "name": "dev",
-    "image": "ghcr.io/ten-framework/ten_agent_build:0.3.5",
+    "image": "ghcr.io/ten-framework/ten_agent_build:0.3.6",
     "customizations": {
         "vscode": {
             "extensions": [
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 2d80597..48255ba 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -24,7 +24,7 @@ jobs:
   ci:
     runs-on: ubuntu-latest
     container:
-      image: ghcr.io/ten-framework/ten_agent_build:0.3.5
+      image: ghcr.io/ten-framework/ten_agent_build:0.3.6
     steps:
       - uses: actions/checkout@v4
         with:
diff --git a/tests/test_5xx.py b/tests/test_5xx.py
index d7b4433..3db8558 100644
--- a/tests/test_5xx.py
+++ b/tests/test_5xx.py
@@ -22,35 +22,26 @@ class ExtensionTester5xx(ExtensionTester):
     def __init__(self):
         super().__init__()
         self.thread = None
-        self.event = threading.Event()
 
     def on_cmd(self, ten_env: TenEnvTester, cmd: Cmd) -> None:
         print(f"on_cmd name {cmd.get_name()}")
         ten_env.return_result(CmdResult.create(StatusCode.ERROR), cmd)
 
-        event_got = self.event.wait(timeout=5)
-        if not event_got:  # timeout
-            print("test error")
-        else:
-            ten_env.stop_test()
-
     def on_start(self, ten_env: TenEnvTester) -> None:
-
-        self.thread = threading.Thread(target=self._async_request, args=[])
+        self.thread = threading.Thread(
+            target=self._async_test, args=[ten_env])
         self.thread.start()
 
         ten_env.on_start_done()
 
-    def _async_request(self) -> None:
+    def _async_test(self, ten_env: TenEnvTester) -> None:
         property_json = {"num": 1, "str": "111"}
         r = httpx.post("http://127.0.0.1:8888/cmd/abc",
                        json=property_json, timeout=5)
         print(r)
 
         if r.status_code >= 500:
-            # TODO: stop test directly once albe to call it in separate thread
-            #   ten_env.stop_test()
-            self.event.set()
+            ten_env.stop_test()
 
 
 def test_5xx():
diff --git a/tests/test_basic.py b/tests/test_basic.py
index a0153ef..5b95879 100644
--- a/tests/test_basic.py
+++ b/tests/test_basic.py
@@ -20,35 +20,27 @@ class ExtensionTesterBasic(ExtensionTester):
     def __init__(self):
         super().__init__()
         self.thread = None
-        self.event = threading.Event()
 
     def on_cmd(self, ten_env: TenEnvTester, cmd: Cmd) -> None:
         print(f"on_cmd name {cmd.get_name()}")
         ten_env.return_result(CmdResult.create(StatusCode.OK), cmd)
 
-        event_got = self.event.wait(timeout=5)
-        if not event_got:  # timeout
-            print("test error")
-        else:
-            ten_env.stop_test()
-
     def on_start(self, ten_env: TenEnvTester) -> None:
 
-        self.thread = threading.Thread(target=self._async_request, args=[])
+        self.thread = threading.Thread(
+            target=self._async_test, args=[ten_env])
         self.thread.start()
 
         ten_env.on_start_done()
 
-    def _async_request(self) -> None:
+    def _async_test(self, ten_env: TenEnvTester) -> None:
         property_json = {"num": 1, "str": "111"}
         r = httpx.post("http://127.0.0.1:8888/cmd/abc",
                        json=property_json, timeout=5)
-        print(r)   
+        print(r)
 
         if r.status_code == httpx.codes.OK:
-            # TODO: stop test directly once albe to call it in separate thread
-            #   ten_env.stop_test()
-            self.event.set()
+            ten_env.stop_test()
 
 
 def test_basic():
diff --git a/tests/test_set_property.py b/tests/test_set_property.py
index 34a3fbb..7f105ce 100644
--- a/tests/test_set_property.py
+++ b/tests/test_set_property.py
@@ -20,35 +20,27 @@ class ExtensionTesterSetProperty(ExtensionTester):
     def __init__(self):
         super().__init__()
         self.thread = None
-        self.event = threading.Event()
 
     def on_cmd(self, ten_env: TenEnvTester, cmd: Cmd) -> None:
         print(f"on_cmd name {cmd.get_name()}")
         ten_env.return_result(CmdResult.create(StatusCode.OK), cmd)
 
-        event_got = self.event.wait(timeout=5)
-        if not event_got:  # timeout
-            print("test error")
-        else:
-            ten_env.stop_test()
-
     def on_start(self, ten_env: TenEnvTester) -> None:
 
-        self.thread = threading.Thread(target=self._async_request, args=[])
+        self.thread = threading.Thread(
+            target=self._async_test, args=[ten_env])
         self.thread.start()
 
         ten_env.on_start_done()
 
-    def _async_request(self) -> None:
+    def _async_test(self, ten_env: TenEnvTester) -> None:
         property_json = {"num": 1, "str": "111"}
         r = httpx.post("http://127.0.0.1:8899/cmd/abc",
                        json=property_json, timeout=5)
         print(r)
 
         if r.status_code == httpx.codes.OK:
-            # TODO: stop test directly once albe to call it in separate thread
-            #   ten_env.stop_test()
-            self.event.set()
+            ten_env.stop_test()
 
 
 def test_set_property():