From 5eb04d15345bbbf2a3ef06a23ccf24937e161f05 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Fri, 17 Apr 2026 00:13:35 -0400 Subject: [PATCH] feat(invoke): add "Use as Ref Image" button and resilient auth fallback Adds a third button to the InvokeAI metadata drawer that uploads the current image to InvokeAI and sets it as a reference image parameter, without resetting any other generation parameters. Also reworks the JWT auth flow so PhotoMap stays in sync with InvokeAI when it is switched between single-user and multi-user modes: the first request always uses whatever token is cached (or none), a 401 triggers a login + retry with Bearer, and a 403 received with a token in hand triggers an anonymous retry and discards the stale cache. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../metadata_modules/invoke_formatter.py | 26 +- photomap/backend/routers/invoke.py | 254 ++++++++- .../static/javascript/invoke-recall.js | 60 ++- tests/backend/test_invoke_metadata.py | 5 +- tests/backend/test_invoke_router.py | 497 ++++++++++++++++++ tests/frontend/invoke-recall.test.js | 54 +- 6 files changed, 846 insertions(+), 50 deletions(-) diff --git a/photomap/backend/metadata_modules/invoke_formatter.py b/photomap/backend/metadata_modules/invoke_formatter.py index f83017c..f2915ed 100644 --- a/photomap/backend/metadata_modules/invoke_formatter.py +++ b/photomap/backend/metadata_modules/invoke_formatter.py @@ -61,19 +61,37 @@ "" ) +# Photo-frame icon — a rectangle with a small sun and a mountain inside, +# drawn in the same stroked style as the remix icon so the three buttons +# share a consistent visual language. +_USE_REF_SVG = ( + '" +) + def _recall_buttons_html() -> str: - """Render the recall / remix button group shown at the bottom of the drawer.""" + """Render the recall / remix / use-ref button group shown at the bottom of the drawer.""" return ( '
' + '" '" - '" "
" diff --git a/photomap/backend/routers/invoke.py b/photomap/backend/routers/invoke.py index 3d11bd2..02ead41 100644 --- a/photomap/backend/routers/invoke.py +++ b/photomap/backend/routers/invoke.py @@ -8,17 +8,27 @@ tuple, load the image metadata, build a recall payload from it, and proxy it to the configured InvokeAI backend's ``/api/v1/recall/{queue_id}`` endpoint. +* ``POST /invokeai/use_ref_image`` — upload the selected image to InvokeAI + and then call the same recall endpoint with the uploaded image as a + reference image parameter, so the next generation uses it for visual + guidance. When the configured InvokeAI backend runs in multi-user mode, the ``username`` / ``password`` fields are used to obtain a JWT bearer token via ``/api/v1/auth/login``. The token is cached in-process and -automatically refreshed on 401. +automatically refreshed on 401. If the backend has since been +reconfigured into single-user mode it will reject a token-bearing +request with a 403 — that causes the cached token to be discarded and +the call to be retried anonymously without requiring a restart. """ from __future__ import annotations import logging +import mimetypes import time +from collections.abc import Awaitable, Callable +from pathlib import Path import httpx from fastapi import APIRouter, HTTPException @@ -46,18 +56,15 @@ _token_username: str | None = None -async def _get_auth_headers(base_url: str, username: str | None, password: str | None) -> dict[str, str]: - """Return an ``Authorization`` header dict, or empty dict for anonymous access. +def _cached_auth_headers(base_url: str, username: str | None) -> dict[str, str]: + """Return ``{"Authorization": "Bearer ..."}`` if we still hold a valid + cached token for this ``(base_url, username)`` pair, else ``{}``. - If a valid cached token exists it is reused. If no token is cached the - caller should first try the request without credentials — the InvokeAI - backend running in single-user mode will accept it. When the caller - receives a 401 it should call ``_invalidate_token_cache`` and call this - function again; the second call will perform the login. + This never talks to the network. Deliberate: the first attempt at any + request always uses whatever auth we already have (or none), so that a + backend that has since been reconfigured into single-user mode is given + a chance to accept the call anonymously. """ - global _cached_token, _token_expires_at, _token_base_url, _token_username # noqa: PLW0603 - - # Reuse cached token if still valid for this backend+user if ( _cached_token and time.monotonic() < _token_expires_at @@ -65,12 +72,15 @@ async def _get_auth_headers(base_url: str, username: str | None, password: str | and _token_username == username ): return {"Authorization": f"Bearer {_cached_token}"} + return {} + - # No credentials configured — anonymous access - if not username or not password: - return {} +async def _login(base_url: str, username: str, password: str) -> dict[str, str]: + """Exchange ``username``/``password`` for a JWT via the InvokeAI auth + endpoint, cache the token, and return the ``Authorization`` header. + """ + global _cached_token, _token_expires_at, _token_base_url, _token_username # noqa: PLW0603 - # No cached token yet — try to obtain one login_url = f"{base_url.rstrip('/')}/api/v1/auth/login" try: async with httpx.AsyncClient(timeout=_HTTP_TIMEOUT) as client: @@ -99,9 +109,50 @@ async def _get_auth_headers(base_url: str, username: str | None, password: str | def _invalidate_token_cache() -> None: """Clear the cached token so the next request re-authenticates.""" - global _cached_token, _token_expires_at # noqa: PLW0603 + global _cached_token, _token_expires_at, _token_base_url, _token_username # noqa: PLW0603 _cached_token = None _token_expires_at = 0.0 + _token_base_url = None + _token_username = None + + +async def _post_with_auth_fallback( + base_url: str, + username: str | None, + password: str | None, + request_fn: Callable[[dict[str, str]], Awaitable[httpx.Response]], +) -> httpx.Response: + """Perform an InvokeAI request with graceful handling of auth transitions. + + ``request_fn`` is an async callable that takes a headers dict and + performs the HTTP call — using a factory lets the caller re-open file + streams (needed for multipart uploads) on a retry. + + Three-step flow: + + 1. First attempt uses whatever token we have cached (or no auth at all). + A freshly-restarted single-user backend then accepts the call even + if credentials are stored in PhotoMap. + 2. If the first attempt returns **401**, the backend demands + authentication: if credentials are configured we log in, cache a + fresh token, and retry. + 3. If the first attempt was made *with* a token and returns **403** + (most commonly "Multiuser mode is disabled. Authentication is not + required…"), the backend was reconfigured to single-user mode — we + invalidate the cached token and retry anonymously. + """ + auth_headers = _cached_auth_headers(base_url, username) + response = await request_fn(auth_headers) + + if response.status_code == 401 and username and password: + _invalidate_token_cache() + auth_headers = await _login(base_url, username, password) + response = await request_fn(auth_headers) + elif response.status_code == 403 and auth_headers: + _invalidate_token_cache() + response = await request_fn({}) + + return response class InvokeAISettings(BaseModel): @@ -124,6 +175,14 @@ class RecallRequest(BaseModel): queue_id: str = Field("default", description="InvokeAI queue id to target") +class UseRefImageRequest(BaseModel): + """Payload posted by the drawer's "Use Ref Image" button.""" + + album_key: str = Field(..., description="Album containing the image") + index: int = Field(..., ge=0, description="Image index within the album") + queue_id: str = Field("default", description="InvokeAI queue id to target") + + @invoke_router.get("/config") async def get_invokeai_config() -> dict: """Return the persisted InvokeAI connection settings. @@ -175,6 +234,17 @@ def _load_raw_metadata(album_key: str, index: int) -> dict: return entry if isinstance(entry, dict) else {} +def _load_image_path(album_key: str, index: int) -> Path: + embeddings = get_embeddings_for_album(album_key) + if not embeddings: + raise HTTPException(status_code=404, detail="Album not found") + indexes = embeddings.indexes + filenames = indexes["sorted_filenames"] + if index < 0 or index >= len(filenames): + raise HTTPException(status_code=404, detail="Index out of range") + return Path(str(filenames[index])) + + def _build_recall_payload(raw_metadata: dict, include_seed: bool) -> dict: if not raw_metadata: raise HTTPException( @@ -222,18 +292,17 @@ async def recall_parameters(request: RecallRequest) -> dict: username = settings["username"] password = settings["password"] - # Try without credentials first (single-user mode). If the backend - # requires authentication (401), obtain a token and retry. - auth_headers = await _get_auth_headers(base_url, username, password) - try: async with httpx.AsyncClient(timeout=_HTTP_TIMEOUT) as client: - response = await client.post(url, json=payload, params={"strict": "true"}, headers=auth_headers) - if response.status_code == 401 and username and password: - _invalidate_token_cache() - auth_headers = await _get_auth_headers(base_url, username, password) - response = await client.post(url, json=payload, params={"strict": "true"}, headers=auth_headers) + async def _do(headers: dict[str, str]) -> httpx.Response: + return await client.post( + url, json=payload, params={"strict": "true"}, headers=headers + ) + + response = await _post_with_auth_fallback( + base_url, username, password, _do + ) except httpx.RequestError as exc: logger.warning("InvokeAI recall request failed: %s", exc) raise HTTPException( @@ -262,3 +331,138 @@ async def recall_parameters(request: RecallRequest) -> dict: "sent": payload, "response": remote, } + + +async def _upload_image_to_invokeai( + client: httpx.AsyncClient, + base_url: str, + image_path: Path, + username: str | None, + password: str | None, +) -> str: + """Upload ``image_path`` to InvokeAI and return the assigned ``image_name``. + + Auth transitions (anonymous ↔ token) are handled transparently by + ``_post_with_auth_fallback``; the multipart stream is re-opened on each + retry since the previous one will have been consumed. + """ + upload_url = f"{base_url.rstrip('/')}/api/v1/images/upload" + mime_type = mimetypes.guess_type(image_path.name)[0] or "image/png" + params = {"image_category": "user", "is_intermediate": "false"} + + async def _do(headers: dict[str, str]) -> httpx.Response: + with image_path.open("rb") as fh: + files = {"file": (image_path.name, fh, mime_type)} + return await client.post( + upload_url, files=files, params=params, headers=headers + ) + + upload_resp = await _post_with_auth_fallback(base_url, username, password, _do) + + if upload_resp.status_code >= 400: + raise HTTPException( + status_code=502, + detail=( + f"InvokeAI image upload returned {upload_resp.status_code}: " + f"{upload_resp.text[:200]}" + ), + ) + + try: + body = upload_resp.json() + except ValueError as exc: + raise HTTPException( + status_code=502, + detail="InvokeAI image upload returned a non-JSON response", + ) from exc + + image_name = body.get("image_name") + if not image_name: + raise HTTPException( + status_code=502, + detail="InvokeAI image upload response did not include image_name", + ) + return image_name + + +@invoke_router.post("/use_ref_image") +async def use_ref_image(request: UseRefImageRequest) -> dict: + """Upload the selected image to InvokeAI and set it as a reference image. + + Implements the two-step flow: first ``POST /api/v1/images/upload`` so + InvokeAI knows the file, then ``POST /api/v1/recall/{queue_id}`` with the + returned ``image_name`` in ``reference_images`` so the next generation + picks it up as visual guidance. + """ + settings = config_manager.get_invokeai_settings() + base_url = settings["url"] + if not base_url: + raise HTTPException( + status_code=400, + detail=( + "InvokeAI backend URL is not configured. Set it in the " + "PhotoMap settings panel." + ), + ) + + image_path = _load_image_path(request.album_key, request.index) + if not image_path.is_file(): + raise HTTPException( + status_code=404, detail=f"Image file not found on disk: {image_path.name}" + ) + + username = settings["username"] + password = settings["password"] + + recall_url = f"{base_url.rstrip('/')}/api/v1/recall/{request.queue_id}" + + try: + async with httpx.AsyncClient(timeout=_HTTP_TIMEOUT) as client: + image_name = await _upload_image_to_invokeai( + client, base_url, image_path, username, password + ) + + # Deliberately omit ``strict=true`` so that the recall only + # *adds* the reference image to whatever the user already has + # set up in InvokeAI rather than resetting every other + # parameter back to defaults. + payload = {"reference_images": [{"image_name": image_name}]} + + async def _do_recall(headers: dict[str, str]) -> httpx.Response: + return await client.post(recall_url, json=payload, headers=headers) + + response = await _post_with_auth_fallback( + base_url, username, password, _do_recall + ) + except httpx.RequestError as exc: + logger.warning("InvokeAI use_ref_image request failed: %s", exc) + raise HTTPException( + status_code=502, + detail=f"Could not reach InvokeAI backend at {base_url}: {exc}", + ) from exc + + if response.status_code >= 400: + logger.warning( + "InvokeAI use_ref_image recall returned %s: %s", + response.status_code, + response.text, + ) + raise HTTPException( + status_code=502, + detail=( + f"InvokeAI backend returned {response.status_code}: " + f"{response.text[:200]}" + ), + ) + + try: + remote = response.json() + except ValueError: + remote = {"raw": response.text} + + return { + "success": True, + "sent": payload, + "uploaded_image_name": image_name, + "response": remote, + } diff --git a/photomap/frontend/static/javascript/invoke-recall.js b/photomap/frontend/static/javascript/invoke-recall.js index c98d0e5..fbe73f0 100644 --- a/photomap/frontend/static/javascript/invoke-recall.js +++ b/photomap/frontend/static/javascript/invoke-recall.js @@ -75,6 +75,21 @@ function showErrorMessage(button, message) { setTimeout(() => banner.remove(), STATUS_RESET_MS * 3); } +async function _raiseForStatus(response) { + let message = `HTTP ${response.status}`; + try { + const body = await response.json(); + if (body && body.detail) { + message = body.detail; + } + } catch { + // ignore JSON parse errors — fall back to the generic message + } + const err = new Error(message); + err.status = response.status; + throw err; +} + export async function sendRecall({ albumKey, index, includeSeed }) { const response = await fetch("invokeai/recall", { method: "POST", @@ -86,18 +101,22 @@ export async function sendRecall({ albumKey, index, includeSeed }) { }), }); if (!response.ok) { - let message = `HTTP ${response.status}`; - try { - const body = await response.json(); - if (body && body.detail) { - message = body.detail; - } - } catch { - // ignore JSON parse errors — fall back to the generic message - } - const err = new Error(message); - err.status = response.status; - throw err; + await _raiseForStatus(response); + } + return response.json(); +} + +export async function sendUseRefImage({ albumKey, index }) { + const response = await fetch("invokeai/use_ref_image", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + album_key: albumKey, + index, + }), + }); + if (!response.ok) { + await _raiseForStatus(response); } return response.json(); } @@ -129,16 +148,21 @@ async function handleRecallClick(button) { button.disabled = true; try { - await sendRecall({ - albumKey, - index: parsed.index, - includeSeed: mode !== "remix", - }); + if (mode === "use_ref") { + await sendUseRefImage({ albumKey, index: parsed.index }); + } else { + await sendRecall({ + albumKey, + index: parsed.index, + includeSeed: mode !== "remix", + }); + } showStatus(button, "success"); } catch (err) { console.error("InvokeAI recall failed:", err); showStatus(button, "error"); - showErrorMessage(button, err && err.message ? err.message : "Recall failed"); + const fallback = mode === "use_ref" ? "Use as Ref Image failed" : "Recall failed"; + showErrorMessage(button, err && err.message ? err.message : fallback); } finally { button.disabled = false; } diff --git a/tests/backend/test_invoke_metadata.py b/tests/backend/test_invoke_metadata.py index 387454c..7c64e54 100644 --- a/tests/backend/test_invoke_metadata.py +++ b/tests/backend/test_invoke_metadata.py @@ -892,5 +892,6 @@ def test_buttons_shown_when_enabled(self, v3_metadata): assert 'class="invoke-recall-controls"' in html assert 'data-recall-mode="recall"' in html assert 'data-recall-mode="remix"' in html - # Asterisk SVG is present for the recall button - assert html.count('class="invoke-recall-btn"') == 2 + assert 'data-recall-mode="use_ref"' in html + assert "Use as Ref Image" in html + assert html.count('class="invoke-recall-btn"') == 3 diff --git a/tests/backend/test_invoke_router.py b/tests/backend/test_invoke_router.py index 9c089f5..c2ac3a2 100644 --- a/tests/backend/test_invoke_router.py +++ b/tests/backend/test_invoke_router.py @@ -26,6 +26,16 @@ def clear_invokeai_config(): manager.set_invokeai_settings(url=None, username=None, password=None) +@pytest.fixture +def clear_token_cache(): + """Wipe the module-level JWT cache before and after each test.""" + from photomap.backend.routers import invoke as invoke_module + + invoke_module._invalidate_token_cache() + yield + invoke_module._invalidate_token_cache() + + def test_get_config_empty(client, clear_invokeai_config): response = client.get("/invokeai/config") assert response.status_code == 200 @@ -202,6 +212,149 @@ async def post(self, url, json, **kwargs): assert "seed" not in captured["json"] +def test_use_ref_image_requires_configured_url(client, clear_invokeai_config): + response = client.post( + "/invokeai/use_ref_image", + json={"album_key": "whatever", "index": 0}, + ) + assert response.status_code == 400 + assert "not configured" in response.json()["detail"].lower() + + +def test_use_ref_image_uploads_then_calls_recall_without_strict( + client, clear_invokeai_config, monkeypatch, tmp_path +): + """Happy path — verify upload + recall ordering, payload, and no strict=true.""" + client.post("/invokeai/config", json={"url": "http://localhost:9090"}) + + # Create a real file on disk so _load_image_path + is_file() pass. + image_file = tmp_path / "pic.png" + image_file.write_bytes(b"\x89PNG\r\n\x1a\nfakebytes") + + from photomap.backend.routers import invoke as invoke_module + + monkeypatch.setattr( + invoke_module, "_load_image_path", lambda album_key, index: image_file + ) + + calls: list[dict] = [] + + class _UploadResponse: + status_code = 200 + text = "" + + def json(self): + return {"image_name": "uploaded-abc.png"} + + class _RecallResponse: + status_code = 200 + text = "{}" + + def json(self): + return {"status": "success"} + + class _StubClient: + def __init__(self, *args, **kwargs): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + async def post(self, url, **kwargs): + call = {"url": url, "params": kwargs.get("params")} + if "files" in kwargs: + call["kind"] = "upload" + # Drain the file stream so httpx-style behavior is preserved. + kwargs["files"]["file"][1].read() + calls.append(call) + return _UploadResponse() + call["kind"] = "recall" + call["json"] = kwargs.get("json") + calls.append(call) + return _RecallResponse() + + monkeypatch.setattr(invoke_module.httpx, "AsyncClient", _StubClient) + + response = client.post( + "/invokeai/use_ref_image", + json={"album_key": "any", "index": 0}, + ) + assert response.status_code == 200, response.text + body = response.json() + assert body["success"] is True + assert body["uploaded_image_name"] == "uploaded-abc.png" + assert body["sent"] == { + "reference_images": [{"image_name": "uploaded-abc.png"}] + } + + # Two upstream calls, in the right order, to the right URLs. + assert len(calls) == 2 + assert calls[0]["kind"] == "upload" + assert calls[0]["url"] == "http://localhost:9090/api/v1/images/upload" + assert calls[0]["params"] == { + "image_category": "user", + "is_intermediate": "false", + } + assert calls[1]["kind"] == "recall" + assert calls[1]["url"] == "http://localhost:9090/api/v1/recall/default" + # CRITICAL: no `strict=true` — sending it would reset every other + # parameter the user already has set up in InvokeAI. + assert calls[1]["params"] is None or "strict" not in (calls[1]["params"] or {}) + assert calls[1]["json"] == { + "reference_images": [{"image_name": "uploaded-abc.png"}] + } + + +def test_use_ref_image_upload_failure_returns_502( + client, clear_invokeai_config, monkeypatch, tmp_path +): + client.post("/invokeai/config", json={"url": "http://localhost:9090"}) + + image_file = tmp_path / "pic.png" + image_file.write_bytes(b"\x89PNG") + + from photomap.backend.routers import invoke as invoke_module + + monkeypatch.setattr( + invoke_module, "_load_image_path", lambda album_key, index: image_file + ) + + class _FailedUpload: + status_code = 500 + text = "disk full" + + def json(self): + return {"detail": "disk full"} + + class _StubClient: + def __init__(self, *args, **kwargs): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + async def post(self, url, **kwargs): + if "files" in kwargs: + kwargs["files"]["file"][1].read() + return _FailedUpload() + raise AssertionError("Recall should not be attempted when upload fails") + + monkeypatch.setattr(invoke_module.httpx, "AsyncClient", _StubClient) + + response = client.post( + "/invokeai/use_ref_image", + json={"album_key": "any", "index": 0}, + ) + assert response.status_code == 502 + assert "upload" in response.json()["detail"].lower() + + def test_recall_upstream_unreachable_returns_502( client, clear_invokeai_config, monkeypatch ): @@ -241,3 +394,347 @@ async def post(self, url, json, **kwargs): ) assert response.status_code == 502 assert "Could not reach InvokeAI backend" in response.json()["detail"] + + +# ── Auth fallback behavior ───────────────────────────────────────────── + + +def _install_recall_stub(monkeypatch): + """Install a trivial metadata stub so recall requests don't need real embeddings.""" + from photomap.backend.routers import invoke as invoke_module + + raw_metadata = { + "metadata_version": 3, + "app_version": "3.5.0", + "positive_prompt": "x", + "model": {"model_name": "m", "base_model": "sd-1"}, + } + monkeypatch.setattr( + invoke_module, "_load_raw_metadata", lambda album_key, index: raw_metadata + ) + + +class _ScriptedClient: + """An httpx.AsyncClient stub whose ``post`` returns the next scripted response. + + Each call records the (url, headers) it saw so tests can assert on the + auth header progression across a retry cycle. + """ + + def __init__(self, script): + self._script = list(script) + self.calls: list[dict] = [] + + def __call__(self, *args, **kwargs): + # Invoked as ``httpx.AsyncClient(...)`` — return self so the ``async + # with`` context manager works. + return self + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + async def post(self, url, **kwargs): + self.calls.append( + { + "url": url, + "headers": dict(kwargs.get("headers") or {}), + "json": kwargs.get("json"), + } + ) + entry = self._script.pop(0) + if callable(entry): + entry = entry(url, kwargs) + return entry + + +class _Resp: + def __init__(self, status_code=200, json_body=None, text="{}"): + self.status_code = status_code + self._json = json_body if json_body is not None else {} + self.text = text + self.headers = {"content-type": "application/json"} + + def json(self): + return self._json + + +def test_recall_first_tries_without_auth_then_logs_in_on_401( + client, clear_invokeai_config, clear_token_cache, monkeypatch +): + """Initial call is anonymous; a 401 response triggers login + retry with Bearer.""" + client.post( + "/invokeai/config", + json={ + "url": "http://localhost:9090", + "username": "alice", + "password": "secret", + }, + ) + _install_recall_stub(monkeypatch) + + from photomap.backend.routers import invoke as invoke_module + + # Script the backend: first /recall returns 401, login returns a token, + # second /recall returns success. Dispatch by URL so we don't depend on + # call ordering between the recall client and the login client. + login_resp = _Resp( + 200, json_body={"token": "abc-token", "expires_in": 3600} + ) + recall_responses = [ + _Resp(401, json_body={"detail": "Not authenticated"}, text="unauth"), + _Resp(200, json_body={"status": "success"}), + ] + + def _route(url, kwargs): + if url.endswith("/auth/login"): + return login_resp + return recall_responses.pop(0) + + script = [_route, _route, _route] # may be invoked up to 3 times + stub = _ScriptedClient(script) + monkeypatch.setattr(invoke_module.httpx, "AsyncClient", stub) + + response = client.post( + "/invokeai/recall", + json={"album_key": "any", "index": 0, "include_seed": True}, + ) + assert response.status_code == 200, response.text + + # First recall attempt must have been anonymous. + recall_calls = [c for c in stub.calls if c["url"].endswith("/recall/default")] + assert len(recall_calls) == 2 + assert "Authorization" not in recall_calls[0]["headers"] + # Second recall attempt must carry the Bearer token from the login response. + assert recall_calls[1]["headers"].get("Authorization") == "Bearer abc-token" + + # A login call must have been made between the two recall attempts. + login_calls = [c for c in stub.calls if c["url"].endswith("/auth/login")] + assert len(login_calls) == 1 + assert login_calls[0]["json"] == {"email": "alice", "password": "secret"} + + +def test_recall_sends_cached_token_on_subsequent_requests( + client, clear_invokeai_config, clear_token_cache, monkeypatch +): + """Once a token is cached, later calls should send it from the first attempt.""" + client.post( + "/invokeai/config", + json={ + "url": "http://localhost:9090", + "username": "alice", + "password": "secret", + }, + ) + + # Pre-seed the token cache as though a previous login had succeeded. + import time as _time + + from photomap.backend.routers import invoke as invoke_module + + invoke_module._cached_token = "cached-token" + invoke_module._token_expires_at = _time.monotonic() + 3600 + invoke_module._token_base_url = "http://localhost:9090" + invoke_module._token_username = "alice" + + _install_recall_stub(monkeypatch) + + def _route(url, kwargs): + return _Resp(200, json_body={"status": "success"}) + + stub = _ScriptedClient([_route]) + monkeypatch.setattr(invoke_module.httpx, "AsyncClient", stub) + + response = client.post( + "/invokeai/recall", + json={"album_key": "any", "index": 0, "include_seed": True}, + ) + assert response.status_code == 200 + + # First and only attempt carried the cached token — no login call needed. + assert len(stub.calls) == 1 + assert stub.calls[0]["headers"].get("Authorization") == "Bearer cached-token" + + +def test_recall_403_with_cached_token_retries_anonymously_and_forgets_token( + client, clear_invokeai_config, clear_token_cache, monkeypatch +): + """If the backend has switched to single-user mode the cached token must be + discarded and the request retried without auth.""" + client.post( + "/invokeai/config", + json={ + "url": "http://localhost:9090", + "username": "alice", + "password": "secret", + }, + ) + + import time as _time + + from photomap.backend.routers import invoke as invoke_module + + # Pre-seed a cached token. + invoke_module._cached_token = "stale-token" + invoke_module._token_expires_at = _time.monotonic() + 3600 + invoke_module._token_base_url = "http://localhost:9090" + invoke_module._token_username = "alice" + + _install_recall_stub(monkeypatch) + + responses = [ + _Resp( + 403, + json_body={ + "detail": "Multiuser mode is disabled. Authentication is not required." + }, + text="forbidden", + ), + _Resp(200, json_body={"status": "success"}), + ] + + def _route(url, kwargs): + assert url.endswith("/recall/default"), url # no login expected + return responses.pop(0) + + stub = _ScriptedClient([_route, _route]) + monkeypatch.setattr(invoke_module.httpx, "AsyncClient", stub) + + response = client.post( + "/invokeai/recall", + json={"album_key": "any", "index": 0, "include_seed": True}, + ) + assert response.status_code == 200, response.text + + # First attempt carried the stale token, second attempt was anonymous. + assert len(stub.calls) == 2 + assert stub.calls[0]["headers"].get("Authorization") == "Bearer stale-token" + assert "Authorization" not in stub.calls[1]["headers"] + + # Cache must have been cleared. + assert invoke_module._cached_token is None + + +def test_recall_anonymous_403_is_not_retried( + client, clear_invokeai_config, clear_token_cache, monkeypatch +): + """A 403 on a call that was already anonymous just surfaces as-is — there + is nothing to retry without.""" + # No credentials configured, so the first request is anonymous. + client.post("/invokeai/config", json={"url": "http://localhost:9090"}) + + _install_recall_stub(monkeypatch) + + from photomap.backend.routers import invoke as invoke_module + + def _route(url, kwargs): + return _Resp(403, json_body={"detail": "forbidden"}, text="forbidden") + + stub = _ScriptedClient([_route]) + monkeypatch.setattr(invoke_module.httpx, "AsyncClient", stub) + + response = client.post( + "/invokeai/recall", + json={"album_key": "any", "index": 0, "include_seed": True}, + ) + assert response.status_code == 502 + assert len(stub.calls) == 1 # no retry, no login + + +def test_use_ref_image_403_with_token_retries_anonymously( + client, clear_invokeai_config, clear_token_cache, monkeypatch, tmp_path +): + """The upload + recall flow must handle the same auth transition as /recall.""" + client.post( + "/invokeai/config", + json={ + "url": "http://localhost:9090", + "username": "alice", + "password": "secret", + }, + ) + + import time as _time + + from photomap.backend.routers import invoke as invoke_module + + invoke_module._cached_token = "stale-token" + invoke_module._token_expires_at = _time.monotonic() + 3600 + invoke_module._token_base_url = "http://localhost:9090" + invoke_module._token_username = "alice" + + image_file = tmp_path / "pic.png" + image_file.write_bytes(b"\x89PNG\r\n\x1a\nfake") + monkeypatch.setattr( + invoke_module, "_load_image_path", lambda album_key, index: image_file + ) + + calls: list[dict] = [] + + class _ClientStub: + def __init__(self, *a, **kw): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def post(self, url, **kwargs): + entry = { + "url": url, + "headers": dict(kwargs.get("headers") or {}), + "kind": "upload" if "files" in kwargs else "recall", + "json": kwargs.get("json"), + } + if "files" in kwargs: + kwargs["files"]["file"][1].read() + calls.append(entry) + + if entry["kind"] == "upload": + # Upload returns 403 on the first attempt (with token), then + # succeeds anonymously. + if entry["headers"].get("Authorization"): + return _Resp( + 403, + json_body={"detail": "Multiuser mode is disabled."}, + text="forbidden", + ) + return _Resp(200, json_body={"image_name": "uploaded.png"}) + + # Recall also: 403 with token → anon success. + if entry["headers"].get("Authorization"): + return _Resp( + 403, + json_body={"detail": "Multiuser mode is disabled."}, + text="forbidden", + ) + return _Resp(200, json_body={"status": "success"}) + + monkeypatch.setattr(invoke_module.httpx, "AsyncClient", _ClientStub) + + response = client.post( + "/invokeai/use_ref_image", + json={"album_key": "any", "index": 0}, + ) + assert response.status_code == 200, response.text + + upload_calls = [c for c in calls if c["kind"] == "upload"] + recall_calls = [c for c in calls if c["kind"] == "recall"] + + # Upload retried anonymously. + assert len(upload_calls) == 2 + assert upload_calls[0]["headers"].get("Authorization") == "Bearer stale-token" + assert "Authorization" not in upload_calls[1]["headers"] + + # Recall also retried anonymously — and since the token was invalidated + # by the upload's 403, the very first recall attempt is already + # anonymous. Either way the final recall succeeds without a token. + assert len(recall_calls) >= 1 + assert "Authorization" not in recall_calls[-1]["headers"] + + # Cache cleared. + assert invoke_module._cached_token is None diff --git a/tests/frontend/invoke-recall.test.js b/tests/frontend/invoke-recall.test.js index a421e62..c01765d 100644 --- a/tests/frontend/invoke-recall.test.js +++ b/tests/frontend/invoke-recall.test.js @@ -18,7 +18,8 @@ jest.unstable_mockModule("../../photomap/frontend/static/javascript/utils.js", ( })); const { state } = await import("../../photomap/frontend/static/javascript/state.js"); -const { parseMetadataUrl, sendRecall } = await import("../../photomap/frontend/static/javascript/invoke-recall.js"); +const { parseMetadataUrl, sendRecall, sendUseRefImage } = + await import("../../photomap/frontend/static/javascript/invoke-recall.js"); describe("invoke-recall.js", () => { describe("parseMetadataUrl", () => { @@ -96,6 +97,34 @@ describe("invoke-recall.js", () => { }); }); + describe("sendUseRefImage", () => { + afterEach(() => { + delete global.fetch; + }); + + it("POSTs album_key / index to /invokeai/use_ref_image without include_seed", async () => { + const fetchMock = jest.fn(() => + Promise.resolve({ + ok: true, + json: () => Promise.resolve({ success: true, uploaded_image_name: "abc.png" }), + }) + ); + global.fetch = fetchMock; + + const result = await sendUseRefImage({ albumKey: "vacation", index: 3 }); + + expect(fetchMock).toHaveBeenCalledTimes(1); + const [url, opts] = fetchMock.mock.calls[0]; + expect(url).toBe("invokeai/use_ref_image"); + expect(opts.method).toBe("POST"); + expect(JSON.parse(opts.body)).toEqual({ + album_key: "vacation", + index: 3, + }); + expect(result.uploaded_image_name).toBe("abc.png"); + }); + }); + describe("button click handling", () => { beforeEach(() => { state.album = "fallback-album"; @@ -108,6 +137,9 @@ describe("invoke-recall.js", () => { + `; }); @@ -163,6 +195,26 @@ describe("invoke-recall.js", () => { expect(body.include_seed).toBe(false); }); + it("posts to /invokeai/use_ref_image when the use_ref button is clicked", async () => { + const fetchMock = jest.fn(() => + Promise.resolve({ + ok: true, + json: () => Promise.resolve({ success: true }), + }) + ); + global.fetch = fetchMock; + + document.querySelector('[data-recall-mode="use_ref"]').click(); + await flushPromises(); + + const [url, opts] = fetchMock.mock.calls[0]; + expect(url).toBe("invokeai/use_ref_image"); + const body = JSON.parse(opts.body); + expect(body.album_key).toBe("my-album"); + expect(body.index).toBe(5); + expect(body.include_seed).toBeUndefined(); + }); + it("shows a red X on failure", async () => { global.fetch = jest.fn(() => Promise.resolve({