Fix result.log() failing with "Event loop is closed" in plain scripts

tests/test_api_client.py

@@ -2,7 +2,9 @@
 import json
 import logging
 import os
+import threading
 import unittest
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from unittest.mock import MagicMock, Mock, patch
 
 import matplotlib.pyplot as plt
@@ -22,7 +24,7 @@
     MissingModelIdError,
     ValidMindAuthError,
 )
-from validmind.utils import md_to_html
+from validmind.utils import md_to_html, run_async
 from validmind.vm_models.figure import Figure
 from validmind.vm_models.result.utils import update_metadata
 
@@ -559,6 +561,42 @@ def test_ping_refreshes_before_request_and_retries_on_401(
         mock_ensure.assert_any_call(force=True)
 
 
+class TestAsyncSessionAcrossEventLoops(unittest.TestCase):
+    """Outside a notebook, run_async uses a new event loop (asyncio.run) per call."""
+
+    def setUp(self):
+        class Handler(BaseHTTPRequestHandler):
+            def do_POST(self):
+                self.rfile.read(int(self.headers.get("Content-Length", 0)))
+                body = b'{"ok": true}'
+                self.send_response(200)
+                self.send_header("Content-Type", "application/json")
+                self.send_header("Content-Length", str(len(body)))
+                self.end_headers()
+                self.wfile.write(body)
+
+            def log_message(self, *args):
+                pass
+
+        self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+        threading.Thread(target=self.server.serve_forever, daemon=True).start()
+        self.addCleanup(self.server.server_close)
+        self.addCleanup(self.server.shutdown)
+        patcher = patch.object(
+            api_client, "_api_host", f"http://127.0.0.1:{self.server.server_port}/"
+        )
+        patcher.start()
+        self.addCleanup(patcher.stop)
+        self.addCleanup(api_client._invalidate_async_session)
+        api_client._invalidate_async_session()
+
+    @patch("validmind.api_client._ensure_fresh_oidc_token")
+    def test_post_works_on_each_new_event_loop(self, mock_ensure):
+        # e.g. vm.init() followed by result.log() in a plain Python script
+        self.assertEqual(run_async(api_client._post, "endpoint"), {"ok": True})
+        self.assertEqual(run_async(api_client._post, "endpoint"), {"ok": True})
+
+
 class TestAPIClientOIDC(unittest.TestCase):
     """OIDC device-flow authentication via vm.init()."""
 

validmind/api_client.py

@@ -50,6 +50,7 @@
 _oidc_refresh_lock = threading.Lock()
 
 __api_session: Optional[aiohttp.ClientSession] = None
+__api_session_loop: Optional[asyncio.AbstractEventLoop] = None
 
 
 def _invalidate_async_session() -> None:
@@ -72,23 +73,10 @@ def _invalidate_async_session() -> None:
 @atexit.register
 def _close_session():
     """Closes the async client session at exit."""
-    if __api_session and not __api_session.closed:
-        try:
-            loop = asyncio.get_event_loop()
-            if loop.is_running():
-                loop.create_task(__api_session.close())
-            else:
-                loop.run_until_complete(__api_session.close())
-        except RuntimeError as e:
-            # ignore RuntimeError when closing the session from the main thread
-            if "no current event loop in thread" in str(e):
-                pass
-            elif "Event loop is closed" in str(e):
-                pass
-            else:
-                raise e
-        except Exception as e:
-            logger.exception("Error closing aiohttp session at exit: %s", e)
+    try:
+        _invalidate_async_session()
+    except Exception as e:
+        logger.exception("Error closing aiohttp session at exit: %s", e)
 
 
 def get_api_host() -> Optional[str]:
@@ -120,15 +108,22 @@ def _get_api_headers() -> Dict[str, str]:
 
 
 def _get_session() -> aiohttp.ClientSession:
-    """Initializes the async client session."""
-    global __api_session
+    """Initializes the async client session for the running event loop."""
+    global __api_session, __api_session_loop
+
+    # An aiohttp session only works on the loop it was created on, and outside a
+    # notebook run_async starts a new loop (asyncio.run) for every call.
+    loop = asyncio.get_running_loop()
+    if __api_session_loop is not loop:
+        _invalidate_async_session()
 
     if not __api_session or __api_session.closed:
         __api_session = aiohttp.ClientSession(
             headers=_get_api_headers(),
             timeout=aiohttp.ClientTimeout(total=int(os.getenv("VM_API_TIMEOUT", 30))),
             trust_env=True,
         )
+        __api_session_loop = loop
 
     return __api_session