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

import asyncio
import json
import sqlite3
from typing import Any, Optional, Union, cast

from pydantic_core import to_jsonable_python

from invokeai.app.services.invoker import Invoker
from invokeai.app.services.session_queue.session_queue_base import SessionQueueBase
from invokeai.app.services.session_queue.session_queue_common import (
    DEFAULT_QUEUE_ID,
    QUEUE_ITEM_STATUS,
    Batch,
    BatchStatus,
    CancelAllExceptCurrentResult,
    CancelByBatchIDsResult,
    CancelByDestinationResult,
    CancelByQueueIDResult,
    ClearResult,
    DeleteAllExceptCurrentResult,
    DeleteByDestinationResult,
    EnqueueBatchResult,
    IsEmptyResult,
    IsFullResult,
    ItemIdsResult,
    NodeFieldValue,
    PruneResult,
    RetryItemsResult,
    SessionQueueCountsByDestination,
    SessionQueueItem,
    SessionQueueItemNotFoundError,
    SessionQueueStatus,
    TooManySessionsError,
    ValueToInsertTuple,
    calc_session_count,
    prepare_values_to_insert,
)
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
from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase

# Round-robin dequeue (multiuser fairness): pick the next pending item from the user who was
# least-recently served.
#
# The "whose turn is it" ordering key is each candidate user's most recent started_at. We compute
# it with a *correlated* MAX subquery rather than a GROUP BY over all started rows: the candidate
# set is one row per user with pending work, and each MAX(started_at) WHERE user_id = ? is
# satisfied by an indexed seek (idx_session_queue_user_started_at) instead of scanning the full
# retained queue history. This keeps dequeue cost proportional to the number of active users, not
# to total history (which is unbounded by default). MAX() ignores NULL started_at values, so users
# with only pending items fall back to the epoch via COALESCE and are served first.
#
# Kept as a module constant so the scaling test can EXPLAIN QUERY PLAN the exact production SQL.
ROUND_ROBIN_DEQUEUE_QUERY = """--sql
    WITH user_next_item AS (
        -- For each user, select their single best pending item (highest priority, then oldest).
        SELECT
            user_id,
            item_id,
            ROW_NUMBER() OVER (
                PARTITION BY user_id
                ORDER BY priority DESC, item_id ASC
            ) AS rn
        FROM session_queue
        WHERE status = 'pending'
    )
    SELECT
        sq.*,
        u.display_name AS user_display_name,
        u.email AS user_email
    FROM session_queue sq
    LEFT JOIN users u ON sq.user_id = u.user_id
    JOIN user_next_item uni ON sq.item_id = uni.item_id AND uni.rn = 1
    ORDER BY
        COALESCE(
            (
                SELECT MAX(served.started_at)
                FROM session_queue served
                WHERE served.user_id = sq.user_id
            ),
            '1970-01-01'
        ) ASC,
        sq.item_id ASC
    LIMIT 1
    """

# FIFO dequeue (single-user mode, or round_robin explicitly disabled): strict priority then
# insertion order.
FIFO_DEQUEUE_QUERY = """--sql
    SELECT
        sq.*,
        u.display_name as user_display_name,
        u.email as user_email
    FROM session_queue sq
    LEFT JOIN users u ON sq.user_id = u.user_id
    WHERE sq.status = 'pending'
    ORDER BY
        sq.priority DESC,
        sq.item_id ASC
    LIMIT 1
    """


