Work around type hinting problems with redis.asyncio

Typing for the asyncio redis client is messed up, and functions
return `T | Awaitable[T]` instead of `T`. Add a workaround
wrapper that asserts that we are in the second case.

See:

  redis/redis-py#3619

for an in-progress upstream fix.

Author: owtaylor

beeai/agents/backport_agent.py

@@ -32,7 +32,7 @@
 from tools.text import CreateTool, InsertTool, StrReplaceTool, ViewTool
 from tools.wicked_git import GitLogSearchTool, GitPatchCreationTool
 from triage_agent import BackportData, ErrorData
-from utils import check_subprocess, get_agent_execution_config, mcp_tools, redis_client, post_private_jira_comment
+from utils import fix_await, check_subprocess, get_agent_execution_config, mcp_tools, redis_client, post_private_jira_comment
 
 logger = logging.getLogger(__name__)
 
@@ -234,7 +234,7 @@ class Task(BaseModel):
 
             while True:
                 logger.info("Waiting for tasks from backport_queue (timeout: 30s)...")
-                element = await redis.brpop("backport_queue", timeout=30)
+                element = await fix_await(redis.brpop(["backport_queue"], timeout=30))
                 if element is None:
                     logger.info("No tasks received, continuing to wait...")
                     continue
@@ -256,13 +256,13 @@ async def retry(task, error):
                             f"Task failed (attempt {task.attempts}/{max_retries}), "
                             f"re-queuing for retry: {backport_data.jira_issue}"
                         )
-                        await redis.lpush("backport_queue", task.model_dump_json())
+                        await fix_await(redis.lpush("backport_queue", task.model_dump_json()))
                     else:
                         logger.error(
                             f"Task failed after {max_retries} attempts, "
                             f"moving to error list: {backport_data.jira_issue}"
                         )
-                        await redis.lpush("error_list", error)
+                        await fix_await(redis.lpush("error_list", error))
 
                 try:
                     logger.info(f"Starting backport processing for {backport_data.jira_issue}")
@@ -298,7 +298,7 @@ async def retry(task, error):
                     rmtree(local_clone)
                     if state.backport_data.success:
                         logger.info(f"Backport successful for {backport_data.jira_issue}, " f"adding to completed list")
-                        await redis.lpush("completed_backport_list", state.backport_data.model_dump_json())
+                        await redis.lpush("completed_backport_list", output.model_dump_json())
                     else:
                         logger.warning(f"Backport failed for {backport_data.jira_issue}: {state.backport_data.error}")
                         await retry(task, state.backport_data.error)

beeai/agents/rebase_agent.py

@@ -28,7 +28,7 @@
 from tools.specfile import AddChangelogEntryTool
 from tools.text import CreateTool, InsertTool, StrReplaceTool, ViewTool
 from triage_agent import RebaseData, ErrorData
-from utils import get_agent_execution_config, mcp_tools, redis_client, run_tool, post_private_jira_comment
+from utils import fix_await, get_agent_execution_config, mcp_tools, redis_client, run_tool, post_private_jira_comment
 
 logger = logging.getLogger(__name__)
 
@@ -237,7 +237,7 @@ class Task(BaseModel):
 
             while True:
                 logger.info("Waiting for tasks from rebase_queue (timeout: 30s)...")
-                element = await redis.brpop("rebase_queue", timeout=30)
+                element = await fix_await(redis.brpop(["rebase_queue"], timeout=30))
                 if element is None:
                     logger.info("No tasks received, continuing to wait...")
                     continue
@@ -260,13 +260,13 @@ async def retry(task, error):
                             f"Task failed (attempt {task.attempts}/{max_retries}), "
                             f"re-queuing for retry: {rebase_data.jira_issue}"
                         )
-                        await redis.lpush("rebase_queue", task.model_dump_json())
+                        await fix_await(redis.lpush("rebase_queue", task.model_dump_json()))
                     else:
                         logger.error(
                             f"Task failed after {max_retries} attempts, "
                             f"moving to error list: {rebase_data.jira_issue}"
                         )
-                        await redis.lpush("error_list", error)
+                        await fix_await(redis.lpush("error_list", error))
 
                 try:
                     logger.info(f"Starting rebase processing for {rebase_data.jira_issue}")
@@ -296,7 +296,7 @@ async def retry(task, error):
                 else:
                     if state.rebase_result.success:
                         logger.info(f"Rebase successful for {rebase_data.jira_issue}, " f"adding to completed list")
-                        await redis.lpush("completed_rebase_list", state.rebase_result.model_dump_json())
+                        await fix_await(redis.lpush("completed_rebase_list", state.rebase_result.model_dump_json()))
                     else:
                         logger.warning(f"Rebase failed for {rebase_data.jira_issue}: {state.rebase_result.error}")
                         await retry(task, state.rebase_result.error)

