fix: durability fixes

agent_core/core/hooks/usage.py

@@ -0,0 +1,121 @@
+"""Deliver async hooks in the request's context and drain them at shutdown.
+
+The loop belongs to a request, not to a long-lived model instance. Worker
+threads inherit that ownership through ContextVars, including session context.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+from contextvars import ContextVar
+import logging
+import threading
+
+_loop = ContextVar("usage_owner_loop", default=None)
+_request_pending = ContextVar("usage_request_pending", default=None)
+_pending: set = set()
+_lock = threading.Lock()
+_logger = logging.getLogger(__name__)
+
+
+def _track(future):
+    # A Task is bound to its loop. Track a thread-safe completion handle so
+    # application shutdown can also drain workers owned by other running loops.
+    completion = concurrent.futures.Future()
+    with _lock:
+        _pending.add(completion)
+
+    def done(item):
+        cancelled = item.cancelled()
+        error = None if cancelled else item.exception()
+        with _lock:
+            _pending.discard(completion)
+        if cancelled:
+            completion.cancel()
+        elif error is not None:
+            completion.set_exception(error)
+            _logger.error(
+                "Usage delivery failed",
+                exc_info=(type(error), error, error.__traceback__),
+            )
+        else:
+            completion.set_result(None)
+
+    future.add_done_callback(done)
+    return future
+
+
+def dispatch_usage(hook, event) -> None:
+    """Submit once; failed hooks are observable and never automatically retried."""
+
+    async def deliver():
+        await hook(event)
+
+    try:
+        running = asyncio.get_running_loop()
+    except RuntimeError:
+        running = None
+    owner = _loop.get() or running
+    if owner is None:
+        # Synchronous callers have no event loop to own delivery. Finish the
+        # hook before returning. Local reporters must be safe across loops.
+        asyncio.run(deliver())
+        return
+    if owner.is_closed():
+        raise RuntimeError("Usage request's event loop is closed")
+    future = (
+        owner.create_task(deliver())
+        if owner is running
+        else asyncio.run_coroutine_threadsafe(deliver(), owner)
+    )
+    _track(future)
+    request = _request_pending.get()
+    if request is not None:
+        request.append(future)
+
+
+async def _join(futures):
+    if futures:
+        await asyncio.gather(
+            *(
+                asyncio.shield(asyncio.wrap_future(f))
+                if isinstance(f, concurrent.futures.Future)
+                else asyncio.shield(f)
+                for f in futures
+            ),
+            return_exceptions=True,
+        )
+
+
+async def to_thread_with_usage(function, /, *args, **kwargs):
+    """Like to_thread, but delivery completes before a successful call returns.
+
+    Cancellation leaves the worker tracked: shutdown drains it before closing
+    the reporter. No request is resubmitted after an uncertain outcome.
+    """
+    request = []
+    owner_token = _loop.set(asyncio.get_running_loop())
+    pending_token = _request_pending.set(request)
+    try:
+        worker = _track(
+            asyncio.create_task(asyncio.to_thread(function, *args, **kwargs))
+        )
+        try:
+            return await asyncio.shield(worker)
+        finally:
+            if worker.done():
+                await _join(request)
+    finally:
+        _request_pending.reset(pending_token)
+        _loop.reset(owner_token)
+
+
+async def drain_usage() -> None:
+    """Wait for workers and hooks, including hooks submitted by finishing workers."""
+    while True:
+        with _lock:
+            pending = list(_pending)
+        if not pending:
+            return
+        await _join(pending)

agent_core/core/impl/image_gen/interface.py

@@ -18,6 +18,7 @@
 if TYPE_CHECKING:
     from agent_core.core.errors import ClassifiedError
 
+from agent_core.core.hooks.usage import dispatch_usage, to_thread_with_usage
 import asyncio
 import base64
 import io
@@ -197,9 +198,7 @@ def _report_usage_async(
     ) -> None:
         """Report image-generation usage if the hook is set (mirrors VLM).
 
