diff --git a/langfuse/callback/langchain.py b/langfuse/callback/langchain.py
index 44d15c556..8749ef914 100644
--- a/langfuse/callback/langchain.py
+++ b/langfuse/callback/langchain.py
@@ -897,13 +897,14 @@ def on_llm_error(
     ) -> Any:
         try:
             self._log_debug_event("on_llm_error", run_id, parent_run_id, error=error)
-            self.runs[run_id] = self.runs[run_id].end(
-                status_message=str(error),
-                level="ERROR",
-                version=self.version,
-                input=kwargs.get("inputs"),
-            )
-            self._update_trace_and_remove_state(run_id, parent_run_id, error)
+            if run_id in self.runs:
+                self.runs[run_id] = self.runs[run_id].end(
+                    status_message=str(error),
+                    level="ERROR",
+                    version=self.version,
+                    input=kwargs.get("inputs"),
+                )
+                self._update_trace_and_remove_state(run_id, parent_run_id, error)
 
         except Exception as e:
             self.log.exception(e)
diff --git a/tests/test_core_sdk.py b/tests/test_core_sdk.py
index e8ac49f9a..2278b0bbc 100644
--- a/tests/test_core_sdk.py
+++ b/tests/test_core_sdk.py
@@ -1500,6 +1500,7 @@ def faulty_mask_func(data):
     langfuse = Langfuse(debug=True, mask=faulty_mask_func)
 
     trace = langfuse.trace(name="test_trace", input={"sensitive": "data"})
+    sleep(0.1)
     trace.update(output={"more": "sensitive"})
     langfuse.flush()
 
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index 715c10e95..d97447f13 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -438,6 +438,7 @@ def test_llama_index_dataset():
     assert run.name == run_name
     assert len(run.dataset_run_items) == 1
     assert run.dataset_run_items[0].dataset_run_id == run.id
+    time.sleep(1)
 
     trace = get_api().trace.get(handler.get_trace_id())