beeai/agents/triage_agent.py

@@ -24,7 +24,7 @@
 from tools.commands import RunShellCommandTool
 from tools.patch_validator import PatchValidatorTool
 from tools.version_mapper import VersionMapperTool
-from utils import get_agent_execution_config, mcp_tools, redis_client, post_private_jira_comment
+from utils import fix_await, get_agent_execution_config, mcp_tools, redis_client, post_private_jira_comment
 
 logger = logging.getLogger(__name__)
 
@@ -303,7 +303,7 @@ class Task(BaseModel):
 
             while True:
                 logger.info("Waiting for tasks from triage_queue (timeout: 30s)...")
-                element = await redis.brpop("triage_queue", timeout=30)
+                element = await fix_await(redis.brpop(["triage_queue"], timeout=30))
                 if element is None:
                     logger.info("No tasks received, continuing to wait...")
                     continue
@@ -322,12 +322,12 @@ async def retry(task, error):
                             f"Task failed (attempt {task.attempts}/{max_retries}), "
                             f"re-queuing for retry: {input.issue}"
                         )
-                        await redis.lpush("triage_queue", task.model_dump_json())
+                        await fix_await(redis.lpush("triage_queue", task.model_dump_json()))
                     else:
                         logger.error(
                             f"Task failed after {max_retries} attempts, " f"moving to error list: {input.issue}"
                         )
-                        await redis.lpush("error_list", error)
+                        await fix_await(redis.lpush("error_list", error))
 
                 try:
                     logger.info(f"Starting triage processing for {input.issue}")
@@ -353,21 +353,21 @@ async def retry(task, error):
                     if output.resolution == Resolution.REBASE:
                         logger.info(f"Triage resolved as REBASE for {input.issue}, " f"adding to rebase queue")
                         task = Task(metadata=output.data.model_dump())
-                        await redis.lpush("rebase_queue", task.model_dump_json())
+                        await fix_await(redis.lpush("rebase_queue", task.model_dump_json()))
                     elif output.resolution == Resolution.BACKPORT:
                         logger.info(f"Triage resolved as BACKPORT for {input.issue}, " f"adding to backport queue")
                         task = Task(metadata=output.data.model_dump())
-                        await redis.lpush("backport_queue", task.model_dump_json())
+                        await fix_await(redis.lpush("backport_queue", task.model_dump_json()))
                     elif output.resolution == Resolution.CLARIFICATION_NEEDED:
                         logger.info(
                             f"Triage resolved as CLARIFICATION_NEEDED for {input.issue}, "
                             f"adding to clarification needed queue"
                         )
                         task = Task(metadata=output.data.model_dump())
-                        await redis.lpush("clarification_needed_queue", task.model_dump_json())
+                        await fix_await(redis.lpush("clarification_needed_queue", task.model_dump_json()))
                     elif output.resolution == Resolution.NO_ACTION:
                         logger.info(f"Triage resolved as NO_ACTION for {input.issue}, " f"adding to no action list")
-                        await redis.lpush("no_action_list", output.data.model_dump_json())
+                        await fix_await(redis.lpush("no_action_list", output.data.model_dump_json()))
                     elif output.resolution == Resolution.ERROR:
                         logger.warning(f"Triage resolved as ERROR for {input.issue}, retrying")
                         await retry(task, output.data.model_dump_json())

beeai/agents/utils.py

@@ -1,11 +1,12 @@
 import asyncio
+import inspect
 import logging
 import os
 import shlex
 import subprocess
 from contextlib import asynccontextmanager
 from pathlib import Path
-from typing import Any, AsyncGenerator, Callable, Tuple
+from typing import Any, AsyncGenerator, Awaitable, Callable, TypeVar, Tuple
 
 import redis.asyncio as redis
 from mcp import ClientSession
@@ -96,6 +97,26 @@ async def redis_client(redis_url: str) -> AsyncGenerator[redis.Redis, None]:
         await client.aclose()
 
 
+T = TypeVar("T")
+
+async def fix_await(v: T | Awaitable[T]) -> T:
+    """
+    Work around typing problems in the asyncio redis client.
+
+    Typing for the asyncio redis client is messed up, and functions
+    return `T | Awaitable[T]` instead of `T`. This function
+    fixes the type error by asserting that the value is awaitable
+    before awaiting it.
+
+    For a proper fix, see: https://github.com/redis/redis-py/pull/3619
+
+
+    Usage: `await fixAwait(redis.get("key"))`
+    """
+    assert inspect.isawaitable(v)
+    return await v
+
+
 @asynccontextmanager
 async def mcp_tools(
     sse_url: str, filter: Callable[[str], bool] | None = None