Add async task multiplexing for concurrent task processing

.pre-commit-config.yaml

@@ -38,3 +38,11 @@ repos:
         stages: [pre-push]
         always_run: true
         pass_filenames: false
+      - id: autobump-version
+        name: Auto-bump patch version for changed packages
+        entry: python scripts/autobump-version.py
+        language: python
+        additional_dependencies: [gitpython, tomlkit]
+        stages: [pre-push]
+        always_run: true
+        pass_filenames: false

compose.yaml

@@ -7,6 +7,8 @@ x-beeai-env: &beeai-env
   COLLECTOR_ENDPOINT: http://otel-collector:4318/v1/traces
   GIT_REPO_BASEPATH: /git-repos
   MAX_RETRIES: 3
+  MAX_CONCURRENT_TASKS: ${MAX_CONCURRENT_TASKS:-1}
+  LOG_BUFFER_SIZE: ${LOG_BUFFER_SIZE:-0}
   DRY_RUN: ${DRY_RUN:-false}
   JIRA_DRY_RUN: ${JIRA_DRY_RUN:-false}
   JIRA_ALLOW_STATUS_CHANGES: ${JIRA_ALLOW_STATUS_CHANGES:-false}

openshift/configmap-agents-env.yml

@@ -1,6 +1,8 @@
 apiVersion: v1
 data:
   MAX_RETRIES: "3"
+  MAX_CONCURRENT_TASKS: "1"
+  LOG_BUFFER_SIZE: "0"
   GIT_REPO_BASEPATH: /git-repos
   # The maximum number of retries for a single step in the agent execution.
   BEEAI_MAX_RETRIES_PER_STEP: "5"

scripts/autobump-version.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+"""Pre-push hook to auto-bump patch version in subpackages with code changes.
+
+Compares the commits being pushed against the remote base. For each configured
+package, if there are code changes but the version hasn't been bumped yet,
+bumps the patch version, commits, and aborts the push so the user can push again
+with the bump included.
+"""
+
+import sys
+from pathlib import Path
+
+import git
+import tomlkit
+
+PACKAGES = [
+    Path("ymir/tools"),
+    Path("ymir/common"),
+]
+
+
+def get_base(repo: git.Repo, local_sha: str, remote_sha: str) -> git.Commit | None:
+    if remote_sha != "0" * 40:
+        return repo.commit(remote_sha)
+    merge_base = repo.merge_base(local_sha, "main")
+    return merge_base[0] if merge_base else None
+
+
+def get_version_at_ref(commit: git.Commit, pyproject: Path) -> str | None:
+    try:
+        blob = commit.tree / str(pyproject)
+    except KeyError:
+        return None
+    data = tomlkit.loads(blob.data_stream.read().decode())
+    return data.get("project", {}).get("version")
+
+
+def has_code_changes(repo: git.Repo, directory: Path, local: git.Commit, base: git.Commit) -> bool:
+    diff = base.diff(local, paths=[str(directory)])
+    return any(not (d.a_path or d.b_path).endswith("pyproject.toml") for d in diff)
+
+
+def main() -> int:
+    push_refs = []
+    for line in sys.stdin:
+        parts = line.strip().split()
+        if len(parts) >= 4:
+            push_refs.append((parts[1], parts[3]))
+
+    if not push_refs:
+        return 0
+
+    repo = git.Repo(search_parent_directories=True)
+    bumped = []
+
+    for local_sha, remote_sha in push_refs:
+        if local_sha == "0" * 40:
+            continue
+
+        local = repo.commit(local_sha)
+        base = get_base(repo, local_sha, remote_sha)
+        if not base:
+            continue
+
+        for package in PACKAGES:
+            pyproject = package / "pyproject.toml"
+
+            if not has_code_changes(repo, package, local, base):
+                continue
+
+            base_version = get_version_at_ref(base, pyproject)
+
+            data = tomlkit.loads(pyproject.read_text())
+            local_version = data["project"]["version"]
+
+            if local_version != base_version:
+                continue
+
+            major, minor, patch = local_version.split(".")
+            new_version = f"{major}.{minor}.{int(patch) + 1}"
+            data["project"]["version"] = new_version
+            pyproject.write_text(tomlkit.dumps(data))
+
+            bumped.append((package, f"{local_version} -> {new_version}"))
+
+    if not bumped:
+        return 0
+
+    files = [str(pkg / "pyproject.toml") for pkg, _ in bumped]
+    repo.index.add(files)
+
+    details = ", ".join(f"{pkg.name} ({ver})" for pkg, ver in bumped)
+    repo.index.commit(f"Bump package version: {details}")
+
+    print(
+        f"Auto-bumped version: {details}. Please push again.",
+        file=sys.stderr,
+    )
+    return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

