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
from core.notification_agent import NotificationReaderAgent, NotificationNoToolCalledError
from core.notification_tools import NotificationTools

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)
        self._notification_reader = NotificationReaderAgent(model_name)
        self._max_retries = 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
        inspection_tools = [
            self._tools.get_issue,
            self._tools.get_pull_request,
            self._tools.get_issue_comments,
            self._tools.get_pull_request_comments,
        ]

        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

            # Run NotificationReaderAgent to pre-screen the notification
            mission = (
                f"Analyze Gitea notification ID {notification_id}.\n"
                f"Repository: {repo_full_name}\n"
                f"Subject Type: {subj_type}\n"
                f"Task Number: {task_number}\n"
                f"Please decide if we should process or skip this notification."
            )
            notification_tools = NotificationTools()
            attempt = 0
            attempt_limit = self._max_retries
            success = False
            while attempt < attempt_limit:
                attempt += 1
                try:
                    await self._notification_reader.decide_notification(
                        mission,
                        inspection_tools,
                        notification_tools
                    )
                    success = True
                    break
                except NotificationNoToolCalledError as e:
                    logger.error(f"Notification reader error on notification {notification_id} (attempt {attempt}/{attempt_limit}): {e}")

            if not success or notification_tools.action == "SKIP":
                reason = notification_tools.arguments.get("reason", "Failed to call routing tool / default skip")
                logger.info(f"Skipping notification {notification_id} for {repo_full_name}#{task_number}. Reason: {reason}")
                if notification_id is not None:
                    self._client.mark_notification_as_read(notification_id)
                    logger.info(f"Marked skipped Gitea notification thread {notification_id} as read.")
                continue

            # Route based on decided action
            if notification_tools.action == "PROCESS_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 notification_tools.action == "PROCESS_PR":
                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.")