-        Best-effort: scheduling onto the running loop can fail when invoked
-        from a worker thread without one, in which case the usage is dropped
-        with a warning rather than breaking generation.
+        Delivery uses the request's owning loop and is drained at shutdown.
         """
         if not self._report_usage:
             return
@@ -212,10 +211,7 @@ def _report_usage_async(
                 output_tokens=output_tokens,
                 cached_tokens=cached_tokens,
             )
-            if self._main_loop is not None:
-                self._main_loop.call_soon_threadsafe(
-                    lambda: asyncio.create_task(self._report_usage(event))
-                )
+            dispatch_usage(self._report_usage, event)
         except Exception as e:
             logger.warning(f"[IMAGE_GEN] Failed to report usage: {e}")
 
@@ -307,8 +303,7 @@ def generate_image(
 
     async def generate_image_async(self, **kwargs) -> List[str]:
         """Async wrapper — runs generate_image in a thread pool."""
-        loop = asyncio.get_event_loop()
-        return await loop.run_in_executor(None, lambda: self.generate_image(**kwargs))
+        return await to_thread_with_usage(self.generate_image, **kwargs)
 
     # ──────────────────────── OpenAI provider ────────────────────────────────
 

agent_core/core/impl/llm/interface.py

@@ -13,7 +13,7 @@
 
 from __future__ import annotations
 
-import asyncio
+from agent_core.core.hooks.usage import dispatch_usage, to_thread_with_usage
 import contextvars
 import re
 import time
@@ -614,13 +614,7 @@ def _report_usage_async(
                 output_tokens=output_tokens,
                 cached_tokens=cached_tokens,
             )
-            # Schedule the async hook on the event loop
-            try:
-                loop = asyncio.get_running_loop()
-                loop.create_task(self._report_usage(event))
-            except RuntimeError:
-                # No running event loop - create one for this call
-                asyncio.run(self._report_usage(event))
+            dispatch_usage(self._report_usage, event)
         except Exception as e:
             logger.warning(f"[LLM] Failed to report usage: {e}")
 
@@ -1260,7 +1254,7 @@ async def generate_response_async(
         )
         token = _active_session_key.set(usage_key)
         try:
-            return await asyncio.to_thread(
+            return await to_thread_with_usage(
                 self._generate_response_sync,
                 system_prompt,
                 user_prompt,
@@ -2335,7 +2329,7 @@ async def generate_response_with_session_async(
         self._begin_call(prompt_name=prompt_name, call_type=call_type, task_id=task_id)
         token = _active_session_key.set(f"{task_id}:{call_type}")
         try:
-            return await asyncio.to_thread(
+            return await to_thread_with_usage(
                 self._generate_response_with_session_sync,
                 task_id,
                 call_type,

agent_core/core/impl/video_gen/interface.py

@@ -24,6 +24,7 @@
 if TYPE_CHECKING:
     from agent_core.core.errors import ClassifiedError
 
+from agent_core.core.hooks.usage import dispatch_usage, to_thread_with_usage
 import asyncio
 import base64
 import json
@@ -244,10 +245,7 @@ def _report_usage_async(
                 output_tokens=output_tokens,
                 cached_tokens=cached_tokens,
             )
-            if self._main_loop is not None:
-                self._main_loop.call_soon_threadsafe(
-                    lambda: asyncio.create_task(self._report_usage(event))
-                )
+            dispatch_usage(self._report_usage, event)
         except Exception as e:
             logger.warning(f"[VIDEO_GEN] Failed to report usage: {e}")
 
@@ -396,8 +394,7 @@ def generate_video(
 
     async def generate_video_async(self, **kwargs) -> List[str]:
         """Async wrapper — runs generate_video in a thread pool."""
-        loop = asyncio.get_running_loop()
-        return await loop.run_in_executor(None, lambda: self.generate_video(**kwargs))
+        return await to_thread_with_usage(self.generate_video, **kwargs)
 
     # ──────────────────────── OpenAI / Sora 2 ────────────────────────────────