ymir/agents/backport_agent.py

@@ -44,9 +44,9 @@
     run_tool,
     wrap_details,
 )
-from ymir.common.base_utils import fix_await, redis_client
+from ymir.common.base_utils import fix_await, redis_client, run_task_loop
 from ymir.common.constants import JiraLabels, RedisQueues
-from ymir.common.logging_setup import configure_logging
+from ymir.common.logging_setup import configure_logging, current_jira_issue, get_trajectory_writeable
 from ymir.common.mock_repos import get_mock_local_tool_env
 from ymir.common.models import (
     BackportData,
@@ -186,7 +186,7 @@ async def create_backport_agent(
                 only_success_invocations=False,
             ),
         ],
-        middlewares=[GlobalTrajectoryMiddleware(pretty=True)],
+        middlewares=[GlobalTrajectoryMiddleware(pretty=True, target=get_trajectory_writeable())],
         role="Red Hat Enterprise Linux developer",
         instructions=await get_instructions(fix_version),
     )
@@ -766,7 +766,7 @@ async def comment_in_jira(state):
 async def main() -> None:
     init_sentry()
 
-    configure_logging(level=logging.INFO)
+    configure_logging(level=logging.INFO, buffer_size=int(os.getenv("LOG_BUFFER_SIZE", 0)))
     resolve_chat_model_override("backport")
 
     span_processor = setup_observability(os.environ["COLLECTOR_ENDPOINT"])
@@ -802,6 +802,7 @@ async def main() -> None:
             return
 
     logger.info("Starting backport agent in queue mode")
+    max_concurrent_tasks = int(os.getenv("MAX_CONCURRENT_TASKS", 1))
     async with redis_client(os.environ["REDIS_URL"]) as redis:
         max_retries = int(os.getenv("MAX_RETRIES", 3))
         # Determine which backport queue to listen to based on container version
@@ -818,21 +819,11 @@ async def main() -> None:
             f"listening to queues: [{backport_queue_todo}, {backport_queue}]"
         )
 
-        while True:
-            redis_logger.info(
-                f"Waiting for tasks from [{backport_queue_todo}, {backport_queue}] (timeout: 30s)..."
-            )
-            element = await fix_await(redis.brpop([backport_queue_todo, backport_queue], timeout=30))
-            if element is None:
-                redis_logger.info("No tasks received, continuing to wait...")
-                continue
-
-            _, payload = element
-            redis_logger.info("Received task from queue.")
-
+        async def process_task(payload):
             task = Task.model_validate_json(payload)
             triage_state = task.metadata
             backport_data = BackportData.model_validate(triage_state["triage_result"]["data"])
+            current_jira_issue.set(backport_data.jira_issue)
             dist_git_branch = triage_state["target_branch"]
             user_triggered = task.user_triggered
             logger.info(
@@ -969,6 +960,13 @@ async def retry(
                         ).model_dump_json(),
                     )
 
+        await run_task_loop(
+            redis,
+            [backport_queue_todo, backport_queue],
+            process_task,
+            max_concurrent=max_concurrent_tasks,
+        )
+
 
 if __name__ == "__main__":
     try:

ymir/agents/build_agent.py

@@ -16,6 +16,7 @@
     is_reasoning_enabled,
     render_template,
 )
+from ymir.common.logging_setup import get_trajectory_writeable
 from ymir.tools.unprivileged.commands import RunShellCommandTool
 from ymir.tools.unprivileged.filesystem import GetCWDTool
 from ymir.tools.unprivileged.text import (
@@ -68,7 +69,7 @@ def create_build_agent(mcp_tools: list[Tool], local_tool_options: dict[str, Any]
             ConditionalRequirement("download_artifacts", only_after=["build_package"]),
             ConditionalRequirement("extract_log_snippets", only_after=["download_artifacts"]),
         ],
-        middlewares=[GlobalTrajectoryMiddleware(pretty=True)],
+        middlewares=[GlobalTrajectoryMiddleware(pretty=True, target=get_trajectory_writeable())],
         role="Red Hat Enterprise Linux developer",
         instructions=get_instructions(),
     )

ymir/agents/cve_applicability_agent.py

@@ -12,6 +12,7 @@
 
 from ymir.agents.reasoning_agent import ReasoningAgent
 from ymir.agents.utils import get_chat_model, get_tool_call_checker_config, is_reasoning_enabled
+from ymir.common.logging_setup import get_trajectory_writeable
 from ymir.common.models import Resolution
 from ymir.tools.unprivileged.commands import RunShellCommandTool
 from ymir.tools.unprivileged.text import SearchTextTool, ViewTool
