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