fix: Correct HybridBackgroundTasks interface and behavior

- Fix add_task method to be synchronous (matching FastAPI BackgroundTasks interface)
- When use_docket=True: Schedule tasks directly in Docket's Redis queue
- When use_docket=False: Use FastAPI's background tasks
- Update tests to reflect correct behavior
- Fix mypy type checking issues
- Maintain proper separation between Docket and FastAPI task execution

Author: abrookins

agent_memory_server/dependencies.py

@@ -1,3 +1,4 @@
+import asyncio
 from collections.abc import Callable
 from typing import Any
 
@@ -13,26 +14,43 @@
 class HybridBackgroundTasks(BackgroundTasks):
     """A BackgroundTasks implementation that can use either Docket or FastAPI background tasks."""
 
-    async def add_task(
-        self, func: Callable[..., Any], *args: Any, **kwargs: Any
-    ) -> None:
+    def add_task(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> None:
         """Run tasks either directly, through Docket, or through FastAPI background tasks"""
         logger.info("Adding task to background tasks...")
 
         if settings.use_docket:
             logger.info("Scheduling task through Docket")
+            # Schedule task directly in Docket's Redis queue
+            self._schedule_docket_task(func, *args, **kwargs)
+        else:
+            logger.info("Using FastAPI background tasks")
+            # Use FastAPI's background tasks directly
+            super().add_task(func, *args, **kwargs)
+
+    def _schedule_docket_task(
+        self, func: Callable[..., Any], *args: Any, **kwargs: Any
+    ) -> None:
+        """Schedule a task in Docket's Redis queue"""
+
+        async def schedule_task():
             from docket import Docket
 
             async with Docket(
                 name=settings.docket_name,
                 url=settings.redis_url,
             ) as docket:
-                # Schedule task through Docket
+                # Schedule task in Docket's queue
                 await docket.add(func)(*args, **kwargs)
-        else:
-            logger.info("Using FastAPI background tasks")
-            # Use FastAPI's background tasks
-            super().add_task(func, *args, **kwargs)
+
+        # Run the async scheduling operation
+        try:
+            # Try to get the current event loop
+            loop = asyncio.get_running_loop()
+            # If we're in an async context, create a task
+            loop.create_task(schedule_task())
+        except RuntimeError:
+            # If no event loop is running, run it synchronously
+            asyncio.run(schedule_task())
 
 
 # Backwards compatibility alias

tests/test_dependencies.py

@@ -1,5 +1,6 @@
 """Tests for the dependencies module, particularly HybridBackgroundTasks."""
 
+import asyncio
 from unittest.mock import Mock, patch
 
 import pytest
@@ -24,8 +25,7 @@ def teardown_method(self):
         """Restore original use_docket setting."""
         settings.use_docket = self.original_use_docket
 
-    @pytest.mark.asyncio
-    async def test_add_task_with_fastapi_background_tasks(self):
+    def test_add_task_with_fastapi_background_tasks(self):
         """Test that tasks are added to FastAPI background tasks when use_docket=False."""
         settings.use_docket = False
 
@@ -35,7 +35,7 @@ def test_func(arg1, arg2=None):
             return f"test_func called with {arg1}, {arg2}"
 
         # Add task
-        await bg_tasks.add_task(test_func, "hello", arg2="world")
+        bg_tasks.add_task(test_func, "hello", arg2="world")
 
         # Verify task was added to FastAPI background tasks
         assert len(bg_tasks.tasks) == 1
@@ -71,8 +71,14 @@ async def execution():
             async def test_func(arg1, arg2=None):
                 return f"test_func called with {arg1}, {arg2}"
 
-            # Add task
-            await bg_tasks.add_task(test_func, "hello", arg2="world")
+            # Add task - this should schedule directly in Docket
+            bg_tasks.add_task(test_func, "hello", arg2="world")
+
+            # Give the async task a moment to complete
+            await asyncio.sleep(0.1)
+
+            # When use_docket=True, tasks should NOT be added to FastAPI background tasks
+            assert len(bg_tasks.tasks) == 0
 
             # Verify Docket was used
             mock_docket_class.assert_called_once_with(
@@ -81,8 +87,7 @@ async def test_func(arg1, arg2=None):
             )
             mock_docket_instance.add.assert_called_once_with(test_func)
 
-    @pytest.mark.asyncio
-    async def test_add_task_logs_correctly_for_fastapi(self):
+    def test_add_task_logs_correctly_for_fastapi(self):
         """Test that FastAPI background tasks work without errors."""
         settings.use_docket = False
 
@@ -92,7 +97,7 @@ def test_func():
             pass
 
         # Should not raise any errors
-        await bg_tasks.add_task(test_func)
+        bg_tasks.add_task(test_func)
 
         # Verify task was added to FastAPI background tasks
         assert len(bg_tasks.tasks) == 1
@@ -122,7 +127,10 @@ def test_func():
                 pass
 
             # Should not raise any errors
-            await bg_tasks.add_task(test_func)
+            bg_tasks.add_task(test_func)
+
+            # Give the async task a moment to complete
+            await asyncio.sleep(0.1)
 
             # Verify Docket was called
             mock_docket_instance.add.assert_called_once_with(test_func)
@@ -178,7 +186,7 @@ def test_func():
 
         # Test with use_docket=False
         settings.use_docket = False
-        await bg_tasks.add_task(test_func)
+        bg_tasks.add_task(test_func)
         assert len(bg_tasks.tasks) == 1
 
         # Clear tasks and test with use_docket=True
@@ -199,14 +207,18 @@ async def execution():
 
             mock_docket_instance.add.return_value = mock_task_callable
 
-            await bg_tasks.add_task(test_func)
+            bg_tasks.add_task(test_func)
+
+            # Give the async task a moment to complete
+            await asyncio.sleep(0.1)
 
             # Should not add to FastAPI tasks when use_docket=True
             assert len(bg_tasks.tasks) == 0
             # Should have called Docket
             mock_docket_class.assert_called_once()
 
-    def test_uses_correct_docket_settings(self):
+    @pytest.mark.asyncio
+    async def test_uses_correct_docket_settings(self):
         """Test that HybridBackgroundTasks uses the correct docket name and URL from settings."""
         settings.use_docket = True
 
@@ -226,12 +238,10 @@ async def execution():
 
             bg_tasks = HybridBackgroundTasks()
 
-            async def run_test():
-                await bg_tasks.add_task(lambda: None)
-
-            import asyncio
+            bg_tasks.add_task(lambda: None)
 
-            asyncio.run(run_test())
+            # Give the async task a moment to complete
+            await asyncio.sleep(0.1)
 
             # Verify Docket was initialized with correct settings
             mock_docket_class.assert_called_once_with(