@@ -45,7 +46,7 @@ def create_applicability_agent(
                 only_success_invocations=False,
             ),
         ],
-        middlewares=[GlobalTrajectoryMiddleware(pretty=True)],
+        middlewares=[GlobalTrajectoryMiddleware(pretty=True, target=get_trajectory_writeable())],
         role="Red Hat security analyst",
     )
 

ymir/agents/log_agent.py

@@ -16,6 +16,7 @@
     is_reasoning_enabled,
     render_template,
 )
+from ymir.common.logging_setup import get_trajectory_writeable
 from ymir.tools.unprivileged.commands import RunShellCommandTool
 from ymir.tools.unprivileged.filesystem import GetCWDTool
 from ymir.tools.unprivileged.specfile import AddChangelogEntryTool
@@ -66,7 +67,7 @@ def create_log_agent(_: list[Tool], local_tool_options: dict[str, Any]) -> Reaso
                 only_success_invocations=False,
             ),
         ],
-        middlewares=[GlobalTrajectoryMiddleware(pretty=True)],
+        middlewares=[GlobalTrajectoryMiddleware(pretty=True, target=get_trajectory_writeable())],
         role="Red Hat Enterprise Linux developer",
         instructions=get_instructions(),
     )

ymir/agents/merge_request_agent.py

@@ -34,7 +34,7 @@
     mcp_tools,
     render_template,
 )
-from ymir.common.logging_setup import configure_logging
+from ymir.common.logging_setup import configure_logging, get_trajectory_writeable
 from ymir.common.models import (
     BuildInputSchema,
     BuildOutputSchema,
@@ -96,7 +96,7 @@ def create_merge_request_agent(mcp_tools: list[Tool], local_tool_options: dict[s
                 only_success_invocations=False,
             ),
         ],
-        middlewares=[GlobalTrajectoryMiddleware(pretty=True)],
+        middlewares=[GlobalTrajectoryMiddleware(pretty=True, target=get_trajectory_writeable())],
         role="Red Hat Enterprise Linux developer",
         instructions=get_instructions(),
     )

ymir/agents/observability.py

@@ -1,6 +1,5 @@
 import atexit
 import contextlib
-from contextvars import ContextVar
 
 import sentry_sdk
 from openinference.instrumentation.beeai import BeeAIInstrumentor
@@ -12,21 +11,21 @@
 from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 
+from ymir.common.logging_setup import current_jira_issue
 
-class AgentSpanProcessor(SpanProcessor):
-    _jira_issue_var: ContextVar[str | None] = ContextVar("jira_issue", default=None)
 
+class AgentSpanProcessor(SpanProcessor):
     def set_jira_issue(self, jira_issue: str | None) -> None:
-        self._jira_issue_var.set(jira_issue)
+        current_jira_issue.set(jira_issue)
 
     @contextlib.contextmanager
     def jira_issue_context(self, jira_issue: str | None):
         """Set the jira issue attribute on all spans created within the context."""
-        token = self._jira_issue_var.set(jira_issue)
+        token = current_jira_issue.set(jira_issue)
         try:
             yield
         finally:
-            self._jira_issue_var.reset(token)
+            current_jira_issue.reset(token)
 
     @contextlib.contextmanager
     def start_transaction(
@@ -40,15 +39,15 @@ def start_transaction(
             transaction.set_data("workflow", workflow)
             transaction.set_data("jira_issue", jira_issue)
 
-            token = self._jira_issue_var.set(jira_issue)
+            token = current_jira_issue.set(jira_issue)
             try:
                 yield
             finally:
-                self._jira_issue_var.reset(token)
+                current_jira_issue.reset(token)
 
     def on_start(self, span: Span, parent_context: Context | None = None) -> None:
         if span.is_recording():
-            jira_issue = self._jira_issue_var.get()
+            jira_issue = current_jira_issue.get()
             if jira_issue:
                 span.set_attribute("jira.issue", jira_issue)
 

ymir/agents/preliminary_testing_agent.py

@@ -30,7 +30,7 @@
     render_template,
     run_tool,
 )
-from ymir.common.logging_setup import configure_logging
+from ymir.common.logging_setup import configure_logging, get_trajectory_writeable
 from ymir.tools.unprivileged.greenwave import FetchGreenWaveTool, FetchTestingFarmResultsTool
 
 logger = logging.getLogger(__file__)
@@ -104,7 +104,7 @@ def create_preliminary_testing_agent(gateway_tools: list) -> ReasoningAgent:
                 only_success_invocations=False,
             ),
         ],
-        middlewares=[GlobalTrajectoryMiddleware(pretty=True)],
+        middlewares=[GlobalTrajectoryMiddleware(pretty=True, target=get_trajectory_writeable())],
     )