chore: import upstream snapshot with attribution

tests/app/routers/test_session_queue_workflow_call.py

@@ -0,0 +1,430 @@
+"""Tests for session queue API behavior with workflow-call queue items."""
+
+import logging
+import uuid
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+from fastapi.testclient import TestClient
+
+from invokeai.app.api.dependencies import ApiDependencies
+from invokeai.app.api_app import app
+from invokeai.app.services.config.config_default import InvokeAIAppConfig
+from invokeai.app.services.invocation_services import InvocationServices
+from invokeai.app.services.invoker import Invoker
+from invokeai.app.services.session_processor.session_processor_common import SessionProcessorStatus
+from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue
+from invokeai.app.services.shared.graph import Graph, GraphExecutionState
+from invokeai.app.services.users.users_common import UserCreateRequest
+from invokeai.app.services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage
+from invokeai.backend.util.logging import InvokeAILogger
+from tests.fixtures.sqlite_database import create_mock_sqlite_database
+
+
+class MockApiDependencies(ApiDependencies):
+    invoker: Invoker
+
+    def __init__(self, invoker: Invoker) -> None:
+        self.invoker = invoker
+
+
+@pytest.fixture
+def setup_jwt_secret():
+    from invokeai.app.services.auth.token_service import set_jwt_secret
+
+    set_jwt_secret("test-secret-key-for-unit-tests-only-do-not-use-in-production")
+
+
+@pytest.fixture
+def client():
+    return TestClient(app)
+
+
+@pytest.fixture
+def mock_services() -> InvocationServices:
+    from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage
+    from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage
+    from invokeai.app.services.boards.boards_default import BoardService
+    from invokeai.app.services.bulk_download.bulk_download_default import BulkDownloadService
+    from invokeai.app.services.client_state_persistence.client_state_persistence_sqlite import (
+        ClientStatePersistenceSqlite,
+    )
+    from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage
+    from invokeai.app.services.images.images_default import ImageService
+    from invokeai.app.services.invocation_cache.invocation_cache_memory import MemoryInvocationCache
+    from invokeai.app.services.invocation_stats.invocation_stats_default import InvocationStatsService
+    from invokeai.app.services.users.users_default import UserService
+    from tests.test_nodes import TestEventService
+
+    configuration = InvokeAIAppConfig(use_memory_db=True, node_cache_size=0)
+    logger = InvokeAILogger.get_logger()
+    db = create_mock_sqlite_database(configuration, logger)
+
+    return InvocationServices(
+        board_image_records=SqliteBoardImageRecordStorage(db=db),
+        board_images=None,  # type: ignore
+        board_records=SqliteBoardRecordStorage(db=db),
+        boards=BoardService(),
+        bulk_download=BulkDownloadService(),
+        configuration=configuration,
+        events=TestEventService(),
+        image_files=None,  # type: ignore
+        image_records=SqliteImageRecordStorage(db=db),
+        images=ImageService(),
+        invocation_cache=MemoryInvocationCache(max_cache_size=0),
+        logger=logging,  # type: ignore
+        model_images=None,  # type: ignore
+        model_manager=None,  # type: ignore
+        download_queue=None,  # type: ignore
+        names=None,  # type: ignore
+        performance_statistics=InvocationStatsService(),
+        session_processor=None,  # type: ignore
+        session_queue=None,  # type: ignore
+        urls=None,  # type: ignore
+        workflow_records=SqliteWorkflowRecordsStorage(db=db),
+        tensors=None,  # type: ignore
+        conditioning=None,  # type: ignore
+        style_preset_records=None,  # type: ignore
+        style_preset_image_files=None,  # type: ignore
+        workflow_thumbnails=None,  # type: ignore
+        model_relationship_records=None,  # type: ignore
+        model_relationships=None,  # type: ignore
+        client_state_persistence=ClientStatePersistenceSqlite(db=db),
+        users=UserService(db),
+        external_generation=None,  # type: ignore
+    )
+
+
+@pytest.fixture
+def mock_invoker(mock_services: InvocationServices) -> Invoker:
+    invoker = Invoker(services=mock_services)
+    queue = SqliteSessionQueue(db=mock_services.board_records._db)
+    mock_services.session_queue = queue
+    mock_services.session_processor = MagicMock()
+    mock_services.session_processor.get_status.return_value = SessionProcessorStatus(
+        is_started=True, is_processing=False
+    )
+    queue.start(invoker)
+    return invoker
+
+
+def _create_user(mock_invoker: Invoker, email: str, display_name: str, is_admin: bool = False) -> str:
+    user = mock_invoker.services.users.create(
+        UserCreateRequest(email=email, display_name=display_name, password="TestPass123", is_admin=is_admin)
+    )
+    return user.user_id
+
+
+def _login(client: TestClient, email: str) -> str:
+    response = client.post("/api/v1/auth/login", json={"email": email, "password": "TestPass123", "remember_me": False})
+    assert response.status_code == 200
+    return response.json()["token"]
+
+
+def _auth(token: str) -> dict[str, str]:
+    return {"Authorization": f"Bearer {token}"}
+
+
+@pytest.fixture
+def enable_multiuser(monkeypatch: Any, mock_invoker: Invoker):
+    mock_invoker.services.configuration.multiuser = True
+    mock_deps = MockApiDependencies(mock_invoker)
+    monkeypatch.setattr("invokeai.app.api.routers.auth.ApiDependencies", mock_deps)
+    monkeypatch.setattr("invokeai.app.api.auth_dependencies.ApiDependencies", mock_deps)
+    monkeypatch.setattr("invokeai.app.api.routers.session_queue.ApiDependencies", mock_deps)
+    yield
+
+
+@pytest.fixture
+def admin_token(setup_jwt_secret: None, enable_multiuser: Any, mock_invoker: Invoker, client: TestClient):
+    _create_user(mock_invoker, "admin@test.com", "Admin", is_admin=True)
+    return _login(client, "admin@test.com")
+
+
+@pytest.fixture
+def user1_token(enable_multiuser: Any, mock_invoker: Invoker, client: TestClient, admin_token: str):
+    _create_user(mock_invoker, "user1@test.com", "User One")
+    return _login(client, "user1@test.com")
+
+
+@pytest.fixture
+def user2_token(enable_multiuser: Any, mock_invoker: Invoker, client: TestClient, admin_token: str):
+    _create_user(mock_invoker, "user2@test.com", "User Two")
+    return _login(client, "user2@test.com")
+
+
+def _insert_queue_item(
+    session_queue: SqliteSessionQueue,
+    *,
+    queue_id: str = "default",
+    user_id: str,
+    status: str,
+    session: GraphExecutionState | None = None,
+    workflow_call_id: str | None = None,
+    parent_item_id: int | None = None,
+    parent_session_id: str | None = None,
+    root_item_id: int | None = None,
+    workflow_call_depth: int | None = None,
+) -> int:
+    session = session or GraphExecutionState(graph=Graph())
+    with session_queue._db.transaction() as cursor:
+        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, status, workflow_call_id, parent_item_id, parent_session_id,
+                root_item_id, workflow_call_depth
+            )
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                queue_id,
+                session.model_dump_json(warnings=False),
+                session.id,
+                str(uuid.uuid4()),
+                None,
+                0,
+                None,
+                None,
+                None,
+                None,
+                user_id,
+                status,
+                workflow_call_id,
+                parent_item_id,
+                parent_session_id,
+                root_item_id,
+                workflow_call_depth,
+            ),
+        )
+        return cursor.lastrowid
+
+
+def test_get_queue_status_reports_waiting_for_owner(
+    client: TestClient, mock_invoker: Invoker, user1_token: str, user2_token: str
+) -> None:
+    user1 = mock_invoker.services.users.get_by_email("user1@test.com")
+    user2 = mock_invoker.services.users.get_by_email("user2@test.com")
+    assert user1 is not None and user2 is not None
+
+    _insert_queue_item(mock_invoker.services.session_queue, user_id=user1.user_id, status="waiting")
+    _insert_queue_item(mock_invoker.services.session_queue, user_id=user2.user_id, status="pending")
+
+    response = client.get("/api/v1/queue/default/status", headers=_auth(user1_token))
+
+    assert response.status_code == 200
+    payload = response.json()
+    assert payload["queue"]["waiting"] == 1
+    assert payload["queue"]["pending"] == 1
+    assert payload["queue"]["in_progress"] == 0
+    assert payload["queue"]["total"] == 2
+    assert payload["queue"]["user_pending"] == 0
+    assert payload["queue"]["user_in_progress"] == 0
+    assert payload["queue"]["item_id"] is None
+
+
+def test_get_queue_item_sanitizes_workflow_call_metadata_for_non_owner(
+    client: TestClient, mock_invoker: Invoker, user1_token: str, user2_token: str
+) -> None:
+    user1 = mock_invoker.services.users.get_by_email("user1@test.com")
+    assert user1 is not None
+
+    item_id = _insert_queue_item(
+        mock_invoker.services.session_queue,
+        user_id=user1.user_id,
+        status="waiting",
+        workflow_call_id="workflow-call-1",
+        parent_item_id=41,
+        parent_session_id="parent-session-1",
+        root_item_id=17,
+        workflow_call_depth=2,
+    )
+
+    response = client.get(f"/api/v1/queue/default/i/{item_id}", headers=_auth(user2_token))
+
+    assert response.status_code == 200
+    payload = response.json()
+    assert payload["status"] == "waiting"
+    assert payload["item_id"] == item_id
+    assert payload["user_id"] == "redacted"
+    assert payload["batch_id"] == "redacted"
+    assert payload["session_id"] == "redacted"
+    assert payload.get("workflow_call_id") is None
+    assert payload.get("parent_item_id") is None
+    assert payload.get("parent_session_id") is None
+    assert payload.get("root_item_id") is None
+    assert payload.get("workflow_call_depth") is None
+
+
+def test_retry_items_by_id_normalizes_child_to_root_at_router(
+    client: TestClient, mock_invoker: Invoker, user1_token: str
+) -> None:
+    user1 = mock_invoker.services.users.get_by_email("user1@test.com")
+    assert user1 is not None
+
+    root_item_id = _insert_queue_item(mock_invoker.services.session_queue, user_id=user1.user_id, status="failed")
+    child_item_id = _insert_queue_item(
+        mock_invoker.services.session_queue,
+        user_id=user1.user_id,
+        status="failed",
+        workflow_call_id="workflow-call-1",
+        parent_item_id=root_item_id,
+        parent_session_id="parent-session-1",
+        root_item_id=root_item_id,
+        workflow_call_depth=1,
+    )
+
+    response = client.put("/api/v1/queue/default/retry_items_by_id", headers=_auth(user1_token), json=[child_item_id])
+
+    assert response.status_code == 200
+    assert response.json()["retried_item_ids"] == [root_item_id]
+
+
+def test_retry_items_by_id_rejects_items_from_other_queue(
+    client: TestClient, mock_invoker: Invoker, user1_token: str
+) -> None:
+    user1 = mock_invoker.services.users.get_by_email("user1@test.com")
+    assert user1 is not None
+
+    item_id = _insert_queue_item(
+        mock_invoker.services.session_queue,
+        queue_id="secondary",
+        user_id=user1.user_id,
+        status="failed",
+    )
+
+    response = client.put("/api/v1/queue/default/retry_items_by_id", headers=_auth(user1_token), json=[item_id])
+
+    assert response.status_code == 404
+
+
+def test_retry_items_by_id_authorizes_workflow_call_root_owner(
+    client: TestClient, mock_invoker: Invoker, user1_token: str, user2_token: str
+) -> None:
+    user1 = mock_invoker.services.users.get_by_email("user1@test.com")
+    user2 = mock_invoker.services.users.get_by_email("user2@test.com")
+    assert user1 is not None and user2 is not None
+
+    root_item_id = _insert_queue_item(mock_invoker.services.session_queue, user_id=user2.user_id, status="failed")
+    child_item_id = _insert_queue_item(
+        mock_invoker.services.session_queue,
+        user_id=user1.user_id,
+        status="failed",
+        workflow_call_id="workflow-call-1",
+        parent_item_id=root_item_id,
+        parent_session_id="parent-session-1",
+        root_item_id=root_item_id,
+        workflow_call_depth=1,
+    )
+
+    response = client.put("/api/v1/queue/default/retry_items_by_id", headers=_auth(user1_token), json=[child_item_id])
+
+    assert response.status_code == 403
+    assert mock_invoker.services.session_queue.get_queue_item(root_item_id).status == "failed"
+
+
+def test_retry_items_by_id_skips_missing_items_and_retries_valid_items(
+    client: TestClient, mock_invoker: Invoker, user1_token: str
+) -> None:
+    user1 = mock_invoker.services.users.get_by_email("user1@test.com")
+    assert user1 is not None
+
+    root_item_id = _insert_queue_item(mock_invoker.services.session_queue, user_id=user1.user_id, status="failed")
+
+    response = client.put(
+        "/api/v1/queue/default/retry_items_by_id",
+        headers=_auth(user1_token),
+        json=[root_item_id + 9999, root_item_id],
+    )
+
+    assert response.status_code == 200
+    assert response.json()["retried_item_ids"] == [root_item_id]
+
+
+def test_cancel_queue_item_cascades_to_waiting_parent_via_router(
+    client: TestClient, mock_invoker: Invoker, user1_token: str
+) -> None:
+    user1 = mock_invoker.services.users.get_by_email("user1@test.com")
+    assert user1 is not None
+
+    parent_item_id = _insert_queue_item(mock_invoker.services.session_queue, user_id=user1.user_id, status="waiting")
+    child_item_id = _insert_queue_item(
+        mock_invoker.services.session_queue,
+        user_id=user1.user_id,
+        status="pending",
+        workflow_call_id="workflow-call-1",
+        parent_item_id=parent_item_id,
+        parent_session_id="parent-session-1",
+        root_item_id=parent_item_id,
+        workflow_call_depth=1,
+    )
+
+    response = client.put(f"/api/v1/queue/default/i/{child_item_id}/cancel", headers=_auth(user1_token))
+
+    assert response.status_code == 200
+    assert response.json()["status"] == "canceled"
+    assert mock_invoker.services.session_queue.get_queue_item(parent_item_id).status == "canceled"
+    assert mock_invoker.services.session_queue.get_queue_item(child_item_id).status == "canceled"
+
+
+def test_cancel_queue_item_rejects_item_from_other_queue(
+    client: TestClient, mock_invoker: Invoker, user1_token: str
+) -> None:
+    user1 = mock_invoker.services.users.get_by_email("user1@test.com")
+    assert user1 is not None
+
+    item_id = _insert_queue_item(
+        mock_invoker.services.session_queue,
+        queue_id="secondary",
+        user_id=user1.user_id,
+        status="pending",
+    )
+
+    response = client.put(f"/api/v1/queue/default/i/{item_id}/cancel", headers=_auth(user1_token))
+
+    assert response.status_code == 404
+
+
+def test_delete_queue_item_rejects_item_from_other_queue(
+    client: TestClient, mock_invoker: Invoker, user1_token: str
+) -> None:
+    user1 = mock_invoker.services.users.get_by_email("user1@test.com")
+    assert user1 is not None
+
+    item_id = _insert_queue_item(
+        mock_invoker.services.session_queue,
+        queue_id="secondary",
+        user_id=user1.user_id,
+        status="failed",
+    )
+
+    response = client.delete(f"/api/v1/queue/default/i/{item_id}", headers=_auth(user1_token))
+
+    assert response.status_code == 404
+
+
+def test_delete_queue_item_authorizes_workflow_call_root_owner(
+    client: TestClient, mock_invoker: Invoker, user1_token: str, user2_token: str
+) -> None:
+    user1 = mock_invoker.services.users.get_by_email("user1@test.com")
+    user2 = mock_invoker.services.users.get_by_email("user2@test.com")
+    assert user1 is not None and user2 is not None
+
+    root_item_id = _insert_queue_item(mock_invoker.services.session_queue, user_id=user2.user_id, status="failed")
+    child_item_id = _insert_queue_item(
+        mock_invoker.services.session_queue,
+        user_id=user1.user_id,
+        status="failed",
+        workflow_call_id="workflow-call-1",
+        parent_item_id=root_item_id,
+        parent_session_id="parent-session-1",
+        root_item_id=root_item_id,
+        workflow_call_depth=1,
+    )
+
+    response = client.delete(f"/api/v1/queue/default/i/{child_item_id}", headers=_auth(user1_token))
+
+    assert response.status_code == 403
+    assert mock_invoker.services.session_queue.get_queue_item(root_item_id).status == "failed"
+    assert mock_invoker.services.session_queue.get_queue_item(child_item_id).status == "failed"