fix: drain non-empty buffers on ExtractionWorker shutdown (#78)

Short conversations that haven't hit the 800-token threshold would
silently lose their buffered extraction requests when shutdown()
cancelled the pending _debounced_process timers. Fix collects all keys
with pending requests before cancelling timers, then awaits
asyncio.gather over _process(key) for each — using the existing
semaphore — before clearing remaining state.

Author: laxmanclo

vektori/utils/async_worker.py

@@ -155,7 +155,30 @@ async def _extract_one(req: ExtractionRequest) -> None:
         await asyncio.gather(*[_extract_one(r) for r in buf.requests])
 
     async def shutdown(self, timeout: float = 30.0) -> None:
-        """Cancel pending timers and wait for any in-flight tasks."""
+        """Drain buffered extractions, then cancel pending timers and clear state."""
+        # Collect keys with buffered-but-not-yet-fired requests before touching timers
+        keys_to_drain = [k for k, buf in self._buffers.items() if buf.requests]
+
+        # Cancel timers for those keys first so _debounced_process can't double-fire
+        for key in keys_to_drain:
+            t = self._timers.get(key)
+            if t and not t.done():
+                t.cancel()
+
+        # Drain remaining buffers under the existing concurrency semaphore
+        if keys_to_drain:
+            try:
+                await asyncio.wait_for(
+                    asyncio.gather(
+                        *[self._process(key) for key in keys_to_drain],
+                        return_exceptions=True,
+                    ),
+                    timeout=timeout,
+                )
+            except asyncio.TimeoutError:
+                logger.warning("Extraction worker shutdown timed out while draining buffers")
+
+        # Cancel any remaining timers (e.g. in-flight tasks not in keys_to_drain)
         pending = [t for t in self._timers.values() if not t.done()]
         for t in pending:
             t.cancel()
@@ -166,6 +189,6 @@ async def shutdown(self, timeout: float = 30.0) -> None:
                     timeout=timeout,
                 )
             except asyncio.TimeoutError:
-                logger.warning("Extraction worker shutdown timed out")
+                logger.warning("Extraction worker shutdown timed out waiting for in-flight tasks")
         self._buffers.clear()
         self._timers.clear()