coding-agent-gitea/core/orchestrator.py

"""Top-level coordinator: polls Gitea unread notifications, queues work, dispatches to agent."""

import asyncio
import logging
import datetime
import json
from pathlib import Path
from typing import Any, Optional

from core.queue import WorkQueue, WorkItem
from core.dispatcher import AgentDispatcher
from gitea.client import GiteaClient
from gitea.models import IssueModel, PullRequestModel, RepositoryModel
from gitea.tools.gitea_tools import GiteaTools
from gitea.config import AGENT_MODEL_ID, AGENT_MAX_RETRIES
from gitea.workspace import WorkspaceManager

logger: logging.Logger = logging.getLogger("agent-orchestrator")


class AgentOrchestrator:
    """Top-level coordinator: polls Gitea notifications, queues work, dispatches to agent."""

    def __init__(
        self,
        client: GiteaClient,
        tools: GiteaTools,
        model_name: str = AGENT_MODEL_ID,
        max_retries: int = AGENT_MAX_RETRIES,
    ) -> None:
        self._client = client
        self._tools = tools
        self._model_name = model_name
        self._work_queue = WorkQueue()
        self._dispatcher = AgentDispatcher(client, tools, model_name, max_retries)

    def _get_state_file_path(self) -> Path:
        """Get the path to the persistent state file."""
        return Path(__file__).parent.parent / "agent_state.json"

    def _read_last_checked(self) -> Optional[str]:
        """Read the last checked timestamp from persistent storage."""
        state_file = self._get_state_file_path()
        if state_file.exists():
            try:
                with open(state_file, "r") as f:
                    data = json.load(f)
                    return data.get("last_checked")
            except Exception as e:
                logger.warning(f"Error reading agent_state.json: {e}")
        return None

    def _write_last_checked(self, timestamp: str) -> None:
        """Write the last checked timestamp to persistent storage."""
        state_file = self._get_state_file_path()
        try:
            with open(state_file, "w") as f:
                json.dump({"last_checked": timestamp}, f)
        except Exception as e:
            logger.error(f"Error writing agent_state.json: {e}")

    async def poll_and_dispatch(self) -> None:
        """Poll Gitea unread notifications, enqueue them, and dispatch to agent."""
        last_checked = self._read_last_checked()
        logger.info(f"Polling unread notifications since: {last_checked or 'beginning'}")

        notifications = self._client.list_unread_notifications(since=last_checked)

        if not notifications:
            logger.info("No new notifications found.")
            return

        logger.info(f"Retrieved {len(notifications)} unread notifications.")

        # Filter and enqueue tasks from notifications
        latest_timestamp = last_checked
        for n in notifications:
            notification_id = n.get("id")
            subject = n.get("subject") or {}
            subj_type = subject.get("type", "").lower()
            subj_url = subject.get("url", "")
            updated_at = n.get("updated_at")

            # Track latest updated_at to advance checkpoint
            if updated_at and (not latest_timestamp or updated_at > latest_timestamp):
                latest_timestamp = updated_at

            repo_info = n.get("repository") or {}
            repo_full_name = repo_info.get("full_name", "")
            if not repo_full_name or not subj_url:
                continue

            owner, repo_name = repo_full_name.split("/")
            try:
                task_number = int(subj_url.rstrip("/").split("/")[-1])
            except (ValueError, IndexError):
                logger.warning(f"Could not parse task number from subject URL: {subj_url}")
                continue

            if subj_type == "issue":
                try:
                    issue = self._client.get_issue(owner, repo_name, task_number)
                    if issue.repository is None:
                        issue = issue.model_copy(update={"repository": RepositoryModel(**repo_info)})

                    item = WorkItem(
                        repo_full_name=repo_full_name,
                        task_type="issue",
                        task_number=task_number,
                        task_info=issue,
                        notification_id=notification_id,
                        priority=0
                    )
                    self._work_queue.enqueue(item)
                except Exception as e:
                    logger.error(f"Failed to fetch issue #{task_number} for notification: {e}")
            elif subj_type in ("pull", "pullrequest"):
                try:
                    pr = self._client.get_pull_request(owner, repo_name, task_number)
                    if pr.repository is None:
                        pr = pr.model_copy(update={"repository": RepositoryModel(**repo_info)})

                    item = WorkItem(
                        repo_full_name=repo_full_name,
                        task_type="pr",
                        task_number=task_number,
                        task_info=pr,
                        notification_id=notification_id,
                        priority=0
                    )
                    self._work_queue.enqueue(item)
                except Exception as e:
                    logger.error(f"Failed to fetch PR #{task_number} for notification: {e}")

        # Process enqueued work
        if not self._work_queue.is_empty:
            await self._process_work()

        # Update last checked timestamp checkpoint
        if latest_timestamp:
            self._write_last_checked(latest_timestamp)

    async def _process_work(self) -> None:
        """Process all queued work, repo by repo, then mark notifications as read."""
        while not self._work_queue.is_empty:
            repo = self._work_queue.get_next_repo()
            if not repo:
                break

            work_items = self._work_queue.get_repo_work(repo)
            self._work_queue.remove_repo_work(repo)

            # Ensure the workspace repository is cloned and sanitized
            workspace = WorkspaceManager()
            repo_path = workspace.get_repo_path(repo)
            if not repo_path.exists():
                workspace.clone_repo(repo)
                logger.info(f"Cloned {repo} to {repo_path}")
            else:
                workspace.sanitize_repo(repo, repo_path)
                logger.info(f"Sanitized existing repo at {repo_path}")

            logger.info(f"Dispatching {len(work_items)} tasks for {repo}")
            results = await self._dispatcher.dispatch(repo, work_items)

            for i, result in enumerate(results):
                item = work_items[i]
                logger.info(f"Completed {item.task_type} #{item.task_number}: {result[:200]}")
                if item.notification_id is not None:
                    self._client.mark_notification_as_read(item.notification_id)
                    logger.info(f"Marked Gitea notification thread {item.notification_id} as read.")