invoke-ai--invokeai/invokeai/app/services/session_queue/session_queue_base.py

from abc import ABC, abstractmethod
from typing import Any, Coroutine, Optional

from invokeai.app.services.session_queue.session_queue_common import (
    QUEUE_ITEM_STATUS,
    Batch,
    BatchStatus,
    CancelAllExceptCurrentResult,
    CancelByBatchIDsResult,
    CancelByDestinationResult,
    CancelByQueueIDResult,
    ClearResult,
    DeleteAllExceptCurrentResult,
    DeleteByDestinationResult,
    EnqueueBatchResult,
    IsEmptyResult,
    IsFullResult,
    ItemIdsResult,
    NodeFieldValue,
    PruneResult,
    RetryItemsResult,
    SessionQueueCountsByDestination,
    SessionQueueItem,
    SessionQueueStatus,
)
from invokeai.app.services.shared.graph import GraphExecutionState
from invokeai.app.services.shared.pagination import CursorPaginatedResults
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection


class SessionQueueBase(ABC):
    """Base class for session queue"""

    @abstractmethod
    def dequeue(self) -> Optional[SessionQueueItem]:
        """Dequeues the next session queue item."""
        pass

    @abstractmethod
    def enqueue_batch(
        self, queue_id: str, batch: Batch, prepend: bool, user_id: str = "system"
    ) -> Coroutine[Any, Any, EnqueueBatchResult]:
        """Enqueues all permutations of a batch for execution for a specific user."""
        pass

    @abstractmethod
    def get_current(self, queue_id: str) -> Optional[SessionQueueItem]:
        """Gets the currently-executing session queue item"""
        pass

    @abstractmethod
    def get_next(self, queue_id: str) -> Optional[SessionQueueItem]:
        """Gets the next session queue item (does not dequeue it)"""
        pass

    @abstractmethod
    def clear(self, queue_id: str, user_id: Optional[str] = None) -> ClearResult:
        """Deletes all session queue items. If user_id is provided, only clears items owned by that user."""
        pass

    @abstractmethod
    def prune(self, queue_id: str, user_id: Optional[str] = None) -> PruneResult:
        """Deletes all completed and errored session queue items. If user_id is provided, only prunes items owned by that user."""
        pass

    @abstractmethod
    def is_empty(self, queue_id: str) -> IsEmptyResult:
        """Checks if the queue is empty"""
        pass

    @abstractmethod
    def is_full(self, queue_id: str) -> IsFullResult:
        """Checks if the queue is empty"""
        pass

    @abstractmethod
    def get_queue_status(
        self,
        queue_id: str,
        user_id: Optional[str] = None,
        acting_user_id: Optional[str] = None,
    ) -> SessionQueueStatus:
        """Gets the status of the queue.

        Aggregate counts (pending/in_progress/.../total) are always global across all users.
        If user_id is provided, the requesting user's own counts are additionally returned in
        the user_pending/user_in_progress fields (left None otherwise).

        acting_user_id is independent of user_id and controls only current-item redaction:
        when set, the returned status omits item_id/session_id/batch_id unless the
        currently-running item belongs to acting_user_id. The redaction is decided from the
        same get_current() snapshot used to embed those identifiers, so it cannot race against
        a concurrent state change.
        """
        pass

    @abstractmethod
    def get_counts_by_destination(
        self, queue_id: str, destination: str, user_id: Optional[str] = None
    ) -> SessionQueueCountsByDestination:
        """Gets the counts of queue items by destination. If user_id is provided, only counts that user's items."""
        pass

    @abstractmethod
    def get_batch_status(self, queue_id: str, batch_id: str, user_id: Optional[str] = None) -> BatchStatus:
        """Gets the status of a batch. If user_id is provided, only counts that user's items."""
        pass

    @abstractmethod
    def complete_queue_item(self, item_id: int) -> SessionQueueItem:
        """Completes a session queue item"""
        pass

    @abstractmethod
    def suspend_queue_item(self, item_id: int) -> SessionQueueItem:
        """Suspends a session queue item while waiting on a child workflow execution."""
        pass

    @abstractmethod
    def resume_queue_item(self, item_id: int) -> SessionQueueItem:
        """Resumes a suspended session queue item by returning it to pending state."""
        pass

    @abstractmethod
    def cancel_queue_item(self, item_id: int) -> SessionQueueItem:
        """Cancels a session queue item"""
        pass

    @abstractmethod
    def delete_queue_item(self, item_id: int) -> None:
        """Deletes a session queue item"""
        pass

    @abstractmethod
    def delete_queue_items_by_id(self, item_ids: list[int]) -> None:
        """Deletes session queue items by ID."""
        pass

    @abstractmethod
    def fail_queue_item(
        self, item_id: int, error_type: str, error_message: str, error_traceback: str
    ) -> SessionQueueItem:
        """Fails a session queue item"""
        pass

    @abstractmethod
    def cancel_by_batch_ids(
        self, queue_id: str, batch_ids: list[str], user_id: Optional[str] = None
    ) -> CancelByBatchIDsResult:
        """Cancels all queue items with matching batch IDs. If user_id is provided, only cancels items owned by that user."""
        pass

    @abstractmethod
    def cancel_by_destination(
        self, queue_id: str, destination: str, user_id: Optional[str] = None
    ) -> CancelByDestinationResult:
        """Cancels all queue items with the given batch destination. If user_id is provided, only cancels items owned by that user."""
        pass

    @abstractmethod
    def delete_by_destination(
        self, queue_id: str, destination: str, user_id: Optional[str] = None
    ) -> DeleteByDestinationResult:
        """Deletes all queue items with the given batch destination. If user_id is provided, only deletes items owned by that user."""
        pass

    @abstractmethod
    def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult:
        """Cancels all queue items with matching queue ID"""
        pass

    @abstractmethod
    def cancel_all_except_current(self, queue_id: str, user_id: Optional[str] = None) -> CancelAllExceptCurrentResult:
        """Cancels all queue items except in-progress items. If user_id is provided, only cancels items owned by that user."""
        pass

    @abstractmethod
    def delete_all_except_current(self, queue_id: str, user_id: Optional[str] = None) -> DeleteAllExceptCurrentResult:
        """Deletes all queue items except in-progress items. If user_id is provided, only deletes items owned by that user."""
        pass

    @abstractmethod
    def list_queue_items(
        self,
        queue_id: str,
        limit: int,
        priority: int,
        cursor: Optional[int] = None,
        status: Optional[QUEUE_ITEM_STATUS] = None,
        destination: Optional[str] = None,
    ) -> CursorPaginatedResults[SessionQueueItem]:
        """Gets a page of session queue items. Do not remove."""
        pass

    @abstractmethod
    def list_all_queue_items(
        self,
        queue_id: str,
        destination: Optional[str] = None,
    ) -> list[SessionQueueItem]:
        """Gets all queue items that match the given parameters"""
        pass

    @abstractmethod
    def get_queue_item_ids(
        self,
        queue_id: str,
        order_dir: SQLiteDirection = SQLiteDirection.Descending,
        user_id: Optional[str] = None,
    ) -> ItemIdsResult:
        """Gets all queue item ids that match the given parameters. If user_id is provided, only returns items for that user."""
        pass

    @abstractmethod
    def get_queue_item(self, item_id: int) -> SessionQueueItem:
        """Gets a session queue item by ID for a given queue"""
        pass

    @abstractmethod
    def set_queue_item_session(self, item_id: int, session: GraphExecutionState) -> SessionQueueItem:
        """Sets the session for a session queue item. Use this to update the session state."""
        pass

    @abstractmethod
    def enqueue_workflow_call_child(
        self,
        parent_queue_item: SessionQueueItem,
        child_session: GraphExecutionState,
        field_values: list[NodeFieldValue] | None = None,
    ) -> SessionQueueItem:
        """Enqueues a child workflow execution linked to a suspended parent queue item."""
        pass

    @abstractmethod
    def cancel_workflow_call_children(
        self, workflow_call_id: str, exclude_item_ids: set[int] | None = None
    ) -> list[int]:
        """Cancels child workflow queue items for a workflow call without canceling the waiting parent chain."""
        pass

    @abstractmethod
    def retry_items_by_id(self, queue_id: str, item_ids: list[int]) -> RetryItemsResult:
        """Retries the given queue items"""
        pass