diff --git a/pyproject.toml b/pyproject.toml
index 273ead1..e58a094 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "synth-sdk"
-version = "0.2.68"
+version = "0.2.69"
 description = ""
 authors = [{name = "Synth AI", email = "josh@usesynth.ai"}]
 license = {text = "MIT"}
diff --git a/setup.py b/setup.py
index b7a9e43..4e52437 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name="synth-sdk",
-    version="0.2.68",
+    version="0.2.69",
     packages=find_packages(),
     install_requires=[
         "opentelemetry-api",
diff --git a/synth_sdk/tracing/upload.py b/synth_sdk/tracing/upload.py
index e2a7f11..d020453 100644
--- a/synth_sdk/tracing/upload.py
+++ b/synth_sdk/tracing/upload.py
@@ -182,20 +182,17 @@ def upload(dataset: Dataset, traces: List[SystemTrace]=[], verbose: bool = False
     async def upload_wrapper(dataset, traces, verbose, show_payload):
         result = await upload_helper(dataset, traces, verbose, show_payload)
         return result
-    
+
+    # If we're in an async context (event loop is running)
     if is_event_loop_running():
         logging.info("Event loop is already running")
-        task = asyncio.create_task(upload_wrapper(dataset, traces, verbose, show_payload))
-        # Wait for the task if called from an async function
-        if asyncio.current_task():
-            return task  # Returning the task to be awaited if in async context
-        else:
-            # Run task synchronously by waiting for it to finish if in sync context
-            return asyncio.get_event_loop().run_until_complete(task)
-        
+        # Return the coroutine directly for async contexts
+        return upload_helper(dataset, traces, verbose, show_payload)
     else:
+        # In sync context, run the coroutine and return the result
         logging.info("Event loop is not running")
-        return asyncio.run(upload_wrapper(dataset, traces, verbose, show_payload))
+        return asyncio.run(upload_helper(dataset, traces, verbose, show_payload))
+
 
 async def upload_helper(dataset: Dataset, traces: List[SystemTrace]=[], verbose: bool = False, show_payload: bool = False):
     """Upload all system traces and dataset to the server."""
diff --git a/testing/ai_agent_async.py b/testing/ai_agent_async.py
index 3dfa624..e213eb7 100644
--- a/testing/ai_agent_async.py
+++ b/testing/ai_agent_async.py
@@ -128,7 +128,7 @@ async def run_test():
         # Upload traces
         try:
             logger.info("Attempting to upload traces")
-            response, _ = await upload(dataset=dataset, verbose=True)
+            response, payload, dataset, traces = await upload(dataset=dataset, verbose=True)
             logger.info("Upload successful!")
             print("Upload successful!")
         except Exception as e:
diff --git a/testing/ai_agent_sync.py b/testing/ai_agent_sync.py
index 0f2f082..19f0000 100644
--- a/testing/ai_agent_sync.py
+++ b/testing/ai_agent_sync.py
@@ -128,7 +128,7 @@ async def run_test():
         # Upload traces
         try:
             logger.info("Attempting to upload traces")
-            response, _ = await upload(dataset=dataset, verbose=True)
+            response, payload, dataset, traces= await upload(dataset=dataset, verbose=True)
             logger.info("Upload successful!")
             print("Upload successful!")
         except Exception as e:
diff --git a/testing/upload_payload_test.py b/testing/upload_payload_test.py
index 3655486..1165f9a 100644
--- a/testing/upload_payload_test.py
+++ b/testing/upload_payload_test.py
@@ -153,7 +153,7 @@ async def test_upload_async(mock_send_system_traces):
     # Upload traces
     logger.info("Attempting to upload traces, async version")
     # TODO: uploads traces directly from list of traces (override the event_store)
-    response, payload = await upload(dataset=dataset, verbose=True, show_payload=True)
+    response, payload, dataset, traces = await upload(dataset=dataset, verbose=True, show_payload=True)
     logger.info("Upload successful!")
 
     # Pytest assertion
@@ -174,7 +174,7 @@ def test_upload_sync(mock_send_system_traces):
     # Upload traces
     logger.info("Attempting to upload traces, non-async version")
     # TODO: uploads traces directly from list of traces (override the event_store)
-    response, payload = upload(dataset=dataset, verbose=True, show_payload=True)
+    response, payload, dataset, traces = upload(dataset=dataset, verbose=True, show_payload=True)
     logger.info("Upload successful!")
 
     # Pytest assertion