class SqliteSessionQueue(SessionQueueBase):
    __invoker: Invoker

    def start(self, invoker: Invoker) -> None:
        self.__invoker = invoker
        self._set_in_progress_to_canceled()
        config = self.__invoker.services.configuration
        if config.clear_queue_on_startup:
            clear_result = self.clear(DEFAULT_QUEUE_ID)
            if clear_result.deleted > 0:
                self.__invoker.services.logger.info(f"Cleared all {clear_result.deleted} queue items")
            return

        if config.max_queue_history is not None:
            deleted = self._prune_terminal_to_limit(DEFAULT_QUEUE_ID, config.max_queue_history)
            if deleted > 0:
                self.__invoker.services.logger.info(
                    f"Pruned {deleted} completed/failed/canceled queue items (kept up to {config.max_queue_history})"
                )

    def __init__(self, db: SqliteDatabase) -> None:
        super().__init__()
        self._db = db

    def _set_in_progress_to_canceled(self) -> None:
        """
        Sets all in_progress or waiting queue items to canceled. Run on app startup, not associated with any queue.
        This is necessary because the invoker may have been killed while processing a queue item or while a parent
        queue item was suspended waiting on a child workflow execution.
        """
        with self._db.transaction() as cursor:
            cursor.execute(
                """--sql
                SELECT item_id
                FROM session_queue
                WHERE status = 'in_progress'
                   OR status = 'waiting';
                """
            )
            interrupted_item_ids = [row[0] for row in cast(list[sqlite3.Row], cursor.fetchall())]
        item_ids_to_cancel: set[int] = set()
        for item_id in interrupted_item_ids:
            item_ids_to_cancel.update(self._get_workflow_call_chain_item_ids(item_id))
        if not item_ids_to_cancel:
            return
        with self._db.transaction() as cursor:
            placeholders = ",".join("?" for _ in item_ids_to_cancel)
            cursor.execute(
                f"""--sql
                UPDATE session_queue
                SET status = 'canceled',
                    status_sequence = COALESCE(status_sequence, 0) + 1
                WHERE item_id IN ({placeholders})
                  AND (
                    status = 'pending'
                    OR status = 'in_progress'
                    OR status = 'waiting'
                  );
                """,
                tuple(item_ids_to_cancel),
            )

    def _prune_terminal_to_limit(self, queue_id: str, keep: int) -> int:
        """Prune terminal items (completed/failed/canceled) to keep at most N most-recent items."""
        with self._db.transaction() as cursor:
            where = """--sql
                WHERE
                queue_id = ?
                AND (
                    status = 'completed'
                    OR status = 'failed'
                    OR status = 'canceled'
                )
                """
            cursor.execute(
                f"""--sql
                SELECT COUNT(*)
                FROM session_queue
                {where}
                AND item_id NOT IN (
                    SELECT item_id
                    FROM session_queue
                    {where}
                    ORDER BY COALESCE(completed_at, updated_at, created_at) DESC, item_id DESC
                    LIMIT ?
                );
                """,
                (queue_id, queue_id, keep),
            )
            count = cursor.fetchone()[0]
            cursor.execute(
                f"""--sql
                DELETE
                FROM session_queue
                {where}
                AND item_id NOT IN (
                    SELECT item_id
                    FROM session_queue
                    {where}
                    ORDER BY COALESCE(completed_at, updated_at, created_at) DESC, item_id DESC
                    LIMIT ?
                );
                """,
                (queue_id, queue_id, keep),
            )
        return count

    def _get_current_queue_size(self, queue_id: str) -> int:
        """Gets the current number of pending queue items"""
        with self._db.transaction() as cursor:
            cursor.execute(
                """--sql
                SELECT count(*)
                FROM session_queue
                WHERE
                queue_id = ?
                AND status = 'pending'
                """,
                (queue_id,),
            )
            count = cast(int, cursor.fetchone()[0])
        return count

    def _get_highest_priority(self, queue_id: str) -> int:
        """Gets the highest priority value in the queue"""
        with self._db.transaction() as cursor:
            cursor.execute(
                """--sql
                SELECT MAX(priority)
                FROM session_queue
                WHERE
                queue_id = ?
                AND status = 'pending'
                """,
                (queue_id,),
            )
            priority = cast(Union[int, None], cursor.fetchone()[0]) or 0
        return priority

    async def enqueue_batch(
        self, queue_id: str, batch: Batch, prepend: bool, user_id: str = "system"
    ) -> EnqueueBatchResult:
        current_queue_size = self._get_current_queue_size(queue_id)
        max_queue_size = self.__invoker.services.configuration.max_queue_size
        max_new_queue_items = max_queue_size - current_queue_size

        priority = 0
        if prepend:
            priority = self._get_highest_priority(queue_id) + 1

        requested_count = await asyncio.to_thread(
            calc_session_count,
            batch=batch,
        )
        values_to_insert = await asyncio.to_thread(
            prepare_values_to_insert,
            queue_id=queue_id,
            batch=batch,
            priority=priority,
            max_new_queue_items=max_new_queue_items,
            user_id=user_id,
        )
        enqueued_count = len(values_to_insert)

        with self._db.transaction() as cursor:
            cursor.executemany(
                """--sql
                    INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination, retried_from_item_id, user_id)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                values_to_insert,
            )
            cursor.execute(
                """--sql
                    SELECT item_id
                    FROM session_queue
                    WHERE batch_id = ?
                    ORDER BY item_id DESC;
                    """,
                (batch.batch_id,),
            )
            item_ids = [row[0] for row in cursor.fetchall()]
        enqueue_result = EnqueueBatchResult(
            queue_id=queue_id,
            requested=requested_count,
            enqueued=enqueued_count,
            batch=batch,
            priority=priority,
            item_ids=item_ids,
        )
        self.__invoker.services.events.emit_batch_enqueued(enqueue_result, user_id=user_id)
        return enqueue_result

    def dequeue(self) -> Optional[SessionQueueItem]:
        config = self.__invoker.services.configuration
        use_round_robin = config.multiuser and config.session_queue_mode == "round_robin"

        query = ROUND_ROBIN_DEQUEUE_QUERY if use_round_robin else FIFO_DEQUEUE_QUERY

        with self._db.transaction() as cursor:
            cursor.execute(query)
            result = cast(Union[sqlite3.Row, None], cursor.fetchone())
        if result is None:
            return None
        queue_item = SessionQueueItem.queue_item_from_dict(dict(result))
        queue_item = self._set_queue_item_status(item_id=queue_item.item_id, status="in_progress")
        return queue_item

    def get_next(self, queue_id: str) -> Optional[SessionQueueItem]:
        with self._db.transaction() as cursor:
            cursor.execute(
                """--sql
                SELECT
                    sq.*,
                    u.display_name as user_display_name,
                    u.email as user_email
                FROM session_queue sq
                LEFT JOIN users u ON sq.user_id = u.user_id
                WHERE
                    sq.queue_id = ?
                    AND sq.status = 'pending'
                ORDER BY
                    sq.priority DESC,
                    sq.created_at ASC
                LIMIT 1
                """,
                (queue_id,),
            )
            result = cast(Union[sqlite3.Row, None], cursor.fetchone())
        if result is None:
            return None
        return SessionQueueItem.queue_item_from_dict(dict(result))

    def get_current(self, queue_id: str) -> Optional[SessionQueueItem]:
        with self._db.transaction() as cursor:
            cursor.execute(
                """--sql
                SELECT
                    sq.*,
                    u.display_name as user_display_name,
                    u.email as user_email
                FROM session_queue sq
                LEFT JOIN users u ON sq.user_id = u.user_id
                WHERE
                    sq.queue_id = ?
                    AND sq.status = 'in_progress'
                LIMIT 1
                """,
                (queue_id,),
            )
            result = cast(Union[sqlite3.Row, None], cursor.fetchone())
        if result is None:
            return None
        return SessionQueueItem.queue_item_from_dict(dict(result))

    def _set_queue_item_status(
        self,
        item_id: int,
        status: QUEUE_ITEM_STATUS,
        error_type: Optional[str] = None,
        error_message: Optional[str] = None,
        error_traceback: Optional[str] = None,
    ) -> SessionQueueItem:
        with self._db.transaction() as cursor:
            cursor.execute(
                """--sql
                SELECT status FROM session_queue WHERE item_id = ?
                """,
                (item_id,),
            )
            row = cursor.fetchone()
        if row is None:
            raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}")
        current_status = row[0]

        # Only update if not already finished (completed, failed or canceled)
        if current_status in ("completed", "failed", "canceled"):
            return self.get_queue_item(item_id)

        with self._db.transaction() as cursor:
            cursor.execute(
                """--sql
                UPDATE session_queue
                SET status = ?, status_sequence = COALESCE(status_sequence, 0) + 1, error_type = ?, error_message = ?, error_traceback = ?
                WHERE item_id = ?
                """,
                (status, error_type, error_message, error_traceback, item_id),
            )

        queue_item = self.get_queue_item(item_id)
        batch_status = self.get_batch_status(queue_id=queue_item.queue_id, batch_id=queue_item.batch_id)
        # The QueueItemStatusChangedEvent ships to user:{queue_item.user_id} and admin rooms.
        # acting_user_id ensures the embedded current-item identifiers are redacted when the
        # in-progress item belongs to someone else, while leaving aggregate counts global.
        # Doing this inside get_queue_status guarantees the redaction decision and the
        # embedded identifiers come from the same get_current() snapshot — eliminating the
        # race where a second read could find None and skip scrubbing stale identifiers.
        queue_status = self.get_queue_status(queue_id=queue_item.queue_id, acting_user_id=queue_item.user_id)

        self.__invoker.services.events.emit_queue_item_status_changed(queue_item, batch_status, queue_status)
        return queue_item

    def _get_workflow_call_child_ids(self, item_id: int) -> list[int]:
        with self._db.transaction() as cursor:
            cursor.execute(
                """--sql
                SELECT item_id
                FROM session_queue
                WHERE parent_item_id = ?
                ORDER BY item_id ASC
                """,
                (item_id,),
            )
            rows = cast(list[sqlite3.Row], cursor.fetchall())
        return [row[0] for row in rows]

    def _get_workflow_call_descendant_ids(self, item_id: int) -> list[int]:
        descendant_ids: list[int] = []
        queue: list[int] = [item_id]
        while queue:
            current_item_id = queue.pop(0)
            child_ids = self._get_workflow_call_child_ids(current_item_id)
            descendant_ids.extend(child_ids)
            queue.extend(child_ids)
        return descendant_ids

    def _get_workflow_call_ancestor_ids(self, item_id: int) -> list[int]:
        ancestor_ids: list[int] = []
        current_queue_item = self.get_queue_item(item_id)
        while current_queue_item.parent_item_id is not None:
            parent_item_id = current_queue_item.parent_item_id
            ancestor_ids.append(parent_item_id)
            current_queue_item = self.get_queue_item(parent_item_id)
        return ancestor_ids

    def _get_workflow_call_chain_item_ids(self, item_id: int) -> list[int]:
        ancestor_ids = self._get_workflow_call_ancestor_ids(item_id)
        root_item_id = ancestor_ids[-1] if ancestor_ids else item_id
        descendant_ids = self._get_workflow_call_descendant_ids(root_item_id)
        chain_item_ids = ancestor_ids + [item_id] + descendant_ids
        deduped_chain_item_ids = list(dict.fromkeys(chain_item_ids))
        return deduped_chain_item_ids

    def _get_current_workflow_call_chain_item_ids(self, queue_id: str) -> set[int]:
        current_queue_item = self.get_current(queue_id)
        if current_queue_item is not None:
            return set(self._get_workflow_call_chain_item_ids(current_queue_item.item_id))

        with self._db.transaction() as cursor:
            cursor.execute(
                """--sql
                SELECT child.item_id
                FROM session_queue child
                JOIN session_queue parent ON parent.item_id = child.parent_item_id
                WHERE
                    child.queue_id = ?
                    AND child.status = 'pending'
                    AND parent.status = 'waiting'
                ORDER BY child.priority DESC, child.created_at ASC, child.item_id ASC
                LIMIT 1
                """,
                (queue_id,),
            )
            row = cursor.fetchone()
        if row is None:
            return set()
        return set(self._get_workflow_call_chain_item_ids(cast(int, row[0])))

    def is_empty(self, queue_id: str) -> IsEmptyResult:
        with self._db.transaction() as cursor:
            cursor.execute(
                """--sql
                SELECT count(*)
                FROM session_queue
                WHERE queue_id = ?
                """,
                (queue_id,),
            )
            is_empty = cast(int, cursor.fetchone()[0]) == 0
        return IsEmptyResult(is_empty=is_empty)

    def is_full(self, queue_id: str) -> IsFullResult:
        with self._db.transaction() as cursor:
            cursor.execute(
                """--sql
                SELECT count(*)
                FROM session_queue
                WHERE queue_id = ?
                """,
                (queue_id,),
            )
            max_queue_size = self.__invoker.services.configuration.max_queue_size
            is_full = cast(int, cursor.fetchone()[0]) >= max_queue_size
        return IsFullResult(is_full=is_full)

    def clear(self, queue_id: str, user_id: Optional[str] = None) -> ClearResult:
        with self._db.transaction() as cursor:
            user_filter = "AND user_id = ?" if user_id is not None else ""
            where = f"""--sql
                WHERE queue_id = ?
                {user_filter}
                """
            params: list[str] = [queue_id]
            if user_id is not None:
                params.append(user_id)
            cursor.execute(
                f"""--sql
                SELECT COUNT(*)
                FROM session_queue
                {where}
                """,
                tuple(params),
            )
            count = cursor.fetchone()[0]
            cursor.execute(
                f"""--sql
                DELETE
                FROM session_queue
                {where}
                """,
                tuple(params),
            )
        self.__invoker.services.events.emit_queue_cleared(queue_id, user_id)
        return ClearResult(deleted=count)

    def delete_queue_items_by_id(self, item_ids: list[int]) -> None:
        if not item_ids:
            return
        placeholders = ", ".join(["?" for _ in item_ids])
        with self._db.transaction() as cursor:
            cursor.execute(
                f"""--sql
                DELETE
                FROM session_queue
                WHERE item_id IN ({placeholders})
                """,
                tuple(item_ids),
            )

    def prune(self, queue_id: str, user_id: Optional[str] = None) -> PruneResult:
        with self._db.transaction() as cursor:
            # Build WHERE clause with optional user_id filter
            user_filter = "AND user_id = ?" if user_id is not None else ""
            where = f"""--sql
                WHERE
                queue_id = ?
                AND (
                    status = 'completed'
                    OR status = 'failed'
                    OR status = 'canceled'
                )
                {user_filter}
                """
            params: list[Any] = [queue_id]
            if user_id is not None:
                params.append(user_id)

            cursor.execute(
                f"""--sql
                SELECT COUNT(*)
                FROM session_queue
                {where};
                """,
                tuple(params),
            )
            count = cursor.fetchone()[0]
            cursor.execute(
                f"""--sql
                DELETE
                FROM session_queue
                {where};
                """,
                tuple(params),
            )
        return PruneResult(deleted=count)

    def cancel_queue_item(self, item_id: int) -> SessionQueueItem:
        chain_item_ids = self._get_workflow_call_chain_item_ids(item_id)
        for chain_item_id in chain_item_ids:
            self._set_queue_item_status(item_id=chain_item_id, status="canceled")
        return self.get_queue_item(item_id)

    def delete_queue_item(self, item_id: int) -> None:
        """Deletes a session queue item"""
        chain_item_ids = self._get_workflow_call_chain_item_ids(item_id)
        if any(
            self.get_queue_item(chain_item_id).status not in {"completed", "failed", "canceled"}
            for chain_item_id in chain_item_ids
        ):
            self.cancel_queue_item(item_id)
        self.delete_queue_items_by_id(chain_item_ids)

    def complete_queue_item(self, item_id: int) -> SessionQueueItem:
        queue_item = self._set_queue_item_status(item_id=item_id, status="completed")
        return queue_item

    def suspend_queue_item(self, item_id: int) -> SessionQueueItem:
        queue_item = self._set_queue_item_status(item_id=item_id, status="waiting")
        return queue_item

    def resume_queue_item(self, item_id: int) -> SessionQueueItem:
        queue_item = self._set_queue_item_status(item_id=item_id, status="pending")
        return queue_item

    def fail_queue_item(
        self,
        item_id: int,
        error_type: str,
        error_message: str,
        error_traceback: str,
    ) -> SessionQueueItem:
        queue_item = self._set_queue_item_status(
            item_id=item_id,
            status="failed",
            error_type=error_type,
            error_message=error_message,
            error_traceback=error_traceback,
        )
        return queue_item

    def cancel_by_batch_ids(
        self, queue_id: str, batch_ids: list[str], user_id: Optional[str] = None
    ) -> CancelByBatchIDsResult:
        with self._db.transaction() as cursor:
            current_queue_item = self.get_current(queue_id)
            placeholders = ", ".join(["?" for _ in batch_ids])

            # Build WHERE clause with optional user_id filter
            user_filter = "AND user_id = ?" if user_id is not None else ""
            where = f"""--sql
                WHERE
                  queue_id == ?
                  AND batch_id IN ({placeholders})
                  AND status != 'canceled'
                  AND status != 'completed'
                  AND status != 'failed'
                  -- We will cancel the current item separately below - skip it here
                  AND status != 'in_progress'
                  {user_filter}
                """
            params = [queue_id] + batch_ids
            if user_id is not None:
                params.append(user_id)

            cursor.execute(
                f"""--sql
                SELECT COUNT(*)
                FROM session_queue
                {where};
                """,
                tuple(params),
            )
            count = cursor.fetchone()[0]
            cursor.execute(
                f"""--sql
                UPDATE session_queue
                SET status = 'canceled',
                    status_sequence = COALESCE(status_sequence, 0) + 1
                {where};
                """,
                tuple(params),
            )

        # Handle current item separately - check ownership if user_id is provided
        if current_queue_item is not None and current_queue_item.batch_id in batch_ids:
            if user_id is None or current_queue_item.user_id == user_id:
                self._set_queue_item_status(current_queue_item.item_id, "canceled")

        return CancelByBatchIDsResult(canceled=count)

    def cancel_by_destination(
        self, queue_id: str, destination: str, user_id: Optional[str] = None
    ) -> CancelByDestinationResult:
        with self._db.transaction() as cursor:
            current_queue_item = self.get_current(queue_id)

            # Build WHERE clause with optional user_id filter
            user_filter = "AND user_id = ?" if user_id is not None else ""
            where = f"""--sql
                WHERE
                  queue_id == ?
                  AND destination == ?
                  AND status != 'canceled'
                  AND status != 'completed'
                  AND status != 'failed'
                  -- We will cancel the current item separately below - skip it here
                  AND status != 'in_progress'
                  {user_filter}
                """
            params = [queue_id, destination]
            if user_id is not None:
                params.append(user_id)

            cursor.execute(
                f"""--sql
                SELECT COUNT(*)
                FROM session_queue
                {where};
                """,
                tuple(params),
            )
            count = cursor.fetchone()[0]
            cursor.execute(
                f"""--sql
                UPDATE session_queue
                SET status = 'canceled',
                    status_sequence = COALESCE(status_sequence, 0) + 1
                {where};
                """,
                tuple(params),
            )

        # Handle current item separately - check ownership if user_id is provided
        if current_queue_item is not None and current_queue_item.destination == destination:
            if user_id is None or current_queue_item.user_id == user_id:
                self._set_queue_item_status(current_queue_item.item_id, "canceled")

        return CancelByDestinationResult(canceled=count)

    def delete_by_destination(
        self, queue_id: str, destination: str, user_id: Optional[str] = None
    ) -> DeleteByDestinationResult:
        with self._db.transaction() as cursor:
            current_queue_item = self.get_current(queue_id)

            # Handle current item separately - check ownership if user_id is provided
            if current_queue_item is not None and current_queue_item.destination == destination:
                if user_id is None or current_queue_item.user_id == user_id:
                    self.cancel_queue_item(current_queue_item.item_id)

            # Build WHERE clause with optional user_id filter
            user_filter = "AND user_id = ?" if user_id is not None else ""
            params = [queue_id, destination]
            if user_id is not None:
                params.append(user_id)

            cursor.execute(
                f"""--sql
                SELECT COUNT(*)
                FROM session_queue
                WHERE
                  queue_id == ?
                  AND destination == ?
                  {user_filter}
                """,
                tuple(params),
            )
            count = cursor.fetchone()[0]
            cursor.execute(
                f"""--sql
                DELETE FROM session_queue
                WHERE
                  queue_id == ?
                  AND destination == ?
                  {user_filter}
                """,
                tuple(params),
            )
        return DeleteByDestinationResult(deleted=count)

    def delete_all_except_current(self, queue_id: str, user_id: Optional[str] = None) -> DeleteAllExceptCurrentResult:
        current_chain_item_ids = self._get_current_workflow_call_chain_item_ids(queue_id)
        with self._db.transaction() as cursor:
            # Build WHERE clause with optional user_id filter
            user_filter = "AND user_id = ?" if user_id is not None else ""
            current_chain_filter = ""
            if current_chain_item_ids:
                placeholders = ", ".join(["?" for _ in current_chain_item_ids])
                current_chain_filter = f"AND item_id NOT IN ({placeholders})"
            where = f"""--sql
                WHERE
                  queue_id == ?
                  AND status IN ('pending', 'waiting')
                  {user_filter}
                  {current_chain_filter}
                """
            params: list[Any] = [queue_id]
            if user_id is not None:
                params.append(user_id)
            params.extend(current_chain_item_ids)

            cursor.execute(
                f"""--sql
                SELECT COUNT(*)
                FROM session_queue
                {where};
                """,
                tuple(params),
            )
            count = cursor.fetchone()[0]
            cursor.execute(
                f"""--sql
                DELETE
                FROM session_queue
                {where};
                """,
                tuple(params),
            )
        return DeleteAllExceptCurrentResult(deleted=count)

    def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult:
        with self._db.transaction() as cursor:
            current_queue_item = self.get_current(queue_id)
            where = """--sql
                WHERE
                  queue_id is ?
                  AND status != 'canceled'
                  AND status != 'completed'
                  AND status != 'failed'
                  -- We will cancel the current item separately below - skip it here
                  AND status != 'in_progress'
                """
            params = [queue_id]
            cursor.execute(
                f"""--sql
                SELECT COUNT(*)
                FROM session_queue
                {where};
                """,
                tuple(params),
            )
            count = cursor.fetchone()[0]
            cursor.execute(
                f"""--sql
                UPDATE session_queue
                SET status = 'canceled',
                    status_sequence = COALESCE(status_sequence, 0) + 1
                {where};
                """,
                tuple(params),
            )

        if current_queue_item is not None and current_queue_item.queue_id == queue_id:
            self._set_queue_item_status(current_queue_item.item_id, "canceled")
        return CancelByQueueIDResult(canceled=count)

    def cancel_all_except_current(self, queue_id: str, user_id: Optional[str] = None) -> CancelAllExceptCurrentResult:
        current_chain_item_ids = self._get_current_workflow_call_chain_item_ids(queue_id)
        with self._db.transaction() as cursor:
            # Build WHERE clause with optional user_id filter
            user_filter = "AND user_id = ?" if user_id is not None else ""
            current_chain_filter = ""
            if current_chain_item_ids:
                placeholders = ", ".join(["?" for _ in current_chain_item_ids])
                current_chain_filter = f"AND item_id NOT IN ({placeholders})"
            where = f"""--sql
                WHERE
                  queue_id == ?
                  AND status IN ('pending', 'waiting')
                  {user_filter}
                  {current_chain_filter}
                """
            params = [queue_id]
            if user_id is not None:
                params.append(user_id)
            params.extend(current_chain_item_ids)

            cursor.execute(
                f"""--sql
                SELECT COUNT(*)
                FROM session_queue
                {where};
                """,
                tuple(params),
            )
            count = cursor.fetchone()[0]
            cursor.execute(
                f"""--sql
                UPDATE session_queue
                SET status = 'canceled',
                    status_sequence = COALESCE(status_sequence, 0) + 1
                {where};
                """,
                tuple(params),
            )
        return CancelAllExceptCurrentResult(canceled=count)

    def get_queue_item(self, item_id: int) -> SessionQueueItem:
        with self._db.transaction() as cursor:
            cursor.execute(
                """--sql
                SELECT
                    sq.*,
                    u.display_name as user_display_name,
                    u.email as user_email
                FROM session_queue sq
                LEFT JOIN users u ON sq.user_id = u.user_id
                WHERE sq.item_id = ?
                """,
                (item_id,),
            )
            result = cast(Union[sqlite3.Row, None], cursor.fetchone())
        if result is None:
            raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}")
        return SessionQueueItem.queue_item_from_dict(dict(result))

    def set_queue_item_session(self, item_id: int, session: GraphExecutionState) -> SessionQueueItem:
        with self._db.transaction() as cursor:
            # Use exclude_none so we don't end up with a bunch of nulls in the graph - this can cause validation errors
            # when the graph is loaded. Graph execution occurs purely in memory - the session saved here is not referenced
            # during execution.
            session_json = session.model_dump_json(warnings=False, exclude_none=True)
            cursor.execute(
                """--sql
                UPDATE session_queue
                SET session = ?
                WHERE item_id = ?
                """,
                (session_json, item_id),
            )
        return self.get_queue_item(item_id)

    def enqueue_workflow_call_child(
        self,
        parent_queue_item: SessionQueueItem,
        child_session: GraphExecutionState,
        field_values: list[NodeFieldValue] | None = None,
    ) -> SessionQueueItem:
        workflow_call_execution = parent_queue_item.session.waiting_workflow_call_execution
        if workflow_call_execution is None:
            raise ValueError("Parent queue item is missing active workflow call execution metadata.")

        session_json = child_session.model_dump_json(warnings=False, exclude_none=True)
        field_values_json = json.dumps(field_values, default=to_jsonable_python) if field_values is not None else None
        root_item_id = parent_queue_item.root_item_id or parent_queue_item.item_id

        with self._db.transaction() as cursor:
            cursor.execute(
                """--sql
                SELECT COUNT(*)
                FROM session_queue
                WHERE queue_id = ? AND status = 'pending'
                """,
                (parent_queue_item.queue_id,),
            )
            pending_count = cast(int, cursor.fetchone()[0])
            if pending_count >= self.__invoker.services.configuration.max_queue_size:
                raise TooManySessionsError(
                    "call_saved_workflow exceeds remaining queue capacity for child workflow executions"
                )

            cursor.execute(
                """--sql
                INSERT INTO session_queue (
                    queue_id,
                    session,
                    session_id,
                    batch_id,
                    field_values,
                    priority,
                    workflow,
                    origin,
                    destination,
                    retried_from_item_id,
                    user_id,
                    workflow_call_id,
                    parent_item_id,
                    parent_session_id,
                    root_item_id,
                    workflow_call_depth,
                    status
                )
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'pending')
                """,
                (
                    parent_queue_item.queue_id,
                    session_json,
                    child_session.id,
                    parent_queue_item.batch_id,
                    field_values_json,
                    parent_queue_item.priority,
                    None,
                    parent_queue_item.origin,
                    parent_queue_item.destination,
                    None,
                    parent_queue_item.user_id,
                    workflow_call_execution.id,
                    parent_queue_item.item_id,
                    parent_queue_item.session_id,
                    root_item_id,
                    workflow_call_execution.depth,
                ),
            )
            item_id = cursor.lastrowid

        queue_item = self.get_queue_item(item_id)
        batch_status = self.get_batch_status(queue_id=queue_item.queue_id, batch_id=queue_item.batch_id)
        queue_status = self.get_queue_status(queue_id=queue_item.queue_id, acting_user_id=queue_item.user_id)
        self.__invoker.services.events.emit_queue_item_status_changed(queue_item, batch_status, queue_status)
        return queue_item

    def cancel_workflow_call_children(
        self, workflow_call_id: str, exclude_item_ids: set[int] | None = None
    ) -> list[int]:
        exclude_item_ids = exclude_item_ids or set()
        with self._db.transaction() as cursor:
            cursor.execute(
                """--sql
                SELECT item_id
                FROM session_queue
                WHERE workflow_call_id = ?
                ORDER BY item_id ASC
                """,
                (workflow_call_id,),
            )
            item_ids = [row[0] for row in cast(list[sqlite3.Row], cursor.fetchall())]
        item_ids_with_descendants: list[int] = []
        for item_id in item_ids:
            item_ids_with_descendants.append(item_id)
            item_ids_with_descendants.extend(self._get_workflow_call_descendant_ids(item_id))
        item_ids = list(dict.fromkeys(item_ids_with_descendants))
        canceled_item_ids: list[int] = []
        for item_id in item_ids:
            if item_id in exclude_item_ids:
                continue
            queue_item = self.get_queue_item(item_id)
            if queue_item.status in {"completed", "failed", "canceled"}:
                continue
            self._set_queue_item_status(item_id=item_id, status="canceled")
            canceled_item_ids.append(item_id)
        return canceled_item_ids

    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]:
        with self._db.transaction() as cursor_:
            item_id = cursor
            query = """--sql
                SELECT *
                FROM session_queue
                WHERE queue_id = ?
            """
            params: list[Union[str, int]] = [queue_id]

            if status is not None:
                query += """--sql
                    AND status = ?
                    """
                params.append(status)

            if destination is not None:
                query += """---sql
                    AND destination = ?
                """
                params.append(destination)

            if item_id is not None:
                query += """--sql
                    AND (priority < ?) OR (priority = ? AND item_id > ?)
                    """
                params.extend([priority, priority, item_id])

            query += """--sql
                ORDER BY
                    priority DESC,
                    item_id ASC
                LIMIT ?
                """
            params.append(limit + 1)
            cursor_.execute(query, params)
            results = cast(list[sqlite3.Row], cursor_.fetchall())
        items = [SessionQueueItem.queue_item_from_dict(dict(result)) for result in results]
        has_more = False
        if len(items) > limit:
            # remove the extra item
            items.pop()
            has_more = True
        return CursorPaginatedResults(items=items, limit=limit, has_more=has_more)

    def list_all_queue_items(
        self,
        queue_id: str,
        destination: Optional[str] = None,
    ) -> list[SessionQueueItem]:
        """Gets all queue items that match the given parameters"""
        with self._db.transaction() as cursor:
            query = """--sql
                SELECT
                    sq.*,
                    u.display_name as user_display_name,
                    u.email as user_email
                FROM session_queue sq
                LEFT JOIN users u ON sq.user_id = u.user_id
                WHERE sq.queue_id = ?
            """
            params: list[Union[str, int]] = [queue_id]

            if destination is not None:
                query += """---sql
                    AND sq.destination = ?
                """
                params.append(destination)

            query += """--sql
                ORDER BY
                    sq.priority DESC,
                    sq.item_id ASC
                ;
                """
            cursor.execute(query, params)
            results = cast(list[sqlite3.Row], cursor.fetchall())
        items = [SessionQueueItem.queue_item_from_dict(dict(result)) for result in results]
        return items

    def get_queue_item_ids(
        self,
        queue_id: str,
        order_dir: SQLiteDirection = SQLiteDirection.Descending,
        user_id: Optional[str] = None,
    ) -> ItemIdsResult:
        with self._db.transaction() as cursor_:
            query = """--sql
                SELECT item_id
                FROM session_queue
                WHERE queue_id = ?
                """
            query_params: list[str] = [queue_id]

            if user_id is not None:
                query += " AND user_id = ?"
                query_params.append(user_id)

            query += f" ORDER BY created_at {order_dir.value}"

            cursor_.execute(query, query_params)
            result = cast(list[sqlite3.Row], cursor_.fetchall())
        item_ids = [row[0] for row in result]

        return ItemIdsResult(item_ids=item_ids, total_count=len(item_ids))

    def get_queue_status(
        self,
        queue_id: str,
        user_id: Optional[str] = None,
        acting_user_id: Optional[str] = None,
    ) -> SessionQueueStatus:
        with self._db.transaction() as cursor:
            # Aggregate counts are always global (across all users). This lets a non-admin's
            # badge show "own / total" — their share of the whole queue — and lets the queue
            # list surface (redacted) entries belonging to other users.
            cursor.execute(
                """--sql
                SELECT status, count(*)
                FROM session_queue
                WHERE queue_id = ?
                GROUP BY status
                """,
                (queue_id,),
            )
            counts_result = cast(list[sqlite3.Row], cursor.fetchall())

            # When user_id is provided, additionally compute that user's own counts so the
            # caller can render the per-user portion of the badge. These are returned in the
            # separate user_pending/user_in_progress fields and never replace the global counts.
            user_counts_result: list[sqlite3.Row] = []
            if user_id is not None:
                cursor.execute(
                    """--sql
                    SELECT status, count(*)
                    FROM session_queue
                    WHERE queue_id = ? AND user_id = ?
                    GROUP BY status
                    """,
                    (queue_id, user_id),
                )
                user_counts_result = cast(list[sqlite3.Row], cursor.fetchall())

        current_item = self.get_current(queue_id=queue_id)
        total = sum(row[1] or 0 for row in counts_result)
        counts: dict[str, int] = {row[0]: row[1] for row in counts_result}

        user_pending: Optional[int] = None
        user_in_progress: Optional[int] = None
        if user_id is not None:
            user_counts: dict[str, int] = {row[0]: row[1] for row in user_counts_result}
            user_pending = user_counts.get("pending", 0)
            user_in_progress = user_counts.get("in_progress", 0)

        # Redaction is decided from the same current_item snapshot used to embed identifiers,
        # so a concurrent transition (e.g. B finishing while A's status changes) cannot leave
        # stale identifiers in the result. The aggregate counts stay global; only the current
        # item's identifiers are gated. acting_user_id (event path) takes precedence over
        # user_id (API path) when deciding the redaction owner; either being None means an
        # admin/global caller who may see the current item.
        owner_user_id = user_id if acting_user_id is None else acting_user_id
        show_current_item = current_item is not None and (
            owner_user_id is None or current_item.user_id == owner_user_id
        )

        return SessionQueueStatus(
            queue_id=queue_id,
            item_id=current_item.item_id if show_current_item else None,
            session_id=current_item.session_id if show_current_item else None,
            batch_id=current_item.batch_id if show_current_item else None,
            pending=counts.get("pending", 0),
            in_progress=counts.get("in_progress", 0),
            waiting=counts.get("waiting", 0),
            completed=counts.get("completed", 0),
            failed=counts.get("failed", 0),
            canceled=counts.get("canceled", 0),
            total=total,
            user_pending=user_pending,
            user_in_progress=user_in_progress,
        )

    def get_batch_status(self, queue_id: str, batch_id: str, user_id: Optional[str] = None) -> BatchStatus:
        with self._db.transaction() as cursor:
            query = """--sql
                SELECT status, count(*), origin, destination
                FROM session_queue
                WHERE queue_id = ? AND batch_id = ?
                """
            params: list[str] = [queue_id, batch_id]
            if user_id is not None:
                query += " AND user_id = ?"
                params.append(user_id)
            query += " GROUP BY status"
            cursor.execute(query, params)
            result = cast(list[sqlite3.Row], cursor.fetchall())
        total = sum(row[1] or 0 for row in result)
        counts: dict[str, int] = {row[0]: row[1] for row in result}
        origin = result[0]["origin"] if result else None
        destination = result[0]["destination"] if result else None

        return BatchStatus(
            batch_id=batch_id,
            origin=origin,
            destination=destination,
            queue_id=queue_id,
            pending=counts.get("pending", 0),
            in_progress=counts.get("in_progress", 0),
            waiting=counts.get("waiting", 0),
            completed=counts.get("completed", 0),
            failed=counts.get("failed", 0),
            canceled=counts.get("canceled", 0),
            total=total,
        )

    def get_counts_by_destination(
        self, queue_id: str, destination: str, user_id: Optional[str] = None
    ) -> SessionQueueCountsByDestination:
        with self._db.transaction() as cursor:
            query = """--sql
                SELECT status, count(*)
                FROM session_queue
                WHERE queue_id = ? AND destination = ?
                """
            params: list[str] = [queue_id, destination]
            if user_id is not None:
                query += " AND user_id = ?"
                params.append(user_id)
            query += " GROUP BY status"
            cursor.execute(query, params)
            counts_result = cast(list[sqlite3.Row], cursor.fetchall())

        total = sum(row[1] or 0 for row in counts_result)
        counts: dict[str, int] = {row[0]: row[1] for row in counts_result}

        return SessionQueueCountsByDestination(
            queue_id=queue_id,
            destination=destination,
            pending=counts.get("pending", 0),
            in_progress=counts.get("in_progress", 0),
            waiting=counts.get("waiting", 0),
            completed=counts.get("completed", 0),
            failed=counts.get("failed", 0),
            canceled=counts.get("canceled", 0),
            total=total,
        )

    def retry_items_by_id(self, queue_id: str, item_ids: list[int]) -> RetryItemsResult:
        """Retries the given queue items"""
        with self._db.transaction() as cursor:
            values_to_insert: list[ValueToInsertTuple] = []
            retried_root_item_ids: list[int] = []
            retried_user_ids: list[str] = []
            retried_item_ids_by_user: dict[str, list[int]] = {}
            seen_root_item_ids: set[int] = set()
            max_new_queue_items = self.__invoker.services.configuration.max_queue_size - self._get_current_queue_size(
                queue_id
            )

            if max_new_queue_items <= 0:
                return RetryItemsResult(queue_id=queue_id, retried_item_ids=[])

            for item_id in item_ids:
                try:
                    queue_item = self.get_queue_item(item_id)
                except SessionQueueItemNotFoundError:
                    continue
                if queue_item.queue_id != queue_id:
                    continue

                if queue_item.status not in ("failed", "canceled"):
                    continue

                root_item_id = queue_item.root_item_id or queue_item.item_id
                if root_item_id in seen_root_item_ids:
                    continue
                seen_root_item_ids.add(root_item_id)

                root_queue_item = self.get_queue_item(root_item_id)
                if root_queue_item.status not in ("failed", "canceled"):
                    continue

                retried_root_item_ids.append(root_item_id)
                retried_user_ids.append(root_queue_item.user_id)
                retried_item_ids_by_user.setdefault(root_queue_item.user_id, []).append(root_item_id)

                field_values_json = (
                    json.dumps(root_queue_item.field_values, default=to_jsonable_python)
                    if root_queue_item.field_values
                    else None
                )
                workflow_json = (
                    json.dumps(root_queue_item.workflow, default=to_jsonable_python)
                    if root_queue_item.workflow
                    else None
                )
                cloned_session = GraphExecutionState(graph=root_queue_item.session.graph)
                cloned_session_json = cloned_session.model_dump_json(warnings=False, exclude_none=True)

                retried_from_item_id = (
                    root_queue_item.retried_from_item_id
                    if root_queue_item.retried_from_item_id is not None
                    else root_queue_item.item_id
                )

                value_to_insert: ValueToInsertTuple = (
                    root_queue_item.queue_id,
                    cloned_session_json,
                    cloned_session.id,
                    root_queue_item.batch_id,
                    field_values_json,
                    root_queue_item.priority,
                    workflow_json,
                    root_queue_item.origin,
                    root_queue_item.destination,
                    retried_from_item_id,
                    root_queue_item.user_id,
                )
                values_to_insert.append(value_to_insert)

                if len(values_to_insert) >= max_new_queue_items:
                    break

            cursor.executemany(
                """--sql
                INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination, retried_from_item_id, user_id)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                values_to_insert,
            )

        retry_result = RetryItemsResult(
            queue_id=queue_id,
            retried_item_ids=retried_root_item_ids,
        )
        self.__invoker.services.events.emit_queue_items_retried(
            retry_result,
            user_ids=list(dict.fromkeys(retried_user_ids)),
            retried_item_ids_by_user=retried_item_ids_by_user,
        )
        return retry_result