Compare commits
71
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
91cbfd81d4 | ||
|
|
537149fd36 | ||
|
|
2f73744113 | ||
|
|
b6f6186a39 | ||
|
|
3e50f9414e | ||
|
|
c3ecffd4cd | ||
|
|
7d564eaa40 | ||
|
|
436a965fa0 | ||
|
|
d278c03f80 | ||
|
|
141c4b3fc9 | ||
|
|
f5e7483570 | ||
|
|
d45067db4d | ||
|
|
4b128ccfc8 | ||
|
|
20a433678c | ||
|
|
c474161fba | ||
|
|
e16131b0b6 | ||
|
|
2a5bb8a71d | ||
|
|
94e0c4e685 | ||
|
|
100ca8d8a2 | ||
|
|
c088a3833c | ||
|
|
1671b21334 | ||
|
|
3565974382 | ||
|
|
193aee3aac | ||
|
|
9370bf82b2 | ||
|
|
9d18ad1c75 | ||
|
|
6fcd119bc8 | ||
|
|
5ccee8f052 | ||
|
|
7d419c2101 | ||
|
|
5087403496 | ||
|
|
d6b888dbd0 | ||
|
|
225658f60f | ||
|
|
d561e35584 | ||
|
|
d3c5dca63f | ||
|
|
9062c238e8 | ||
|
|
b81974b52f | ||
|
|
1c7d66ce34 | ||
|
|
76af4f30e8 | ||
|
|
701d166b3f | ||
|
|
49c4f28e76 | ||
|
|
21d25458b9 | ||
|
|
0fd473923b | ||
|
|
4ac849c28e | ||
|
|
2ee53f67e4 | ||
|
|
5016d2c5d8 | ||
|
|
baa40cf16c | ||
|
|
7ba4956e8a | ||
|
|
8423eada67 | ||
|
|
4ee5adcef9 | ||
|
|
9e4d098af2 | ||
|
|
c86ded1bad | ||
|
|
4c77c82b90 | ||
|
|
a6c1130b54 | ||
|
|
a363751352 | ||
|
|
99919db60d | ||
|
|
fac9880aaf | ||
|
|
bb3711eca1 | ||
|
|
86e301345f | ||
|
|
3f48ca15e0 | ||
|
|
fdc2da4c57 | ||
|
|
9993da806e | ||
|
|
b5bbbfd724 | ||
|
|
73abeb1109 | ||
|
|
f75cb4b0cc | ||
|
|
c2f16445f2 | ||
|
|
2bf5421313 | ||
|
|
93688cac99 | ||
|
|
be313e8068 | ||
|
|
6b5597d8cb | ||
|
|
dd1f9e7138 | ||
|
|
86bb0a3cad | ||
|
|
96a633d3ce |
+4
-3
@@ -23,9 +23,10 @@ venv/
|
||||
.ruff_cache/
|
||||
tests/
|
||||
tests-unit/
|
||||
scripts/*
|
||||
# ...except the pack installer, which the image build invokes.
|
||||
!scripts/install_custom_nodes.sh
|
||||
# scripts/ ships WHOLE: the image build runs scripts/install_custom_nodes.sh and
|
||||
# the studio pod runs scripts/library_manifest.py as the build-manifest sidecar
|
||||
# (indexes every org's output/ into library.json). Excluding it silently no-ops
|
||||
# the sidecar's `[ -f ]` guard, so nothing gets indexed — keep the tree intact.
|
||||
notebooks/
|
||||
*.swp
|
||||
.env
|
||||
|
||||
@@ -97,6 +97,24 @@ shared Hanzo IAM (standard OIDC code flow). No embedding.
|
||||
visor_client}.py`): `/v1/workers/register` (heartbeat), `/v1/workers` (list),
|
||||
`/v1/worker/execute` (in-cluster push fast path). Worker↔coordinator trust via
|
||||
`X-Worker-Token` (`STUDIO_WORKER_TOKEN`, KMS-sourced).
|
||||
- **Per-GPU gpu-jobs queue** (`middleware/{gpu_dispatch,studio_home}.py`, `server.py`) —
|
||||
the ONE submit path onto a GPU. A model-less coordinator's `POST /prompt` enqueues to
|
||||
the cloud tasks engine (`POST {STUDIO_CLOUD_API_URL}/v1/tasks/namespaces/gpu-jobs/activities`,
|
||||
user bearer → org-scoped); a BYO worker (`hanzo gpu connect`) claims + runs it.
|
||||
**Targeting**: the Home "Run on <gpu>" chip sets `X-Target-GPU: <identity>`;
|
||||
`dispatch_if_worker` enqueues on that machine's own lane `taskQueue="gpu:<identity>"`
|
||||
(namespace stays `gpu-jobs`), else the shared `gpu-jobs` lane any worker drains —
|
||||
`X-Target-GPU` rides through `_queue_prompt` so fix/compose/video/rerun all honor it.
|
||||
**Authoritative visible queue**: `/v1/render-queue` + `/v1/nodes` read the tasks
|
||||
engine's activity list — REAL claiming `identity`, status
|
||||
SCHEDULED/STARTED/COMPLETED/FAILED → queued/running/done/failed, `lastHeartbeatTime`
|
||||
liveness, per-node depth — not a local shadow; `render_jobs.json` survives only as a
|
||||
~2s pre-claim placeholder. **No hidden runs**: in `--worker-mode` a direct `/prompt`
|
||||
POST is refused 403 unless it carries `X-Worker-Token` — the only accepted execution is
|
||||
a claimed job handed in over `POST /v1/worker/execute`
|
||||
(`worker_client.reject_untrusted_worker_submit`, one policy gate). Legacy push
|
||||
(`prompt_router`) is OFF unless `STUDIO_LEGACY_PUSH_ROUTER=1`. Tests:
|
||||
`middleware_test/{gpu_dispatch_test,worker_seam_test}.py`.
|
||||
- **Durable render queue** (`middleware/tasks_queue.py`, `docs/federation.md`):
|
||||
`SqlitePromptQueue` — crash-durable render queue in a single SQLite file
|
||||
(stdlib `sqlite3`, WAL, **zero external processes**). `STUDIO_QUEUE_DB=<path>`
|
||||
@@ -111,6 +129,25 @@ shared Hanzo IAM (standard OIDC code flow). No embedding.
|
||||
compete safely. Tested: `middleware_test/tasks_queue_test.py` (incl.
|
||||
crash-recovery across instances). Future (federation.md §6): a remote queue
|
||||
backend served by Hanzo Tasks + a Go/unified-binary (HIP-0106) migration.
|
||||
- **Stacks** (`middleware/stacks_store.py`, `middleware/studio_home.py`): Procreate-
|
||||
style gallery folders — a per-org SQLite store (`orgs/<org>/stacks.db`, WAL, every
|
||||
mutation under `BEGIN IMMEDIATE`, the same `tasks_queue.py` pattern) that is ONLY an
|
||||
organizational layer; assets stay the source of truth (library.json + files). Schema
|
||||
= the spec's `stacks` + `stack_assets` join; `UNIQUE(asset_id)` ⇒ one Stack per
|
||||
asset. Tenancy is structural (one DB file per org → cross-org 404s by construction).
|
||||
REST (`/v1` house style): `GET/POST /v1/stacks`, `GET/PATCH/DELETE /v1/stacks/:id`
|
||||
(`?mode=only|images`), `POST|DELETE /v1/stacks/:id/assets`, `PATCH /v1/stacks/:id/order`,
|
||||
`POST /v1/stacks/:id/unstack`, `GET/PATCH /v1/stacks/settings`; registered via
|
||||
injected resolvers (org/owner/workspace + the library soft-deleter) so the module
|
||||
stays decoupled. Remove ≠ delete (returns to the gallery); Delete-Stack defaults to
|
||||
preserving; "…and Images" maps to the library soft-delete (`status=deleted`, files
|
||||
kept). Generation integration: `fix/compose/rerun/template` accept a `stack` id and
|
||||
the ingest point (`_index_upload` → `stacks_store.absorb`) assigns membership when
|
||||
outputs land (ALL outputs of a multi-output job); a per-org `auto_stack_batches`
|
||||
setting auto-creates `"<prompt> — <date>"` Stacks. UI is the vanilla-JS home
|
||||
(`studio_home.html`): layered cards, drag/multi-select/mobile-tap parity (no DnD ever
|
||||
required), a11y labels + live region. Tested: `studio_home_test/test_stacks.py`
|
||||
(routes, tenancy, one-per-asset, order persistence, delete/unstack, ingest→absorb).
|
||||
- **Engine selector** (`middleware/engine_selector.py`): `/v1/engines` lists
|
||||
execution targets for the org (`local` + registered compute_config workers) and
|
||||
`PUT /v1/engines/default` sets a per-org default (stored on
|
||||
|
||||
@@ -23,7 +23,8 @@ test:
|
||||
tests-unit/folder_paths_test \
|
||||
tests-unit/middleware_test \
|
||||
tests-unit/app_test \
|
||||
tests-unit/prompt_server_test
|
||||
tests-unit/prompt_server_test \
|
||||
tests-unit/studio_home_test
|
||||
|
||||
# Deploy: roll the freshly-built image onto the studio Service CR.
|
||||
deploy:
|
||||
|
||||
@@ -456,12 +456,24 @@ def start_studio(asyncio_loop=None):
|
||||
|
||||
# In worker mode, add worker execution routes and register with coordinator
|
||||
if args.worker_mode:
|
||||
from middleware.worker_client import add_worker_routes
|
||||
from middleware.worker_client import add_worker_routes, WORKER_TOKEN
|
||||
# FAIL CLOSED: a worker box executes jobs the coordinator hands it over a
|
||||
# secret-authenticated seam. Without STUDIO_WORKER_TOKEN that seam (and the
|
||||
# worker-mode /prompt gate) would accept un-tokened callers — silently
|
||||
# reopening the hidden-run hole. Refuse to start rather than run wide open.
|
||||
if not WORKER_TOKEN:
|
||||
sys.exit("worker-mode requires STUDIO_WORKER_TOKEN (coordinator shared secret) — refusing to start")
|
||||
add_worker_routes(prompt_server.routes, prompt_server)
|
||||
logging.info("Worker mode enabled — worker_id=%s coordinator=%s",
|
||||
args.worker_id, args.coordinator_url)
|
||||
|
||||
threading.Thread(target=prompt_worker, daemon=True, args=(prompt_server.prompt_queue, prompt_server,)).start()
|
||||
_worker_thread = threading.Thread(target=prompt_worker, daemon=True, args=(prompt_server.prompt_queue, prompt_server,))
|
||||
_worker_thread.start()
|
||||
# Exposed so /ready can gate on the render worker actually being alive: if this
|
||||
# daemon thread dies (an unhandled exception in the executor), the queue still
|
||||
# accepts /prompt but nothing runs — readiness must fail so the pod leaves rotation
|
||||
# and a roll never promotes a wedged pod.
|
||||
prompt_server.prompt_worker_thread = _worker_thread
|
||||
|
||||
if args.quick_test_for_ci:
|
||||
exit(0)
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Colour fidelity for Fix outputs.
|
||||
|
||||
A Qwen-Image-Edit fix runs at denoise 1.0 — it re-paints EVERY pixel, so the
|
||||
render's global white-balance / exposure / saturation drift off the source, and
|
||||
that drift COMPOUNDS across fix-of-a-fix. This restores the render's global colour
|
||||
statistics to its source (Reinhard mean/std transfer, per RGB channel): the local
|
||||
edit stays, the global cast goes. Matching an already-correct image to its source
|
||||
is the identity, so a render that did not drift is left untouched.
|
||||
|
||||
Applied POD-SIDE when a BYO-GPU worker's finished render is ingested
|
||||
(server.upload_output). It needs NO node on the GPU worker and touches no render
|
||||
graph, so it can never make a render fail. Fail-safe throughout: any error, missing
|
||||
source, or size mismatch leaves the raw bytes exactly as uploaded.
|
||||
"""
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
# Fraction of the correction to apply. <1.0 leaves headroom for an INTENTIONAL
|
||||
# global colour change in the instruction: the drift being corrected is a global
|
||||
# cast, while a normal edit is local and barely moves the whole-frame statistics,
|
||||
# so a moderate-high default corrects drift while an intended recolour mostly
|
||||
# survives. Tunable without a rebuild via STUDIO_COLORMATCH_STRENGTH.
|
||||
STRENGTH = float(os.environ.get("STUDIO_COLORMATCH_STRENGTH", "0.85"))
|
||||
|
||||
|
||||
def color_match(target: Image.Image, reference: Image.Image, strength: float = STRENGTH) -> Image.Image:
|
||||
"""Pull ``target``'s global colour balance onto ``reference``'s — a per-channel
|
||||
Reinhard mean/std transfer in RGB. target-to-itself is the identity."""
|
||||
t = np.asarray(target.convert("RGB"), np.float32)
|
||||
ref = reference.convert("RGB")
|
||||
if ref.size != target.size:
|
||||
ref = ref.resize(target.size, Image.LANCZOS)
|
||||
r = np.asarray(ref, np.float32)
|
||||
out = t.copy()
|
||||
for c in range(3):
|
||||
tm, ts = t[..., c].mean(), t[..., c].std() + 1e-5
|
||||
rm, rs = r[..., c].mean(), r[..., c].std() + 1e-5
|
||||
out[..., c] = (t[..., c] - tm) * (rs / ts) + rm
|
||||
out = t + (out - t) * float(strength)
|
||||
return Image.fromarray(np.clip(out, 0, 255).astype(np.uint8))
|
||||
|
||||
|
||||
def match_in_place(out_path: str, src_path: str, strength: float = STRENGTH) -> bool:
|
||||
"""Colour-match the fix render at ``out_path`` to its source, atomically in
|
||||
place (tmp + replace). Returns True if rewritten, False (no-op) on any error or
|
||||
missing input — never raises."""
|
||||
try:
|
||||
if strength <= 0 or not (os.path.isfile(out_path) and os.path.isfile(src_path)):
|
||||
return False
|
||||
with Image.open(out_path) as o, Image.open(src_path) as s:
|
||||
matched = color_match(o, s, strength)
|
||||
# Tmp keeps a real image extension so PIL can infer the format; pass it
|
||||
# explicitly too (the ".cm.tmp" suffix alone is unknown to PIL).
|
||||
ext = os.path.splitext(out_path)[1].lower()
|
||||
fmt = {".png": "PNG", ".jpg": "JPEG", ".jpeg": "JPEG", ".webp": "WEBP"}.get(ext, "PNG")
|
||||
tmp = f"{out_path}.cm.tmp{ext or '.png'}"
|
||||
matched.save(tmp, format=fmt)
|
||||
os.replace(tmp, out_path)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
+170
-17
@@ -8,13 +8,22 @@ shared keys. Falls back to the in-pod queue when there is no live worker.
|
||||
"""
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import urllib.request
|
||||
|
||||
import folder_paths
|
||||
|
||||
log = logging.getLogger("studio.gpu_dispatch")
|
||||
|
||||
CLOUD = os.environ.get("STUDIO_CLOUD_API_URL", "https://api.hanzo.ai").rstrip("/")
|
||||
JOBS_NS = os.environ.get("STUDIO_GPU_JOBS_NS", "gpu-jobs")
|
||||
|
||||
# The ComfyUI output-node class types — the ONE Python list, defined in this base
|
||||
# module (no studio_home dependency) so studio_home can import it without a cycle.
|
||||
# The client mirror lives in studio_home.html (a separate JS runtime).
|
||||
SAVE_NODES = ("SaveImage", "SaveVideo", "SaveWEBM")
|
||||
# Public base of THIS studio deployment — where the BYO-GPU worker returns finished
|
||||
# outputs (POST /upload/output → orgs/{org}/output → gallery). The worker uploads
|
||||
# here with the user's IAM token; no S3/rclone credentials ever touch the box.
|
||||
@@ -92,11 +101,22 @@ def _node_label(m: dict) -> str:
|
||||
|
||||
def _fleet_machines(tok: str) -> list:
|
||||
"""The caller-org's machines from the cloud fleet. The token scopes the fleet to
|
||||
the caller's org, so this is inherently tenant-isolated."""
|
||||
req = urllib.request.Request(f"{CLOUD}/v1/machines", headers={"Authorization": tok})
|
||||
d = json.loads(urllib.request.urlopen(req, timeout=10).read())
|
||||
ms = d.get("machines", [])
|
||||
return ms if isinstance(ms, list) else []
|
||||
the caller's org, so this is inherently tenant-isolated. Retried a few times so a
|
||||
transient edge fault (502, TLS reset, timeout) does NOT read as "no GPU" and
|
||||
bounce the render onto the model-less pod — the cause of intermittent
|
||||
"Prompt outputs failed validation"."""
|
||||
last = None
|
||||
for attempt in range(3):
|
||||
try:
|
||||
req = urllib.request.Request(f"{CLOUD}/v1/machines", headers={"Authorization": tok})
|
||||
d = json.loads(urllib.request.urlopen(req, timeout=10).read())
|
||||
ms = d.get("machines", [])
|
||||
return ms if isinstance(ms, list) else []
|
||||
except Exception as e: # any transient fault (HTTP/TLS/timeout/non-JSON) is retryable
|
||||
last = e
|
||||
if attempt < 2:
|
||||
time.sleep(0.4 * (attempt + 1))
|
||||
raise last
|
||||
|
||||
|
||||
def _online_gpu_nodes(tok: str) -> list:
|
||||
@@ -108,23 +128,135 @@ def _has_online_gpu(tok: str) -> bool:
|
||||
return bool(_online_gpu_nodes(tok))
|
||||
|
||||
|
||||
# Server credentials that read the WHOLE fleet, not just the caller-org's slice like
|
||||
# a user token. KMS-sourced (the SAME envs worker/commerce already use); consulted
|
||||
# ONLY by the advisory capacity gate below — never by a render decision.
|
||||
_SERVICE_TOKEN_ENVS = ("STUDIO_WORKER_TOKEN", "STUDIO_COMMERCE_TOKEN")
|
||||
|
||||
|
||||
def _service_token() -> str:
|
||||
"""A service bearer for a fleet read, or "" if none is configured."""
|
||||
for env in _SERVICE_TOKEN_ENVS:
|
||||
v = os.environ.get(env, "").strip()
|
||||
if v:
|
||||
return v if v.lower().startswith("bearer ") else f"Bearer {v}"
|
||||
return ""
|
||||
|
||||
|
||||
def has_online_gpu_advisory(user_tok: str) -> bool:
|
||||
"""Fleet availability for the UI capacity gate ONLY (the ⚠ warning + the
|
||||
/v1/gpu-status badge) — NEVER a render decision. Reads with a SERVICE identity
|
||||
first so every org sees the TRUE worker availability: a BYO box is often
|
||||
registered under a sibling org and is invisible to an org-scoped user-token read,
|
||||
which made a karma session believe no GPU existed. Falls back to the user token
|
||||
when no service credential is set (local dev). Job tenancy is untouched — dispatch
|
||||
still enqueues under the user token."""
|
||||
for tok in (_service_token(), user_tok):
|
||||
if not tok:
|
||||
continue
|
||||
try:
|
||||
if _has_online_gpu(tok):
|
||||
return True
|
||||
except Exception:
|
||||
continue
|
||||
return False
|
||||
|
||||
|
||||
# Model-loader class → (the widget naming the file, the folder_paths list it lives in).
|
||||
# Covers the loaders the studio graphs use: Qwen fix/compose (UNET+CLIP+VAE) and the
|
||||
# SD-family templates (checkpoint). Enough to answer "can THIS box load THIS graph".
|
||||
_LOADERS = {
|
||||
"UNETLoader": ("unet_name", "diffusion_models"),
|
||||
"CheckpointLoaderSimple": ("ckpt_name", "checkpoints"),
|
||||
"CheckpointLoader": ("ckpt_name", "checkpoints"),
|
||||
"ImageOnlyCheckpointLoader": ("ckpt_name", "checkpoints"),
|
||||
"CLIPLoader": ("clip_name", "text_encoders"),
|
||||
"VAELoader": ("vae_name", "vae"),
|
||||
}
|
||||
|
||||
|
||||
def has_models_for(prompt: dict) -> bool:
|
||||
"""True only if THIS instance has EVERY model file the graph actually references
|
||||
(its UNET / checkpoint / CLIP / VAE loader inputs). The decision is per-GRAPH, not
|
||||
"any model present": a cloud coordinator may carry a stray checkpoint (e.g. an SD1.5
|
||||
stopgap) yet lack the Qwen edit models a fix/compose graph needs — validating THAT
|
||||
locally is the cryptic "not in []" / "Prompt outputs failed validation". So a box
|
||||
renders locally only a graph it can actually load; every other graph goes to the
|
||||
fleet (or a retryable 503). Errs toward the fleet (returns False) on any doubt. A
|
||||
worker node (worker_mode) never reaches this path."""
|
||||
try:
|
||||
for node in (prompt or {}).values():
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
spec = _LOADERS.get(node.get("class_type", ""))
|
||||
if not spec:
|
||||
continue
|
||||
name = (node.get("inputs") or {}).get(spec[0])
|
||||
if not isinstance(name, str) or not name:
|
||||
continue
|
||||
try:
|
||||
have = folder_paths.get_filename_list(spec[1])
|
||||
except Exception:
|
||||
have = []
|
||||
if name not in have:
|
||||
return False # a model this graph needs is absent → cannot render here
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def dispatch_if_worker(request, org_id: str, prompt_id: str, prompt: dict) -> bool:
|
||||
"""If the caller's org has an online GPU node in the fleet, enqueue a
|
||||
studio.render job to gpu-jobs and return True. Return False to run in-pod.
|
||||
Never raises — any failure falls back to local."""
|
||||
"""Enqueue this render to the shared GPU lane (gpu-jobs) and return True; return
|
||||
False only to run it in-pod.
|
||||
|
||||
A model-less cloud coordinator NEVER renders locally — it ALWAYS enqueues and
|
||||
lets the job WAIT for a worker, even when the org-scoped fleet read shows none.
|
||||
BYO workers can be registered under a sibling org while draining the SHARED
|
||||
gpu-jobs queue, so an org-blind read (e.g. an org whose GPUs live under a peer
|
||||
org) must not bounce the render onto the GPU-less pod — that self-POST is what
|
||||
died with an instant "Prompt outputs failed validation". A queued job that waits
|
||||
for a worker beats an instant fake failure. A box WITH its own models still
|
||||
renders locally when it sees no online worker (local dev / a worker node).
|
||||
|
||||
Never raises — a genuine enqueue failure falls back, and server.py then returns a
|
||||
retryable 503 on a model-less pod rather than a doomed local validation."""
|
||||
try:
|
||||
tok = _bearer(request)
|
||||
online = _online_gpu_nodes(tok) if tok else []
|
||||
if not online:
|
||||
if not tok:
|
||||
return False
|
||||
# Which node runs it: the single online GPU is definitive; with several the
|
||||
# claimer is decided at pick-up time, so record the ambiguous "gpu".
|
||||
node = _node_label(online[0]) if len(online) == 1 else "gpu"
|
||||
# Per-GPU targeting: the UI's "Run on <gpu>" affordance sets X-Target-GPU to a
|
||||
# machine identity. A targeted render is enqueued on THAT machine's own lane
|
||||
# ("gpu:<identity>") which only that worker claims; an untargeted render rides
|
||||
# the shared "gpu-jobs" lane any worker drains. The NAMESPACE is "gpu-jobs"
|
||||
# either way — targeting is a taskQueue, not a second queue. One enqueue path.
|
||||
target = (request.headers.get("X-Target-GPU") or "").strip()
|
||||
online = _online_gpu_nodes(tok)
|
||||
# Honor a target ONLY if it is one of the caller-org's OWN online GPUs. This
|
||||
# blocks a forged lane AND sanitizes the value before it becomes a taskQueue
|
||||
# (and later a queue-row label): an unknown/offline target falls back to the
|
||||
# shared lane rather than enqueuing onto a "gpu:<bogus>" lane no worker will
|
||||
# ever claim (a 30-min schedule-to-start hang) — and a "<img onerror>" target
|
||||
# never reaches the queue. Cross-org is already impossible (per-org shard); this
|
||||
# is defense in depth.
|
||||
if target and target not in {_node_label(m) for m in online}:
|
||||
target = ""
|
||||
task_queue = f"gpu:{target}" if target else JOBS_NS
|
||||
# No visible worker AND this box can load THIS graph → render locally. No
|
||||
# visible worker on a box that lacks the graph's models → still enqueue and
|
||||
# wait (never a doomed local validation), even with a stray unrelated model.
|
||||
# A validated target ALWAYS enqueues (the user asked for that online GPU).
|
||||
if not target and not online and has_models_for(prompt):
|
||||
return False
|
||||
# Which node runs it: the explicit target; else a single visible GPU is
|
||||
# definitive; with several — or an org-blind read that saw none — the claimer
|
||||
# is decided at pick-up: "gpu".
|
||||
node = target or (_node_label(online[0]) if len(online) == 1 else "gpu")
|
||||
staged = _collect_input_images(prompt, org_id)
|
||||
body = json.dumps({
|
||||
"activityId": prompt_id,
|
||||
"runId": prompt_id,
|
||||
"activityType": {"name": "studio.render"},
|
||||
"taskQueue": JOBS_NS,
|
||||
"taskQueue": task_queue,
|
||||
# 1200s: a heavy Qwen edit on a cold GB10 (model reload + full 30-step
|
||||
# sample) can exceed 600s; the worker heartbeats each step, so this is
|
||||
# the ceiling for a genuinely long single render, not idle slack.
|
||||
@@ -133,6 +265,17 @@ def dispatch_if_worker(request, org_id: str, prompt_id: str, prompt: dict) -> bo
|
||||
# render. Unset, the tasks default (~1h) reaped real renders mid-run and
|
||||
# the queue re-ran them for hours.
|
||||
"startToCloseTimeout": "14400s",
|
||||
# Stale jobs EXPIRE instead of waiting forever for a worker. If no GPU
|
||||
# picks this up within 30 min it is dropped — so a worker outage can't
|
||||
# leave a backlog that silently replays (re-billing old renders) the moment
|
||||
# a worker reconnects. A live user re-runs; nothing lingers unbounded.
|
||||
"scheduleToStartTimeout": "1800s",
|
||||
# ONE attempt, no auto-retry. A GPU render is expensive and often
|
||||
# deterministic — re-running a failed/crashed render (e.g. spark OOM mid-
|
||||
# sample) just re-burns money on the SAME output ("the same reject over and
|
||||
# over"). A failure surfaces to the user, who re-runs deliberately. This is
|
||||
# the cap that the startToCloseTimeout alone did NOT provide.
|
||||
"retryPolicy": {"maximumAttempts": 1},
|
||||
"input": {
|
||||
"prompt": prompt,
|
||||
"org": org_id,
|
||||
@@ -140,7 +283,7 @@ def dispatch_if_worker(request, org_id: str, prompt_id: str, prompt: dict) -> bo
|
||||
# Uploaded inputs live in orgs/{org}/input on THIS pod; the worker
|
||||
# renders on its own disk and cannot read them — ship them along so
|
||||
# LoadImage resolves there.
|
||||
"inputs": _collect_input_images(prompt, org_id),
|
||||
"inputs": staged,
|
||||
},
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
@@ -148,8 +291,18 @@ def dispatch_if_worker(request, org_id: str, prompt_id: str, prompt: dict) -> bo
|
||||
data=body,
|
||||
headers={"Authorization": tok, "Content-Type": "application/json"},
|
||||
)
|
||||
urllib.request.urlopen(req, timeout=15).read()
|
||||
# 90s, not 15s: a fix/compose with a large (e.g. 4k) reference inlines tens of
|
||||
# MB, and 15s wasn't enough to PUT the enqueue body — it timed out, dispatch
|
||||
# fell back, and the model-less pod returned "GPU worker unavailable" on a
|
||||
# perfectly good render. The enqueue itself is quick; this only covers upload.
|
||||
urllib.request.urlopen(req, timeout=90).read()
|
||||
_track_dispatched(org_id, prompt_id, prompt, node=node)
|
||||
# Correlated by pid (the job id) with studio.home's build/queue lines; names are
|
||||
# content-addressed checksums, so this pins the exact source bytes shipped — no
|
||||
# signed URLs, tokens or image data are logged.
|
||||
log.info("studio.trace gpu.publish org=%s pid=%s node=%s inputs=%s refs=%s",
|
||||
org_id, prompt_id, node, len(staged),
|
||||
",".join(sorted(s["name"] for s in staged)) or None)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
@@ -170,7 +323,7 @@ def _track_dispatched(org_id: str, prompt_id: str, prompt: dict, node: str = "gp
|
||||
if not isinstance(n, dict):
|
||||
continue
|
||||
ct = n.get("class_type", "")
|
||||
if ct == "SaveImage":
|
||||
if ct in SAVE_NODES:
|
||||
full_prefix = n.get("inputs", {}).get("filename_prefix") or "render"
|
||||
elif "LoadImage" in ct:
|
||||
im = n.get("inputs", {}).get("image")
|
||||
|
||||
@@ -9,10 +9,13 @@ signing keys are cached. Signature (RS256 by default), ``exp``, ``iss`` and
|
||||
tenancy / rate-limit / billing paths.
|
||||
|
||||
Browser sessions use the standard OIDC Authorization Code flow: an
|
||||
unauthenticated HTML request is redirected to the IAM ``authorize`` endpoint;
|
||||
IAM bounces back to ``/callback`` which exchanges the code for a token and sets
|
||||
an httpOnly session cookie. No frontend changes are required — the prebuilt SPA
|
||||
just rides the cookie.
|
||||
unauthenticated browser navigation is redirected (302) to ``/login`` — the one
|
||||
server-owned entry point — which starts the flow at the IAM ``authorize``
|
||||
endpoint; IAM bounces back to ``/callback`` which exchanges the code for a token
|
||||
and sets an httpOnly session cookie. A programmatic API/XHR/JSON request instead
|
||||
gets a 401 (it handles its own re-auth), so the SPA's ``fetch`` calls never
|
||||
follow a redirect and parse HTML as JSON. No frontend changes are required — the
|
||||
prebuilt SPA just rides the cookie.
|
||||
|
||||
Localhost connections bypass auth (single-user dev). Static assets, health
|
||||
checks and the worker coordinator endpoints are exempt.
|
||||
@@ -94,6 +97,48 @@ def _is_public_path(path: str) -> bool:
|
||||
))
|
||||
|
||||
|
||||
def _is_api_request(request: web.Request) -> bool:
|
||||
"""A programmatic API/XHR/JSON caller (answered with 401) as opposed to a
|
||||
top-level browser navigation (answered with a 302 into the login flow).
|
||||
|
||||
A human who types a URL or follows a link must never be shown a raw 401, so
|
||||
an *ambiguous* non-browser client (plain ``curl``, ``Accept: */*``, no
|
||||
fetch-metadata) counts as a navigation and is redirected. Only a request that
|
||||
self-identifies as fetch/XHR/JSON is refused — otherwise the SPA's own
|
||||
``fetch`` calls would follow the redirect to the IdP and parse HTML as JSON.
|
||||
"""
|
||||
# Fetch-metadata (every evergreen browser stamps it): a user navigation is
|
||||
# ``Sec-Fetch-Mode: navigate``; ``fetch()``/XHR carry cors|same-origin|no-cors.
|
||||
mode = request.headers.get("Sec-Fetch-Mode", "").lower()
|
||||
if mode:
|
||||
return mode != "navigate"
|
||||
# Legacy AJAX marker (jQuery / axios / htmx).
|
||||
if request.headers.get("X-Requested-With", "").lower() == "xmlhttprequest":
|
||||
return True
|
||||
# Content negotiation: an explicit JSON preference with no HTML acceptance.
|
||||
accept = request.headers.get("Accept", "")
|
||||
return "application/json" in accept and "text/html" not in accept
|
||||
|
||||
|
||||
def _safe_next(raw: str | None) -> str:
|
||||
"""A same-origin return path, or ``/``. Rejects absolute URLs, protocol- and
|
||||
backslash-relative forms (open-redirect vectors) and the login route itself
|
||||
(which would loop)."""
|
||||
if not raw or not raw.startswith("/") or raw[:2] in ("//", "/\\"):
|
||||
return "/"
|
||||
if raw == "/login" or raw.startswith(("/login?", "/login/")):
|
||||
return "/"
|
||||
return raw
|
||||
|
||||
|
||||
def _login_redirect(request: web.Request) -> web.Response:
|
||||
"""Send an unauthenticated browser to the ONE server-owned login entry point
|
||||
(``/login``), preserving where it was headed via ``?next`` so the OIDC round
|
||||
trip returns it there. The authorize-URL construction lives in ``/login``
|
||||
(``handle_login``) alone — this gate only decides *who* is sent there."""
|
||||
return web.HTTPFound("/login?" + urlencode({"next": str(request.rel_url)}))
|
||||
|
||||
|
||||
def _cache_key(token: str) -> str:
|
||||
return hashlib.sha256(token.encode()).hexdigest()[:32]
|
||||
|
||||
@@ -339,7 +384,9 @@ def create_iam_auth_middleware(iam_url: str, localhost_bypass: bool = True):
|
||||
Authorization Code flow (PKCE + signed state + ``/callback``) is built
|
||||
server-side — there is exactly one way to start a login, and the client
|
||||
never hand-rolls an authorize URL. Already-authenticated callers skip
|
||||
straight to the app; everyone else is bounced to IAM and returns to ``/``.
|
||||
straight to the app; everyone else is bounced to IAM and returns to the
|
||||
sanitized ``?next`` target (default ``/``) — this is where the auth gate
|
||||
sends an unauthenticated browser navigation, carrying its original path.
|
||||
"""
|
||||
token = None
|
||||
auth_header = request.headers.get("Authorization", "")
|
||||
@@ -353,7 +400,7 @@ def create_iam_auth_middleware(iam_url: str, localhost_bypass: bool = True):
|
||||
return web.HTTPFound("/")
|
||||
except Exception:
|
||||
pass # IAM/JWKS hiccup — fall through and re-authenticate cleanly.
|
||||
return _authorize_redirect(request, target="/")
|
||||
return _authorize_redirect(request, target=_safe_next(request.query.get("next")))
|
||||
|
||||
async def handle_callback(request: web.Request) -> web.Response:
|
||||
"""Exchange the authorization code for a token and set the session cookie."""
|
||||
@@ -428,12 +475,16 @@ def create_iam_auth_middleware(iam_url: str, localhost_bypass: bool = True):
|
||||
if not token and path == "/ws":
|
||||
token = request.rel_url.query.get("token")
|
||||
|
||||
wants_html = "text/html" in request.headers.get("Accept", "")
|
||||
# A browser navigation is redirected into the login flow (302 → /login,
|
||||
# carrying ?next); a programmatic API/XHR/JSON caller gets a 401 it can
|
||||
# handle itself. One predicate decides, one way, for both no-token and
|
||||
# bad-token cases below.
|
||||
api = _is_api_request(request)
|
||||
|
||||
if not token:
|
||||
if wants_html:
|
||||
return _authorize_redirect(request)
|
||||
return web.json_response({"error": "Authentication required", "iam_url": iam_url}, status=401)
|
||||
if api:
|
||||
return web.json_response({"error": "Authentication required", "iam_url": iam_url}, status=401)
|
||||
return _login_redirect(request)
|
||||
|
||||
try:
|
||||
user = await _validate(token)
|
||||
@@ -441,9 +492,9 @@ def create_iam_auth_middleware(iam_url: str, localhost_bypass: bool = True):
|
||||
return web.json_response({"error": "Authentication service unavailable"}, status=503)
|
||||
|
||||
if user is None:
|
||||
if wants_html:
|
||||
return _authorize_redirect(request)
|
||||
return web.json_response({"error": "Invalid or expired token"}, status=401)
|
||||
if api:
|
||||
return web.json_response({"error": "Invalid or expired token"}, status=401)
|
||||
return _login_redirect(request)
|
||||
|
||||
request["iam_user"] = _with_active_org(user, request.cookies.get(_ACTIVE_ORG_COOKIE))
|
||||
# The verified raw access token. Studio records a completed render in the
|
||||
|
||||
@@ -0,0 +1,286 @@
|
||||
"""Render quality gate — the ONE guard that keeps corrupted renders out of the library.
|
||||
|
||||
Every render lands through ``POST /v1/library/upload`` → ``studio_home._index_upload``.
|
||||
The moment a NEW asset row is created there, this fires an out-of-band judgment
|
||||
(never blocks the upload response): downscale the just-stored image, ask a cloud
|
||||
vision model to answer PASS/FAIL, and on a clear FAIL move the asset's library
|
||||
status ``draft → flagged`` so it drops out of every default grid — the same
|
||||
visibility semantics as ``deleted``, but recoverable from the UI's Flagged view.
|
||||
|
||||
It FAILS OPEN by construction. A PASS, an inconclusive reply, a 402/429, a timeout,
|
||||
an unreachable judge, an unreadable image — anything that is not an unambiguous
|
||||
FAIL — leaves the render a normal ``draft``. The gate can only ever hide clearly
|
||||
broken work; it can never hide good work by accident, and it never blocks a render
|
||||
from landing.
|
||||
|
||||
Env (all optional):
|
||||
STUDIO_GATE "off"/"0"/"false"/"no" disables the gate entirely. Default on.
|
||||
STUDIO_GATE_MODEL vision model id. Default "zen-vl" (the cloud's VL model).
|
||||
STUDIO_GATE_URL full chat-completions URL override, e.g. a local engine at
|
||||
http://127.0.0.1:1234/v1/chat/completions. Default = the
|
||||
cloud gateway (gpu_dispatch.CLOUD) + /v1/chat/completions.
|
||||
|
||||
One auth path: the judgment carries the SAME bearer the upload carried, so the
|
||||
cloud derives org + billing from it — no shared keys, tenant-isolated like dispatch.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
|
||||
from middleware import worklog
|
||||
from middleware.gpu_dispatch import CLOUD, _bearer
|
||||
|
||||
log = logging.getLogger("studio.gate")
|
||||
|
||||
# The strict binary judge. The model is told to answer EXACTLY PASS or FAIL; the
|
||||
# failure classes are the observed corruption modes of the render pipeline.
|
||||
PROMPT = (
|
||||
"You are a render QA gate for fashion imagery. Look at the image. Answer exactly "
|
||||
"PASS or FAIL. FAIL if: skin texture is corrupted (mottled, blotchy, speckled, "
|
||||
"noise-stippled), the image has ghosting/double-exposure, a frame-within-frame/"
|
||||
"poster border, or glyph/pixel-garbage regions. Otherwise PASS."
|
||||
)
|
||||
|
||||
# The motion judge — the SAME strict binary contract as PROMPT, but the failure classes
|
||||
# are the observed corruption modes of the video pipeline (it is shown start/mid/end frames).
|
||||
MOTION = (
|
||||
"You are a render QA gate for generated video, shown frames sampled across its "
|
||||
"duration. Answer exactly PASS or FAIL. FAIL if: frame-to-frame flicker, "
|
||||
"identity/clothing morphing between frames, frozen/static motion where movement is "
|
||||
"expected, temporal ghosting or trails, or glyph/pixel-garbage regions. Otherwise PASS."
|
||||
)
|
||||
|
||||
_OFF = ("off", "0", "false", "no")
|
||||
_FAIL_RE = re.compile(r"\bfail\b", re.I)
|
||||
_PASS_RE = re.compile(r"\bpass\b", re.I)
|
||||
|
||||
# Keep a reference to in-flight tasks so a fire-and-forget judgment is not
|
||||
# garbage-collected before it runs (per asyncio.ensure_future docs).
|
||||
_TASKS: set = set()
|
||||
|
||||
|
||||
def enabled() -> bool:
|
||||
return os.environ.get("STUDIO_GATE", "").strip().lower() not in _OFF
|
||||
|
||||
|
||||
def _model() -> str:
|
||||
return os.environ.get("STUDIO_GATE_MODEL", "").strip() or "zen-vl"
|
||||
|
||||
|
||||
def _url() -> str:
|
||||
return os.environ.get("STUDIO_GATE_URL", "").strip() or (CLOUD + "/v1/chat/completions")
|
||||
|
||||
|
||||
def parse_verdict(reply: str) -> str:
|
||||
""""pass" | "fail" | "inconclusive" from the judge's reply. Case-insensitive
|
||||
PASS/FAIL as a whole token (so "FAIL — mottled skin" reads as fail but "failure"
|
||||
does not spuriously flag). A reply carrying BOTH tokens, or NEITHER, is
|
||||
inconclusive — and inconclusive fails open."""
|
||||
fails = [m.start() for m in _FAIL_RE.finditer(reply or "")]
|
||||
passes = [m.start() for m in _PASS_RE.finditer(reply or "")]
|
||||
# The LAST token wins: a reasoning judge weighs both words before concluding
|
||||
# ("is this PASS or FAIL... verdict: FAIL") — co-occurrence is deliberation,
|
||||
# position is the verdict.
|
||||
if fails and (not passes or fails[-1] > passes[-1]):
|
||||
return "fail"
|
||||
if passes and (not fails or passes[-1] > fails[-1]):
|
||||
return "pass"
|
||||
return "inconclusive"
|
||||
|
||||
|
||||
def _encode(im) -> str | None:
|
||||
"""A PIL image as a ~512px JPEG, base64 — the ONE frame-encode contract, small
|
||||
enough for a fast vision call. Returns None on any failure."""
|
||||
try:
|
||||
im = im.convert("RGB")
|
||||
im.thumbnail((512, 512))
|
||||
buf = io.BytesIO()
|
||||
im.save(buf, "JPEG", quality=80)
|
||||
return base64.b64encode(buf.getvalue()).decode()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _downscale(path: Path) -> str | None:
|
||||
"""The stored image as a ~512px JPEG, base64. None if it can't be read (→ fail open)."""
|
||||
try:
|
||||
from PIL import Image
|
||||
return _encode(Image.open(path))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _video_frame(path: str, t: float) -> str | None:
|
||||
"""One ~512px JPEG (base64) decoded at ``t`` seconds — seek then take the next frame.
|
||||
A fresh container per frame keeps the decode robust (no abandoned-generator state)."""
|
||||
try:
|
||||
import av
|
||||
with av.open(path) as c:
|
||||
st = c.streams.video[0]
|
||||
if t > 0 and st.time_base:
|
||||
c.seek(int(t / st.time_base), stream=st)
|
||||
for frame in c.decode(st):
|
||||
return _encode(frame.to_image())
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _frames(path: Path) -> list[str] | None:
|
||||
"""Up to three ~512px JPEG frames (base64) to judge, or None when unreadable. A video
|
||||
(.mp4/.webm) is sampled at start / middle / end so the motion judge sees temporal
|
||||
coherence; any other file is the single downscaled image. None → the caller decides
|
||||
(an image fails open; a video is marked ungated, never a silent pass)."""
|
||||
if path.suffix.lower() in (".mp4", ".webm"):
|
||||
try:
|
||||
import av
|
||||
with av.open(str(path)) as probe:
|
||||
st = probe.streams.video[0]
|
||||
if st.duration and st.time_base:
|
||||
dur = float(st.duration * st.time_base)
|
||||
elif probe.duration:
|
||||
dur = float(probe.duration) / 1_000_000
|
||||
else:
|
||||
dur = 0.0
|
||||
targets = [0.0, dur / 2, dur] if dur > 0 else [0.0]
|
||||
out = [b for b in (_video_frame(str(path), t) for t in targets) if b]
|
||||
return out or None
|
||||
except Exception:
|
||||
return None
|
||||
b = _downscale(path)
|
||||
return [b] if b else None
|
||||
|
||||
|
||||
def _reply_text(data: dict) -> str:
|
||||
"""The assistant text from an OpenAI-style completion, tolerating a content
|
||||
string, the content-parts array some gateways return, and REASONING models
|
||||
whose verdict lands in reasoning_content while content comes back null (a
|
||||
tiny max_tokens made zen-vl spend the whole budget thinking — every verdict
|
||||
parsed empty and the gate failed open, silently passing corrupted renders)."""
|
||||
try:
|
||||
msg = data["choices"][0]["message"]
|
||||
except (KeyError, IndexError, TypeError):
|
||||
return ""
|
||||
content = msg.get("content") or ""
|
||||
if isinstance(content, list):
|
||||
content = "".join(p.get("text", "") for p in content if isinstance(p, dict))
|
||||
if not str(content).strip():
|
||||
content = msg.get("reasoning_content") or ""
|
||||
return str(content)
|
||||
|
||||
|
||||
async def judge(images: list[str], bearer: str, prompt: str) -> str | None:
|
||||
"""Ask the vision judge to score N frames against ``prompt`` (PROMPT for a still image,
|
||||
MOTION for a video's sampled frames) — one text part + N image parts. Returns its reply
|
||||
text, or None on ANY error (non-200, bad JSON, timeout, unreachable) — caller fails open."""
|
||||
content = [{"type": "text", "text": prompt}]
|
||||
for b64 in images:
|
||||
content.append({"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + b64}})
|
||||
payload = {
|
||||
"model": _model(),
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"temperature": 0,
|
||||
"max_tokens": 768, # reasoning models think before the verdict — 8 starved them into fail-open
|
||||
}
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if bearer:
|
||||
headers["Authorization"] = bearer
|
||||
try:
|
||||
async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=45)) as s:
|
||||
async with s.post(_url(), json=payload, headers=headers) as r:
|
||||
if r.status != 200:
|
||||
log.warning("gate: judge HTTP %s", r.status)
|
||||
return None
|
||||
return _reply_text(await r.json())
|
||||
except Exception as e:
|
||||
log.warning("gate: judge unreachable: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
def _flag(org: str, root: Path, rel: str, reply: str) -> bool:
|
||||
"""Move a still-``draft`` asset to ``flagged`` and record it on its work-log row.
|
||||
Only ``draft`` transitions — a late verdict must never yank a render a human has
|
||||
since approved/published (or one already flagged/deleted). One writer of
|
||||
library.json (studio_home's helpers). Returns True if the asset was flagged."""
|
||||
from middleware import studio_home as sh # lazy: avoid an import cycle
|
||||
lib = sh._load_library(root)
|
||||
for a in lib.get("assets", []):
|
||||
if a.get("path") == rel and a.get("status") == "draft":
|
||||
a["status"] = "flagged"
|
||||
a["updatedAt"] = int(time.time())
|
||||
sh._save_library(root, lib)
|
||||
worklog.add_event(org, rel, "flagged",
|
||||
"quality gate: " + (reply or "").strip()[:80])
|
||||
log.info("gate: FLAGGED %s/%s", org, rel)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _skip(org: str, root: Path, rel: str) -> None:
|
||||
"""A video whose frames could not be sampled is marked visibly UNGATED — NEVER a
|
||||
silent pass: set the asset row's ``gate`` field to ``skipped`` (the card renders an
|
||||
'ungated' badge from it) and log an ``ungated`` work-log event. Same single-writer
|
||||
library pattern as ``_flag``; leaves status untouched."""
|
||||
from middleware import studio_home as sh # lazy: avoid an import cycle
|
||||
lib = sh._load_library(root)
|
||||
hit = False
|
||||
for a in lib.get("assets", []):
|
||||
if a.get("path") == rel:
|
||||
a["gate"] = "skipped"
|
||||
a["updatedAt"] = int(time.time())
|
||||
hit = True
|
||||
if hit:
|
||||
sh._save_library(root, lib)
|
||||
worklog.add_event(org, rel, "ungated", "gate could not sample frames")
|
||||
log.info("gate: UNGATED %s/%s", org, rel)
|
||||
|
||||
|
||||
async def run(org: str, root: Path, rel: str, bearer: str) -> None:
|
||||
"""Judge one just-stored render — a still image against PROMPT, a video against MOTION
|
||||
over its sampled frames — and flag it on a clear FAIL. A video whose frames can't be
|
||||
sampled is marked UNGATED (never a silent pass); an unreadable image fails open. Never
|
||||
raises — a background task that fails open on everything else."""
|
||||
try:
|
||||
p = Path(root) / rel
|
||||
if not p.is_file():
|
||||
return
|
||||
video = p.suffix.lower() in (".mp4", ".webm")
|
||||
frames = _frames(p)
|
||||
if frames is None:
|
||||
if video:
|
||||
_skip(org, root, rel) # honest: an ungated video is badged, never passed silently
|
||||
else:
|
||||
log.info("gate: unreadable %s/%s — pass (fail-open)", org, rel)
|
||||
return
|
||||
reply = await judge(frames, bearer, MOTION if video else PROMPT)
|
||||
verdict = parse_verdict(reply or "")
|
||||
if verdict == "fail":
|
||||
_flag(org, root, rel, reply or "")
|
||||
else:
|
||||
log.info("gate: %s %s/%s", verdict, org, rel)
|
||||
except Exception as e: # a background task must never surface an exception
|
||||
log.info("gate: error %s/%s: %s — pass (fail-open)", org, rel, e)
|
||||
|
||||
|
||||
def schedule(request, org: str, root: Path, rel: str) -> None:
|
||||
"""Fire the judgment out-of-band. Captures the caller's bearer NOW (the request
|
||||
is gone by the time the task runs) and schedules ``run`` on the event loop.
|
||||
Never blocks the upload response; never raises."""
|
||||
if not enabled():
|
||||
return
|
||||
try:
|
||||
bearer = _bearer(request)
|
||||
task = asyncio.ensure_future(run(org, root, rel, bearer))
|
||||
_TASKS.add(task)
|
||||
task.add_done_callback(_TASKS.discard)
|
||||
except Exception as e:
|
||||
log.info("gate: schedule failed %s/%s: %s", org, rel, e)
|
||||
@@ -0,0 +1,63 @@
|
||||
/* Hanzo Studio shell — the ONE shared chrome, loaded by BOTH the Studio home
|
||||
* (middleware/studio_home.html) and the Editor (injected into the engine's
|
||||
* index.html by server.get_root). Self-contained dark palette + `hz` id/class
|
||||
* prefix so it never collides with the editor's own theme. */
|
||||
#hzbar{position:fixed;bottom:18px;right:18px;z-index:2147483000;display:flex;align-items:center;gap:8px;
|
||||
flex-wrap:wrap;justify-content:flex-end;max-width:calc(100vw - 36px);
|
||||
font:400 14px/1.4 Inter,system-ui,-apple-system,sans-serif}
|
||||
#hzbar button,#hzbar a{font:inherit}
|
||||
.hz-pill{display:inline-flex;align-items:center;gap:7px;background:#131317;color:#ececf1;border:1px solid #232329;
|
||||
border-radius:999px;padding:7px 13px;cursor:pointer;text-decoration:none;white-space:nowrap}
|
||||
.hz-pill:hover{border-color:#3a3a48}
|
||||
.hz-pill .hz-dot{width:8px;height:8px;border-radius:50%;background:#555;flex:0 0 auto}
|
||||
.hz-pill.on .hz-dot{background:#3ddc84}
|
||||
.hz-pill.hz-chat{background:#f4f4f6;color:#111;border-color:#f4f4f6;font-weight:600}
|
||||
.hz-topup{color:#6aa5ff;text-decoration:none;font-size:.82em;border-left:1px solid #232329;padding-left:7px}
|
||||
.hz-topup:hover{text-decoration:underline}
|
||||
/* right-docked chat sidebar (both views) */
|
||||
#hzchat{position:fixed;top:0;right:0;height:100vh;width:380px;max-width:94vw;background:#0e0e12;color:#ececf1;
|
||||
border-left:1px solid #232329;z-index:2147483001;display:flex;flex-direction:column;
|
||||
transform:translateX(100%);transition:transform .22s ease;box-shadow:-14px 0 44px #0007}
|
||||
#hzchat.on{transform:translateX(0)}
|
||||
#hzchat .hz-h{display:flex;align-items:center;gap:9px;padding:13px 15px;border-bottom:1px solid #232329;flex-wrap:wrap}
|
||||
#hzchat .hz-h .hz-av{width:26px;height:26px;border-radius:50%;background:#26262e;display:flex;align-items:center;
|
||||
justify-content:center;font-size:.72rem;font-weight:600;flex:0 0 auto}
|
||||
#hzchat .hz-h .hz-who{font-size:.82rem;min-width:0}
|
||||
#hzchat .hz-h .hz-who b{display:block;font-size:.86rem;white-space:nowrap;overflow:hidden;text-overflow:ellipsis}
|
||||
#hzchat .hz-h .hz-who span{color:#8a8a95;font-size:.72rem}
|
||||
#hzchat .hz-h .hz-x{margin-left:auto;cursor:pointer;color:#8a8a95;font-size:1.3rem;line-height:1;background:none;border:none}
|
||||
#hzmsgs{flex:1;overflow:auto;padding:14px;display:flex;flex-direction:column;gap:10px}
|
||||
.hz-m{max-width:88%;padding:9px 12px;border-radius:12px;font-size:.86rem;line-height:1.45;white-space:pre-wrap;word-break:break-word}
|
||||
.hz-m.u{align-self:flex-end;background:#f4f4f6;color:#111;border-bottom-right-radius:3px}
|
||||
.hz-m.a{align-self:flex-start;background:#1a1a20;border:1px solid #232329;border-bottom-left-radius:3px}
|
||||
.hz-m.sys{align-self:center;color:#8a8a95;font-size:.78rem;background:none}
|
||||
#hzform{display:flex;gap:8px;padding:12px;border-top:1px solid #232329}
|
||||
#hzform input{flex:1;background:#0a0a0d;border:1px solid #232329;border-radius:9px;color:#ececf1;padding:9px 11px;font:inherit;font-size:.85rem;outline:none}
|
||||
#hzform button{background:#f4f4f6;color:#111;border:none;border-radius:9px;padding:0 14px;font-weight:600;cursor:pointer}
|
||||
/* GPUs panel (both views) */
|
||||
#hzgpus{position:fixed;bottom:64px;right:18px;z-index:2147483001;width:320px;max-width:92vw;background:#131317;
|
||||
color:#ececf1;border:1px solid #232329;border-radius:12px;padding:12px 14px;display:none;box-shadow:0 18px 50px #0009}
|
||||
#hzgpus.on{display:block}
|
||||
#hzgpus h4{font-size:.82rem;font-weight:600;margin:0 0 8px;display:flex;align-items:center;gap:7px}
|
||||
.hz-node{display:flex;align-items:flex-start;gap:9px;padding:7px 0;border-top:1px solid #1e1e24;font-size:.8rem}
|
||||
.hz-node:first-of-type{border-top:none}
|
||||
.hz-node .hz-dot{width:8px;height:8px;border-radius:50%;background:#555;margin-top:5px;flex:0 0 auto}
|
||||
.hz-node.on .hz-dot{background:#3ddc84}
|
||||
.hz-node .hz-nn{flex:1;min-width:0}
|
||||
.hz-node .hz-nn b{font-weight:500}
|
||||
.hz-node .hz-nsub{color:#8a8a95;font-size:.72rem;margin-top:2px}
|
||||
/* Tablet: a touch narrower sidebar; the pill bar + panels already fit. */
|
||||
@media (max-width:900px){#hzchat{width:340px}}
|
||||
/* Phone: chat becomes a full-width bottom sheet (same pattern the old chatpanel
|
||||
used); the pill bar wraps upward from the bottom-right so the toggle, wallet
|
||||
and GPUs badge all stay reachable; the GPUs panel spans the width. */
|
||||
@media (max-width:600px){
|
||||
#hzbar{bottom:14px;right:14px;gap:7px}
|
||||
.hz-pill{padding:7px 11px}
|
||||
.hz-pill .hz-lbl{display:none} /* icon-only Chat/GPUs → one line, clears content */
|
||||
#hzchat{top:auto;bottom:0;width:100vw;max-width:100vw;height:85vh;
|
||||
border-left:none;border-top:1px solid #232329;border-radius:16px 16px 0 0;
|
||||
transform:translateY(100%);box-shadow:0 -14px 44px #0007}
|
||||
#hzchat.on{transform:translateY(0)}
|
||||
#hzgpus{left:12px;right:12px;width:auto;bottom:66px}
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
/* Hanzo Studio shell — the ONE shared chrome for BOTH views.
|
||||
*
|
||||
* Studio home (/studio) includes this via <script src="/studio/shell.js">
|
||||
* Editor (/ ?advanced=1) gets the same tag injected into the engine's
|
||||
* index.html by server.get_root
|
||||
*
|
||||
* It renders (identically in both views): a Studio/Editor toggle, a wallet chip,
|
||||
* a GPUs badge + panel, and a dockable right chat SIDEBAR backed by the ONE chat
|
||||
* path (POST /v1/copilot/chat). Same-origin cookies carry the IAM session, so org
|
||||
* + login are inherited — no auth is rebuilt here. In the editor it also honours
|
||||
* ?load=<worklog id> / ?template=<name> to open a graph, and offers Save as
|
||||
* template. Guarded end to end: any failure logs and no-ops, never throws. */
|
||||
(function () {
|
||||
"use strict";
|
||||
if (window.__hzShell) return;
|
||||
window.__hzShell = 1;
|
||||
|
||||
var STUDIO = location.pathname === "/studio" || location.pathname === "/studio/";
|
||||
var $ = function (id) { return document.getElementById(id); };
|
||||
var el = function (tag, attrs, html) {
|
||||
var e = document.createElement(tag);
|
||||
if (attrs) for (var k in attrs) e.setAttribute(k, attrs[k]);
|
||||
if (html != null) e.innerHTML = html;
|
||||
return e;
|
||||
};
|
||||
var esc = function (s) { return String(s == null ? "" : s).replace(/[<>&]/g, function (c) { return { "<": "<", ">": ">", "&": "&" }[c]; }); };
|
||||
async function getJSON(url, opt) {
|
||||
var r = await fetch(url, Object.assign({ credentials: "same-origin" }, opt || {}));
|
||||
var j = {}; try { j = await r.json(); } catch (_) {}
|
||||
return { ok: r.ok, status: r.status, body: j };
|
||||
}
|
||||
|
||||
// ── chrome: bottom-right pill cluster + docked chat sidebar + GPUs panel ──────
|
||||
var bar = el("div", { id: "hzbar" });
|
||||
var toggle = el("a", { class: "hz-pill", href: STUDIO ? "/?advanced=1" : "/studio",
|
||||
title: STUDIO ? "Open the node editor" : "Back to Studio" }, STUDIO ? "⇄ Editor" : "⇄ Studio");
|
||||
var tmplBtn = STUDIO ? null : el("button", { class: "hz-pill", title: "Save the current graph as a reusable template" }, "★ Template");
|
||||
var wallet = el("a", { class: "hz-pill", id: "hzwallet", href: "https://pay.hanzo.ai", target: "_blank", rel: "noopener",
|
||||
title: "Wallet — top up at pay.hanzo.ai" }, "wallet …");
|
||||
var gpus = el("button", { class: "hz-pill", id: "hzgpubtn", title: "GPUs" }, '<span class="hz-dot"></span><span class="hz-lbl">GPUs</span>');
|
||||
var chatBtn = el("button", { class: "hz-pill hz-chat", title: "Chat" }, '💬<span class="hz-lbl"> Chat</span>');
|
||||
bar.appendChild(toggle);
|
||||
if (tmplBtn) bar.appendChild(tmplBtn);
|
||||
bar.appendChild(wallet); bar.appendChild(gpus); bar.appendChild(chatBtn);
|
||||
|
||||
var gpuPanel = el("div", { id: "hzgpus" }, '<h4>GPUs</h4><div id="hznodes"></div>');
|
||||
var chat = el("aside", { id: "hzchat" });
|
||||
chat.innerHTML =
|
||||
'<div class="hz-h"><span class="hz-av" id="hzav">·</span>' +
|
||||
'<span class="hz-who"><b id="hzuser">…</b><span id="hzorg"></span></span>' +
|
||||
'<button class="hz-x" id="hzclose" title="Close">×</button></div>' +
|
||||
'<div id="hzmsgs"></div>' +
|
||||
'<form id="hzform"><input id="hzin" placeholder="Ask the copilot…" autocomplete="off"><button type="submit">Send</button></form>';
|
||||
|
||||
function mount() {
|
||||
document.body.appendChild(bar);
|
||||
document.body.appendChild(gpuPanel);
|
||||
document.body.appendChild(chat);
|
||||
wire();
|
||||
loadSession(); loadWallet(); loadGpus();
|
||||
if (!STUDIO) editorBoot();
|
||||
}
|
||||
|
||||
// ── chat: the ONE path — POST /v1/copilot/chat; ops applied via the editor's
|
||||
// existing applier when present (no second chat mechanism) ────────────────────
|
||||
var HISTORY = [];
|
||||
function addMsg(role, text) {
|
||||
var m = el("div", { class: "hz-m " + role }, esc(text));
|
||||
$("hzmsgs").appendChild(m); $("hzmsgs").scrollTop = $("hzmsgs").scrollHeight;
|
||||
return m;
|
||||
}
|
||||
async function graphContext() {
|
||||
var app = window.app;
|
||||
if (!app || typeof app.graphToPrompt !== "function") return {};
|
||||
try {
|
||||
var p = await app.graphToPrompt();
|
||||
var ctx = { graph_json: p.output || {} };
|
||||
try { ctx.catalog = { types: Object.keys((window.LiteGraph && window.LiteGraph.registered_node_types) || {}) }; } catch (_) {}
|
||||
return ctx;
|
||||
} catch (_) { return {}; }
|
||||
}
|
||||
async function send(text) {
|
||||
HISTORY.push({ role: "user", content: text });
|
||||
addMsg("u", text);
|
||||
var pending = addMsg("sys", "…");
|
||||
var payload = Object.assign({ messages: HISTORY.slice(-20) }, await graphContext());
|
||||
var r = await getJSON("/v1/copilot/chat", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(payload) });
|
||||
pending.remove();
|
||||
var reply = (r.body && r.body.reply) || (r.body && r.body.error) || (r.ok ? "" : "chat failed (" + r.status + ")");
|
||||
if (reply) { HISTORY.push({ role: "assistant", content: reply }); addMsg("a", reply); }
|
||||
var ops = (r.body && r.body.ops) || [];
|
||||
if (ops.length && window.HanzoCopilot && typeof window.HanzoCopilot.applyOps === "function") {
|
||||
try { window.HanzoCopilot.applyOps(ops); } catch (e) { addMsg("sys", "couldn't apply changes: " + e.message); }
|
||||
} else if (ops.length && STUDIO) {
|
||||
addMsg("sys", "Open the Editor to apply " + ops.length + " graph change" + (ops.length > 1 ? "s" : "") + ".");
|
||||
}
|
||||
}
|
||||
|
||||
function wire() {
|
||||
chatBtn.onclick = function () { chat.classList.toggle("on"); if (chat.classList.contains("on")) $("hzin").focus(); };
|
||||
$("hzclose").onclick = function () { chat.classList.remove("on"); };
|
||||
$("hzform").onsubmit = function (e) {
|
||||
e.preventDefault();
|
||||
var v = $("hzin").value.trim(); if (!v) return;
|
||||
$("hzin").value = ""; send(v).catch(function (err) { addMsg("sys", "error: " + err.message); });
|
||||
};
|
||||
gpus.onclick = function () { gpuPanel.classList.toggle("on"); if (gpuPanel.classList.contains("on")) loadGpus(); };
|
||||
if (tmplBtn) tmplBtn.onclick = saveTemplate;
|
||||
}
|
||||
|
||||
// ── session · wallet · GPUs ───────────────────────────────────────────────────
|
||||
async function loadSession() {
|
||||
var r = await getJSON("/v1/session"); var s = r.body || {};
|
||||
var nm = (s.user || s.email || "You").trim();
|
||||
$("hzav").textContent = (nm[0] || "·").toUpperCase();
|
||||
$("hzuser").textContent = nm;
|
||||
$("hzorg").textContent = s.org ? ("org · " + s.org) : "";
|
||||
}
|
||||
function money(n, cur) {
|
||||
if (n == null || isNaN(n)) return "—";
|
||||
var sym = (cur || "").toLowerCase() === "usd" ? "$" : "";
|
||||
return sym + Number(n).toFixed(2) + (sym ? "" : " " + (cur || "").toUpperCase());
|
||||
}
|
||||
async function loadWallet() {
|
||||
var r = await getJSON("/v1/wallet"); var w = r.body || {};
|
||||
wallet.innerHTML = esc(money(w.balance, w.currency)) + '<span class="hz-topup">Top up ↗</span>';
|
||||
wallet.title = "Wallet" + (w.org ? " · " + w.org : "") + " — top up at pay.hanzo.ai";
|
||||
}
|
||||
async function loadGpus() {
|
||||
var nres = await getJSON("/v1/nodes"); var qres = await getJSON("/v1/queue/status");
|
||||
var nodes = (nres.body && nres.body.nodes) || [];
|
||||
var pod = (nres.body && nres.body.pod) || null;
|
||||
if (pod) nodes = nodes.concat([pod]);
|
||||
var items = (qres.body && qres.body.items) || {};
|
||||
var byNode = {};
|
||||
Object.keys(items).forEach(function (id) {
|
||||
var it = items[id]; var k = it.node || "gpu";
|
||||
(byNode[k] = byNode[k] || { running: null, depth: 0 });
|
||||
byNode[k].depth++;
|
||||
if (it.status === "running") byNode[k].running = it.prompt || it.kind || "a render";
|
||||
});
|
||||
var online = nodes.filter(function (n) { return n.online; }).length;
|
||||
gpus.innerHTML = '<span class="hz-dot"></span><span class="hz-lbl">GPUs </span>' + online + "/" + nodes.length;
|
||||
gpus.classList.toggle("on", online > 0);
|
||||
$("hznodes").innerHTML = nodes.length ? nodes.map(function (n) {
|
||||
var s = byNode[n.name] || { running: null, depth: 0 };
|
||||
// Authoritative depth from /v1/nodes (gpu-jobs activities) when present; the
|
||||
// running-render label still comes from the org's live queue/status.
|
||||
var depth = (typeof n.depth === "number") ? n.depth : s.depth;
|
||||
var sub = (n.online ? "online" : "offline") +
|
||||
(s.running ? " · running: " + esc(String(s.running).slice(0, 40)) : (n.online ? " · idle" : "")) +
|
||||
" · queue " + depth;
|
||||
return '<div class="hz-node ' + (n.online ? "on" : "") + '"><span class="hz-dot"></span>' +
|
||||
'<div class="hz-nn"><b>' + esc(n.name) + '</b><div class="hz-nsub">' + sub + "</div></div></div>";
|
||||
}).join("") : '<div class="hz-nsub">No GPU nodes connected. Run <b>hanzo gpu connect</b> on a box to add one.</div>';
|
||||
}
|
||||
|
||||
// ── editor-only: open a graph (?load / ?template), Save as template ────────────
|
||||
function waitForApp(ms) {
|
||||
return new Promise(function (resolve, reject) {
|
||||
var t0 = Date.now();
|
||||
(function poll() {
|
||||
if (window.app && typeof window.app.loadGraphData === "function") return resolve(window.app);
|
||||
if (Date.now() - t0 > ms) return reject(new Error("editor not ready"));
|
||||
setTimeout(poll, 150);
|
||||
})();
|
||||
});
|
||||
}
|
||||
function editorLoad(g) {
|
||||
var app = window.app; if (!app || !g) return false;
|
||||
if (g.workflow && typeof app.loadGraphData === "function") { app.loadGraphData(g.workflow); return true; }
|
||||
if (g.output && typeof app.loadApiJson === "function") { app.loadApiJson(g.output, "hanzo"); return true; }
|
||||
if ((g.nodes || g.last_node_id != null) && typeof app.loadGraphData === "function") { app.loadGraphData(g); return true; }
|
||||
if (typeof app.loadApiJson === "function") { app.loadApiJson(g, "hanzo"); return true; }
|
||||
if (typeof app.loadGraphData === "function") { app.loadGraphData(g); return true; }
|
||||
return false;
|
||||
}
|
||||
async function editorBoot() {
|
||||
var q = new URLSearchParams(location.search);
|
||||
var load = q.get("load"), template = q.get("template");
|
||||
if (!load && !template) return;
|
||||
try {
|
||||
await waitForApp(30000);
|
||||
var url = load ? "/v1/workflow?id=" + encodeURIComponent(load) : "/v1/templates/" + encodeURIComponent(template);
|
||||
var r = await getJSON(url);
|
||||
if (!r.ok) { addMsg("sys", "couldn't open: " + ((r.body && r.body.error) || r.status)); return; }
|
||||
editorLoad(r.body.graph || r.body);
|
||||
} catch (e) { console.warn("[hz shell]", e); }
|
||||
}
|
||||
async function saveTemplate() {
|
||||
var app = window.app;
|
||||
if (!app || typeof app.graphToPrompt !== "function") { alert("Open a workflow in the editor first."); return; }
|
||||
var name = prompt("Save this graph as a template named:");
|
||||
if (!name || !name.trim()) return;
|
||||
var g;
|
||||
try { g = await app.graphToPrompt(); } catch (e) { alert("Couldn't read the graph: " + e.message); return; }
|
||||
var r = await getJSON("/v1/templates", { method: "POST", headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ name: name.trim(), graph: { workflow: g.workflow, output: g.output } }) });
|
||||
alert(r.ok ? '★ Saved template "' + (r.body.name || name.trim()) + '"' : "Save failed: " + ((r.body && r.body.error) || r.status));
|
||||
}
|
||||
|
||||
if (document.body) mount();
|
||||
else document.addEventListener("DOMContentLoaded", mount);
|
||||
})();
|
||||
@@ -0,0 +1,598 @@
|
||||
"""Stacks — Procreate-style organizational folders for the Studio gallery.
|
||||
|
||||
A Stack groups library assets into a named folder (drag one image onto another,
|
||||
or multi-select → Stack). Stacks are ONLY an organizational layer: the assets
|
||||
themselves stay exactly where they are (library.json rows + files), and lineage,
|
||||
jobs and storage remain the source of truth. Removing an image from a Stack
|
||||
returns it to the gallery; it is never a delete. Deleting a Stack defaults to
|
||||
preserving every image.
|
||||
|
||||
Backed by ONE per-org SQLite file (stdlib sqlite3, WAL) at ``orgs/<org>/stacks.db``
|
||||
— the proven ``tasks_queue.py`` pattern: WAL + NORMAL + busy_timeout, and every
|
||||
mutation runs under ``BEGIN IMMEDIATE`` so create / add / move / remove / reorder /
|
||||
delete are atomic. Tenancy is structural: a different org is a different FILE, so
|
||||
a stack from another org is unreachable by construction (every route 404s across
|
||||
orgs). One stack per asset is enforced by ``UNIQUE(asset_id)`` on the membership
|
||||
table — a membership row exists only while the asset lives in a (non-deleted)
|
||||
Stack, so "one non-deleted Stack per asset" falls out of the constraint.
|
||||
|
||||
The schema is exactly the spec's ``stacks`` + ``stack_assets`` join, plus two
|
||||
small sidecar tables that keep generation-integration and the per-org setting in
|
||||
the same tenant-scoped file:
|
||||
|
||||
stack_dest pending "outputs of this render → this Stack" routing (prefix→id)
|
||||
settings per-org key/value (e.g. auto_stack_batches)
|
||||
|
||||
Standard library only — no new dependency.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
import folder_paths
|
||||
|
||||
_UNTITLED = "Untitled Stack"
|
||||
_MAX_NAME = 80
|
||||
_MAX_DESC = 2000
|
||||
_MAX_ASSET = 512
|
||||
_MAX_ASSETS = 2000 # a Stack is a folder, not a database — bound it
|
||||
_DEST_TTL = 6 * 3600.0 # a render lands within minutes; prune stale routing
|
||||
|
||||
_SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS stacks (
|
||||
id TEXT PRIMARY KEY,
|
||||
organization_id TEXT NOT NULL,
|
||||
workspace_id TEXT,
|
||||
owner_id TEXT,
|
||||
name TEXT NOT NULL,
|
||||
description TEXT NOT NULL DEFAULT '',
|
||||
cover_asset_id TEXT,
|
||||
created_at INTEGER NOT NULL,
|
||||
updated_at INTEGER NOT NULL,
|
||||
archived_at INTEGER,
|
||||
deleted_at INTEGER
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS stack_assets (
|
||||
id TEXT PRIMARY KEY,
|
||||
stack_id TEXT NOT NULL,
|
||||
asset_id TEXT NOT NULL UNIQUE, -- one Stack per asset
|
||||
position INTEGER NOT NULL,
|
||||
added_by TEXT,
|
||||
created_at INTEGER NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS ix_stack_assets_stack ON stack_assets(stack_id, position);
|
||||
CREATE TABLE IF NOT EXISTS stack_dest (
|
||||
prefix TEXT PRIMARY KEY,
|
||||
stack_id TEXT NOT NULL,
|
||||
added_by TEXT,
|
||||
created_at REAL NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS settings (
|
||||
key TEXT PRIMARY KEY,
|
||||
value TEXT NOT NULL
|
||||
);
|
||||
"""
|
||||
|
||||
_conns: dict[str, sqlite3.Connection] = {}
|
||||
_locks: dict[str, threading.Lock] = {}
|
||||
_registry_lock = threading.Lock()
|
||||
|
||||
|
||||
def _db_path(org: str) -> Path:
|
||||
"""``orgs/<org>/stacks.db`` — a per-org sibling of the output tree that holds
|
||||
library.json/worklog.json/templates, so a Stack DB is tenant-scoped exactly
|
||||
like every other per-org file (same resolution as ``_templates_root``)."""
|
||||
out = Path(folder_paths.get_org_output_directory(org))
|
||||
base = out.parent if out.parent.name == org else out
|
||||
base.mkdir(parents=True, exist_ok=True)
|
||||
return base / "stacks.db"
|
||||
|
||||
|
||||
def _conn(org: str) -> tuple[sqlite3.Connection, threading.Lock]:
|
||||
"""The org's connection + its write lock, opened once and reused (one durable
|
||||
handle per tenant file, like the render queue)."""
|
||||
path = str(_db_path(org))
|
||||
with _registry_lock:
|
||||
c = _conns.get(path)
|
||||
if c is None:
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
c = sqlite3.connect(path, check_same_thread=False, isolation_level=None)
|
||||
c.row_factory = sqlite3.Row
|
||||
c.execute("PRAGMA journal_mode=WAL")
|
||||
c.execute("PRAGMA synchronous=NORMAL")
|
||||
c.execute("PRAGMA busy_timeout=5000")
|
||||
c.executescript(_SCHEMA)
|
||||
_conns[path] = c
|
||||
_locks[path] = threading.Lock()
|
||||
return c, _locks[path]
|
||||
|
||||
|
||||
def _now() -> int:
|
||||
return int(time.time())
|
||||
|
||||
|
||||
def _clean_ids(raw) -> list[str]:
|
||||
"""A client asset-id list → safe, de-duped, order-preserving, bounded."""
|
||||
out: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for x in (raw if isinstance(raw, list) else []):
|
||||
if not isinstance(x, str):
|
||||
continue
|
||||
s = x.strip()[:_MAX_ASSET]
|
||||
if s and s not in seen:
|
||||
seen.add(s)
|
||||
out.append(s)
|
||||
if len(out) >= _MAX_ASSETS:
|
||||
break
|
||||
return out
|
||||
|
||||
|
||||
def _row(stack: sqlite3.Row) -> dict:
|
||||
return {
|
||||
"id": stack["id"],
|
||||
"organizationId": stack["organization_id"],
|
||||
"workspaceId": stack["workspace_id"],
|
||||
"ownerId": stack["owner_id"],
|
||||
"name": stack["name"],
|
||||
"description": stack["description"] or "",
|
||||
"coverAssetId": stack["cover_asset_id"],
|
||||
"createdAt": stack["created_at"],
|
||||
"updatedAt": stack["updated_at"],
|
||||
"archivedAt": stack["archived_at"],
|
||||
}
|
||||
|
||||
|
||||
# ── membership primitives (all run inside a caller-held BEGIN IMMEDIATE) ─────────
|
||||
def _detach(c: sqlite3.Connection, asset_id: str, now: int) -> None:
|
||||
"""Remove an asset from whatever Stack currently holds it, keeping that Stack's
|
||||
cover valid — the move-between-Stacks and create-steals-asset invariant. Does
|
||||
NOT touch the asset itself (files/library rows are never Stacks' concern)."""
|
||||
row = c.execute("SELECT stack_id FROM stack_assets WHERE asset_id=?", (asset_id,)).fetchone()
|
||||
if row is None:
|
||||
return
|
||||
src = row["stack_id"]
|
||||
c.execute("DELETE FROM stack_assets WHERE asset_id=?", (asset_id,))
|
||||
cover = c.execute("SELECT cover_asset_id FROM stacks WHERE id=?", (src,)).fetchone()
|
||||
if cover and cover["cover_asset_id"] == asset_id:
|
||||
nxt = c.execute("SELECT asset_id FROM stack_assets WHERE stack_id=? ORDER BY position LIMIT 1",
|
||||
(src,)).fetchone()
|
||||
c.execute("UPDATE stacks SET cover_asset_id=?, updated_at=? WHERE id=?",
|
||||
(nxt["asset_id"] if nxt else None, now, src))
|
||||
else:
|
||||
c.execute("UPDATE stacks SET updated_at=? WHERE id=?", (now, src))
|
||||
|
||||
|
||||
def _append(c: sqlite3.Connection, stack_id: str, asset_ids: list[str], added_by: str, now: int) -> None:
|
||||
"""Append assets to a Stack in order, each first detached from its old Stack."""
|
||||
nextpos = c.execute("SELECT COALESCE(MAX(position)+1,0) AS n FROM stack_assets WHERE stack_id=?",
|
||||
(stack_id,)).fetchone()["n"]
|
||||
for aid in asset_ids:
|
||||
_detach(c, aid, now)
|
||||
c.execute("INSERT INTO stack_assets(id, stack_id, asset_id, position, added_by, created_at)"
|
||||
" VALUES(?,?,?,?,?,?)",
|
||||
(uuid.uuid4().hex, stack_id, aid, nextpos, added_by, now))
|
||||
nextpos += 1
|
||||
|
||||
|
||||
def _members(c: sqlite3.Connection, stack_id: str) -> list[dict]:
|
||||
rows = c.execute("SELECT asset_id, position, added_by, created_at FROM stack_assets"
|
||||
" WHERE stack_id=? ORDER BY position", (stack_id,)).fetchall()
|
||||
return [{"assetId": r["asset_id"], "position": r["position"],
|
||||
"addedBy": r["added_by"], "createdAt": r["created_at"]} for r in rows]
|
||||
|
||||
|
||||
# ── public API (each opens its own BEGIN IMMEDIATE; tenancy = the org's file) ────
|
||||
def create(org: str, name: str, asset_ids, *, owner: str = "", workspace: str | None = None) -> dict:
|
||||
name = (name or "").strip()[:_MAX_NAME] or _UNTITLED
|
||||
ids = _clean_ids(asset_ids)
|
||||
sid = uuid.uuid4().hex
|
||||
now = _now()
|
||||
c, lock = _conn(org)
|
||||
with lock:
|
||||
c.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
c.execute("INSERT INTO stacks(id, organization_id, workspace_id, owner_id, name,"
|
||||
" description, cover_asset_id, created_at, updated_at)"
|
||||
" VALUES(?,?,?,?,?,?,?,?,?)",
|
||||
(sid, org, workspace, owner, name, "", ids[0] if ids else None, now, now))
|
||||
_append(c, sid, ids, owner, now)
|
||||
c.execute("COMMIT")
|
||||
except Exception:
|
||||
c.execute("ROLLBACK")
|
||||
raise
|
||||
return get(org, sid)
|
||||
|
||||
|
||||
def get(org: str, sid: str) -> dict | None:
|
||||
c, _ = _conn(org)
|
||||
row = c.execute("SELECT * FROM stacks WHERE id=? AND organization_id=? AND deleted_at IS NULL",
|
||||
(sid, org)).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
out = _row(row)
|
||||
out["assets"] = _members(c, sid)
|
||||
out["count"] = len(out["assets"])
|
||||
return out
|
||||
|
||||
|
||||
def list_stacks(org: str, *, include_archived: bool = True) -> list[dict]:
|
||||
"""All of the org's live Stacks, newest-touched first, each with its member
|
||||
asset ids (the gallery hides stacked assets from the loose grid and draws the
|
||||
layered thumbnail from the cover + first members)."""
|
||||
c, _ = _conn(org)
|
||||
q = "SELECT * FROM stacks WHERE organization_id=? AND deleted_at IS NULL"
|
||||
if not include_archived:
|
||||
q += " AND archived_at IS NULL"
|
||||
q += " ORDER BY updated_at DESC"
|
||||
out = []
|
||||
for row in c.execute(q, (org,)).fetchall():
|
||||
s = _row(row)
|
||||
members = _members(c, row["id"])
|
||||
s["assets"] = members
|
||||
s["assetIds"] = [m["assetId"] for m in members]
|
||||
s["count"] = len(members)
|
||||
out.append(s)
|
||||
return out
|
||||
|
||||
|
||||
def patch(org: str, sid: str, fields: dict) -> dict | None:
|
||||
now = _now()
|
||||
c, lock = _conn(org)
|
||||
with lock:
|
||||
c.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
row = c.execute("SELECT id FROM stacks WHERE id=? AND organization_id=? AND deleted_at IS NULL",
|
||||
(sid, org)).fetchone()
|
||||
if row is None:
|
||||
c.execute("COMMIT")
|
||||
return None
|
||||
updates: list[tuple[str, object]] = []
|
||||
if "name" in fields:
|
||||
updates.append(("name", (str(fields.get("name") or "").strip()[:_MAX_NAME]) or _UNTITLED))
|
||||
if "description" in fields:
|
||||
updates.append(("description", str(fields.get("description") or "")[:_MAX_DESC]))
|
||||
if "cover" in fields or "coverAssetId" in fields:
|
||||
cover = str(fields.get("cover") or fields.get("coverAssetId") or "").strip()[:_MAX_ASSET]
|
||||
# A cover must be a member — else the card would show a ghost.
|
||||
member = c.execute("SELECT 1 FROM stack_assets WHERE stack_id=? AND asset_id=?",
|
||||
(sid, cover)).fetchone()
|
||||
if cover and member is None:
|
||||
raise ValueError("cover must be an asset in this Stack")
|
||||
updates.append(("cover_asset_id", cover or None))
|
||||
if "archived" in fields:
|
||||
updates.append(("archived_at", now if fields.get("archived") else None))
|
||||
updates.append(("updated_at", now))
|
||||
cols = ", ".join(f"{col}=?" for col, _ in updates)
|
||||
args = [val for _, val in updates] + [sid, org]
|
||||
c.execute(f"UPDATE stacks SET {cols} WHERE id=? AND organization_id=?", args)
|
||||
c.execute("COMMIT")
|
||||
except Exception:
|
||||
c.execute("ROLLBACK")
|
||||
raise
|
||||
return get(org, sid)
|
||||
|
||||
|
||||
def add_assets(org: str, sid: str, asset_ids, *, added_by: str = "") -> dict | None:
|
||||
ids = _clean_ids(asset_ids)
|
||||
now = _now()
|
||||
c, lock = _conn(org)
|
||||
with lock:
|
||||
c.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
row = c.execute("SELECT cover_asset_id FROM stacks WHERE id=? AND organization_id=? AND deleted_at IS NULL",
|
||||
(sid, org)).fetchone()
|
||||
if row is None:
|
||||
c.execute("COMMIT")
|
||||
return None
|
||||
_append(c, sid, ids, added_by, now)
|
||||
if row["cover_asset_id"] is None and ids:
|
||||
c.execute("UPDATE stacks SET cover_asset_id=? WHERE id=?", (ids[0], sid))
|
||||
c.execute("UPDATE stacks SET updated_at=? WHERE id=?", (now, sid))
|
||||
c.execute("COMMIT")
|
||||
except Exception:
|
||||
c.execute("ROLLBACK")
|
||||
raise
|
||||
return get(org, sid)
|
||||
|
||||
|
||||
def remove_assets(org: str, sid: str, asset_ids) -> dict | None:
|
||||
"""Remove assets from a Stack — they return to the main gallery, never deleted.
|
||||
Positions compact; a removed cover falls back to the new first member."""
|
||||
ids = _clean_ids(asset_ids)
|
||||
now = _now()
|
||||
c, lock = _conn(org)
|
||||
with lock:
|
||||
c.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
row = c.execute("SELECT cover_asset_id FROM stacks WHERE id=? AND organization_id=? AND deleted_at IS NULL",
|
||||
(sid, org)).fetchone()
|
||||
if row is None:
|
||||
c.execute("COMMIT")
|
||||
return None
|
||||
for aid in ids:
|
||||
c.execute("DELETE FROM stack_assets WHERE stack_id=? AND asset_id=?", (sid, aid))
|
||||
# compact positions to keep them contiguous and stable
|
||||
rest = c.execute("SELECT id FROM stack_assets WHERE stack_id=? ORDER BY position",
|
||||
(sid,)).fetchall()
|
||||
for i, r in enumerate(rest):
|
||||
c.execute("UPDATE stack_assets SET position=? WHERE id=?", (i, r["id"]))
|
||||
cover = row["cover_asset_id"]
|
||||
if cover in ids:
|
||||
nxt = c.execute("SELECT asset_id FROM stack_assets WHERE stack_id=? ORDER BY position LIMIT 1",
|
||||
(sid,)).fetchone()
|
||||
c.execute("UPDATE stacks SET cover_asset_id=? WHERE id=?",
|
||||
(nxt["asset_id"] if nxt else None, sid))
|
||||
c.execute("UPDATE stacks SET updated_at=? WHERE id=?", (now, sid))
|
||||
c.execute("COMMIT")
|
||||
except Exception:
|
||||
c.execute("ROLLBACK")
|
||||
raise
|
||||
return get(org, sid)
|
||||
|
||||
|
||||
def reorder(org: str, sid: str, asset_ids) -> dict | None:
|
||||
"""Persist a new order. The given ids must be exactly the Stack's current
|
||||
membership (same set) — a partial/foreign list is rejected so the order can
|
||||
never silently drop or invent a member."""
|
||||
ids = _clean_ids(asset_ids)
|
||||
now = _now()
|
||||
c, lock = _conn(org)
|
||||
with lock:
|
||||
c.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
row = c.execute("SELECT id FROM stacks WHERE id=? AND organization_id=? AND deleted_at IS NULL",
|
||||
(sid, org)).fetchone()
|
||||
if row is None:
|
||||
c.execute("COMMIT")
|
||||
return None
|
||||
cur = {r["asset_id"] for r in c.execute("SELECT asset_id FROM stack_assets WHERE stack_id=?", (sid,)).fetchall()}
|
||||
if set(ids) != cur:
|
||||
raise ValueError("order must list exactly the Stack's assets")
|
||||
for i, aid in enumerate(ids):
|
||||
c.execute("UPDATE stack_assets SET position=? WHERE stack_id=? AND asset_id=?", (i, sid, aid))
|
||||
c.execute("UPDATE stacks SET updated_at=? WHERE id=?", (now, sid))
|
||||
c.execute("COMMIT")
|
||||
except Exception:
|
||||
c.execute("ROLLBACK")
|
||||
raise
|
||||
return get(org, sid)
|
||||
|
||||
|
||||
def _dissolve(org: str, sid: str) -> list[str] | None:
|
||||
"""Soft-delete the container and free its members (membership rows removed so
|
||||
the assets are unstacked and one-per-asset stays true). Returns the freed asset
|
||||
ids, or None if the Stack isn't the caller-org's live Stack. Backs both Unstack
|
||||
and Delete-Stack."""
|
||||
now = _now()
|
||||
c, lock = _conn(org)
|
||||
with lock:
|
||||
c.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
row = c.execute("SELECT id FROM stacks WHERE id=? AND organization_id=? AND deleted_at IS NULL",
|
||||
(sid, org)).fetchone()
|
||||
if row is None:
|
||||
c.execute("COMMIT")
|
||||
return None
|
||||
members = [r["asset_id"] for r in
|
||||
c.execute("SELECT asset_id FROM stack_assets WHERE stack_id=? ORDER BY position", (sid,)).fetchall()]
|
||||
c.execute("DELETE FROM stack_assets WHERE stack_id=?", (sid,))
|
||||
c.execute("UPDATE stacks SET deleted_at=?, updated_at=? WHERE id=?", (now, now, sid))
|
||||
c.execute("COMMIT")
|
||||
except Exception:
|
||||
c.execute("ROLLBACK")
|
||||
raise
|
||||
return members
|
||||
|
||||
|
||||
def unstack(org: str, sid: str) -> list[str] | None:
|
||||
"""Delete the container, preserve every asset to the main gallery."""
|
||||
return _dissolve(org, sid)
|
||||
|
||||
|
||||
def delete(org: str, sid: str, mode: str = "only") -> dict | None:
|
||||
"""Delete a Stack. ``only`` (default) returns its images to the gallery;
|
||||
``images`` also asks the caller to soft-delete them. Files are never touched.
|
||||
Returns ``{ok, mode, assetIds}`` (assetIds = the freed/soft-deletable assets),
|
||||
or None across orgs / already gone."""
|
||||
freed = _dissolve(org, sid)
|
||||
if freed is None:
|
||||
return None
|
||||
return {"ok": True, "mode": ("images" if mode == "images" else "only"), "assetIds": freed}
|
||||
|
||||
|
||||
def stack_of(org: str, asset_id: str) -> str | None:
|
||||
c, _ = _conn(org)
|
||||
row = c.execute("SELECT stack_id FROM stack_assets WHERE asset_id=?", (asset_id,)).fetchone()
|
||||
return row["stack_id"] if row else None
|
||||
|
||||
|
||||
# ── generation integration: route a render's outputs into a chosen Stack ─────────
|
||||
def route_prefix(org: str, output_prefix: str, sid: str, *, added_by: str = "") -> None:
|
||||
"""Remember "every output of the render at ``output_prefix`` goes into Stack
|
||||
``sid``". Recorded at dispatch; consumed by :func:`absorb` when each output
|
||||
lands (so ALL images of a multi-output job land in the same Stack). Stale
|
||||
routing is pruned — a render lands within minutes."""
|
||||
if not output_prefix or not sid:
|
||||
return
|
||||
now = time.time()
|
||||
c, lock = _conn(org)
|
||||
with lock:
|
||||
c.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
c.execute("DELETE FROM stack_dest WHERE created_at < ?", (now - _DEST_TTL,))
|
||||
c.execute("INSERT INTO stack_dest(prefix, stack_id, added_by, created_at) VALUES(?,?,?,?)"
|
||||
" ON CONFLICT(prefix) DO UPDATE SET stack_id=excluded.stack_id,"
|
||||
" added_by=excluded.added_by, created_at=excluded.created_at",
|
||||
(output_prefix, sid, added_by, now))
|
||||
c.execute("COMMIT")
|
||||
except Exception:
|
||||
c.execute("ROLLBACK")
|
||||
raise
|
||||
|
||||
|
||||
_COUNTER = re.compile(r"_\d+_\.\w+$")
|
||||
|
||||
|
||||
def absorb(org: str, rel: str, *, added_by: str = "") -> str | None:
|
||||
"""A just-ingested asset ``rel`` → the Stack its render was routed to, if any.
|
||||
Matches a landed ``<prefix>_<counter>_.png`` back to its ``output_prefix`` (and
|
||||
the ingested-path case where the stored path IS the prefix). Idempotent via the
|
||||
membership UNIQUE. Returns the Stack id it joined, or None."""
|
||||
if not rel:
|
||||
return None
|
||||
c, _ = _conn(org)
|
||||
cand = c.execute("SELECT stack_id FROM stack_dest WHERE prefix=?", (rel,)).fetchone()
|
||||
if cand is None:
|
||||
prefix = _COUNTER.sub("", rel)
|
||||
if prefix != rel:
|
||||
cand = c.execute("SELECT stack_id FROM stack_dest WHERE prefix=?", (prefix,)).fetchone()
|
||||
if cand is None:
|
||||
return None
|
||||
sid = cand["stack_id"]
|
||||
if add_assets(org, sid, [rel], added_by=added_by or "generation") is None:
|
||||
return None
|
||||
return sid
|
||||
|
||||
|
||||
# ── per-org settings (auto-stack batch generations) ─────────────────────────────
|
||||
def get_setting(org: str, key: str, default=None):
|
||||
c, _ = _conn(org)
|
||||
row = c.execute("SELECT value FROM settings WHERE key=?", (key,)).fetchone()
|
||||
if row is None:
|
||||
return default
|
||||
try:
|
||||
return json.loads(row["value"])
|
||||
except Exception:
|
||||
return default
|
||||
|
||||
|
||||
def set_setting(org: str, key: str, value) -> None:
|
||||
c, lock = _conn(org)
|
||||
with lock:
|
||||
c.execute("INSERT INTO settings(key, value) VALUES(?,?)"
|
||||
" ON CONFLICT(key) DO UPDATE SET value=excluded.value",
|
||||
(key, json.dumps(value)))
|
||||
|
||||
|
||||
# ── HTTP routes: /v1/stacks… (house style: /v1, never /api). Dependencies are
|
||||
# INJECTED (org / owner / workspace resolvers + the library soft-deleter) so this
|
||||
# module stays decoupled from studio_home — no import cycle, one stacks concern. ──
|
||||
def add_stacks_routes(routes: web.RouteTableDef, server, *, org_of, owner_of, workspace_of,
|
||||
soft_delete_assets) -> None:
|
||||
|
||||
async def _json(request):
|
||||
try:
|
||||
return await request.json()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@routes.get("/v1/stacks")
|
||||
async def list_route(request: web.Request):
|
||||
org = org_of(request)
|
||||
return web.json_response({"org": org, "stacks": list_stacks(org)})
|
||||
|
||||
@routes.post("/v1/stacks")
|
||||
async def create_route(request: web.Request):
|
||||
body = await _json(request)
|
||||
if body is None:
|
||||
return web.json_response({"error": "invalid JSON"}, status=400)
|
||||
stack = create(org_of(request), body.get("name"), body.get("assetIds"),
|
||||
owner=owner_of(request), workspace=workspace_of(request))
|
||||
return web.json_response(stack)
|
||||
|
||||
# /v1/stacks/settings MUST precede /v1/stacks/{id} — else "settings" reads as an id.
|
||||
@routes.get("/v1/stacks/settings")
|
||||
async def get_settings_route(request: web.Request):
|
||||
org = org_of(request)
|
||||
return web.json_response({"org": org, "autoStack": bool(get_setting(org, "auto_stack_batches", False))})
|
||||
|
||||
@routes.patch("/v1/stacks/settings")
|
||||
async def set_settings_route(request: web.Request):
|
||||
body = await _json(request)
|
||||
if body is None:
|
||||
return web.json_response({"error": "invalid JSON"}, status=400)
|
||||
org = org_of(request)
|
||||
if "autoStack" in body:
|
||||
set_setting(org, "auto_stack_batches", bool(body.get("autoStack")))
|
||||
return web.json_response({"org": org, "autoStack": bool(get_setting(org, "auto_stack_batches", False))})
|
||||
|
||||
@routes.get("/v1/stacks/{id}")
|
||||
async def get_route(request: web.Request):
|
||||
stack = get(org_of(request), request.match_info["id"])
|
||||
if stack is None:
|
||||
return web.json_response({"error": "not found"}, status=404)
|
||||
return web.json_response(stack)
|
||||
|
||||
@routes.patch("/v1/stacks/{id}")
|
||||
async def patch_route(request: web.Request):
|
||||
body = await _json(request)
|
||||
if body is None:
|
||||
return web.json_response({"error": "invalid JSON"}, status=400)
|
||||
try:
|
||||
stack = patch(org_of(request), request.match_info["id"], body)
|
||||
except ValueError as e:
|
||||
return web.json_response({"error": str(e)}, status=400)
|
||||
if stack is None:
|
||||
return web.json_response({"error": "not found"}, status=404)
|
||||
return web.json_response(stack)
|
||||
|
||||
@routes.delete("/v1/stacks/{id}")
|
||||
async def delete_route(request: web.Request):
|
||||
org = org_of(request)
|
||||
mode = "images" if request.query.get("mode") == "images" else "only"
|
||||
result = delete(org, request.match_info["id"], mode)
|
||||
if result is None:
|
||||
return web.json_response({"error": "not found"}, status=404)
|
||||
if mode == "images" and result["assetIds"]:
|
||||
soft_delete_assets(org, result["assetIds"])
|
||||
return web.json_response(result)
|
||||
|
||||
@routes.post("/v1/stacks/{id}/assets")
|
||||
async def add_assets_route(request: web.Request):
|
||||
body = await _json(request)
|
||||
if body is None:
|
||||
return web.json_response({"error": "invalid JSON"}, status=400)
|
||||
stack = add_assets(org_of(request), request.match_info["id"], body.get("assetIds"),
|
||||
added_by=owner_of(request))
|
||||
if stack is None:
|
||||
return web.json_response({"error": "not found"}, status=404)
|
||||
return web.json_response(stack)
|
||||
|
||||
@routes.delete("/v1/stacks/{id}/assets")
|
||||
async def remove_assets_route(request: web.Request):
|
||||
body = await _json(request)
|
||||
if body is None:
|
||||
return web.json_response({"error": "invalid JSON"}, status=400)
|
||||
stack = remove_assets(org_of(request), request.match_info["id"], body.get("assetIds"))
|
||||
if stack is None:
|
||||
return web.json_response({"error": "not found"}, status=404)
|
||||
return web.json_response(stack)
|
||||
|
||||
@routes.patch("/v1/stacks/{id}/order")
|
||||
async def order_route(request: web.Request):
|
||||
body = await _json(request)
|
||||
if body is None:
|
||||
return web.json_response({"error": "invalid JSON"}, status=400)
|
||||
try:
|
||||
stack = reorder(org_of(request), request.match_info["id"], body.get("assetIds"))
|
||||
except ValueError as e:
|
||||
return web.json_response({"error": str(e)}, status=400)
|
||||
if stack is None:
|
||||
return web.json_response({"error": "not found"}, status=404)
|
||||
return web.json_response(stack)
|
||||
|
||||
@routes.post("/v1/stacks/{id}/unstack")
|
||||
async def unstack_route(request: web.Request):
|
||||
freed = unstack(org_of(request), request.match_info["id"])
|
||||
if freed is None:
|
||||
return web.json_response({"error": "not found"}, status=404)
|
||||
return web.json_response({"ok": True, "assetIds": freed})
|
||||
+1085
-174
File diff suppressed because it is too large
Load Diff
+2034
-185
File diff suppressed because it is too large
Load Diff
@@ -7,6 +7,7 @@ When Studio runs with --worker-mode, this module handles:
|
||||
3. Reporting device capabilities (GPU model, VRAM, etc.)
|
||||
"""
|
||||
import asyncio
|
||||
import hmac
|
||||
import logging
|
||||
import os
|
||||
import time as _time
|
||||
@@ -17,6 +18,7 @@ import aiohttp
|
||||
from aiohttp import web
|
||||
|
||||
import execution
|
||||
from studio.cli_args import args
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -31,10 +33,31 @@ WORKER_TOKEN = os.environ.get("STUDIO_WORKER_TOKEN", "")
|
||||
|
||||
|
||||
def verify_worker_token(request) -> bool:
|
||||
"""True if the request carries the coordinator shared secret (or none is set)."""
|
||||
"""True iff the request carries the coordinator shared secret. FAILS CLOSED: in
|
||||
worker mode the secret is MANDATORY (main.py refuses to boot without it, and the
|
||||
KMS /v1 migration has emptied that secret before), so an empty/unset token is a
|
||||
valid single-trust-domain dev signal ONLY when NOT in worker mode — never a bypass
|
||||
that reopens the hidden-run hole on a worker box. Constant-time comparison."""
|
||||
if not WORKER_TOKEN:
|
||||
return True
|
||||
return request.headers.get("X-Worker-Token", "") == WORKER_TOKEN
|
||||
return not args.worker_mode
|
||||
return hmac.compare_digest(request.headers.get("X-Worker-Token", ""), WORKER_TOKEN)
|
||||
|
||||
|
||||
def reject_untrusted_worker_submit(request, worker_mode: bool):
|
||||
"""The ONE policy gate for worker-mode /prompt (decomplected: one function, one
|
||||
place). In worker mode this process is a render BACKEND, not a front — the only
|
||||
accepted submit is a job claimed off the org's gpu-jobs queue and handed in over the
|
||||
coordinator seam (POST /v1/worker/execute with X-Worker-Token). A direct /prompt POST
|
||||
without a valid token is refused (403), so no render can reach a GPU without appearing
|
||||
in the org's visible queue. Returns an aiohttp Response to send, or None to allow.
|
||||
Not worker mode → always allowed (returns None)."""
|
||||
if worker_mode and not verify_worker_token(request):
|
||||
return web.json_response({"error": {
|
||||
"type": "worker_mode_forbidden",
|
||||
"message": "worker-mode renders run only from the gpu-jobs queue "
|
||||
"(POST /v1/worker/execute); direct /prompt is disabled.",
|
||||
"details": "", "extra_info": {}}, "node_errors": {}}, status=403)
|
||||
return None
|
||||
|
||||
|
||||
def _worker_headers() -> dict:
|
||||
@@ -172,6 +195,8 @@ def add_worker_routes(routes: web.RouteTableDef, prompt_server):
|
||||
return web.json_response({"error": "No prompt provided"}, status=400)
|
||||
|
||||
prompt = json_data["prompt"]
|
||||
if not isinstance(prompt, dict):
|
||||
return web.json_response({"error": "prompt must be an object"}, status=400)
|
||||
prompt_id = str(json_data.get("prompt_id", uuid.uuid4()))
|
||||
|
||||
partial_execution_targets = json_data.get("partial_execution_targets")
|
||||
@@ -197,8 +222,19 @@ def add_worker_routes(routes: web.RouteTableDef, prompt_server):
|
||||
if sensitive_val in extra_data:
|
||||
sensitive[sensitive_val] = extra_data.pop(sensitive_val)
|
||||
|
||||
number = float(json_data.get("number", 0))
|
||||
try:
|
||||
number = float(json_data.get("number", 0))
|
||||
except (TypeError, ValueError):
|
||||
return web.json_response({"error": "number must be numeric"}, status=400)
|
||||
extra_data["create_time"] = int(_time.time() * 1000)
|
||||
# `org` scopes outputs on THIS worker's LOCAL disk only — a worker-LOCAL label,
|
||||
# never a gallery destination. The durable upload path re-derives the gallery org
|
||||
# from the user's IAM token (not this body field), so a forged body `org` can at
|
||||
# most mislabel a local temp dir, never write into another tenant's gallery. The
|
||||
# prompt worker binds folder_paths.set_execution_org from extra_data["org_id"].
|
||||
org = json_data.get("org") or extra_data.get("org_id")
|
||||
if org:
|
||||
extra_data["org_id"] = org
|
||||
|
||||
# Queue locally — the worker's prompt_worker thread will pick it up
|
||||
prompt_server.prompt_queue.put(
|
||||
|
||||
@@ -59,6 +59,20 @@ def _save(org: str, rows: list[dict]) -> None:
|
||||
tmp.replace(p)
|
||||
|
||||
|
||||
def has_output(org: str, path: str) -> bool:
|
||||
"""Whether any row already owns this landed file — the ingest upsert key
|
||||
(re-ingesting must repair, never duplicate). A dispatch row's output_prefix
|
||||
lacks the SaveImage counter suffix (``fixes/x_ab12cd`` owns
|
||||
``fixes/x_ab12cd_00001_.png``), so the match is prefix-aware, not equality —
|
||||
equality alone doubled every mirrored render with an ingest row."""
|
||||
for r in load(org):
|
||||
pre = r.get("output_prefix") or ""
|
||||
if path and (path == pre or path == r.get("output") or
|
||||
(pre and path.startswith(pre + "_"))):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def record(org: str, *, kind: str, prompt: str, refs: list, uploads: list,
|
||||
parents: list, output_prefix: str, params: dict | None = None,
|
||||
pid: str | None = None, lane: str | None = None, node: str | None = None,
|
||||
@@ -92,6 +106,22 @@ def record(org: str, *, kind: str, prompt: str, refs: list, uploads: list,
|
||||
return row
|
||||
|
||||
|
||||
def remove(org: str, ids) -> int:
|
||||
"""Delete work-log rows by id (a queued/cancelled/failed/done entry the user wants
|
||||
out of their queue). Returns how many were removed. The log is just a record of
|
||||
dispatch requests, so dropping a row is safe — it never touches an output image
|
||||
(those live in the library and are archived/deleted there). Idempotent."""
|
||||
want = {str(i) for i in (ids if isinstance(ids, (list, tuple, set)) else [ids]) if i}
|
||||
if not want:
|
||||
return 0
|
||||
rows = load(org)
|
||||
kept = [r for r in rows if str(r.get("id")) not in want]
|
||||
removed = len(rows) - len(kept)
|
||||
if removed:
|
||||
_save(org, kept)
|
||||
return removed
|
||||
|
||||
|
||||
def set_status(org: str, pid: str, status: str, note: str | None = None) -> bool:
|
||||
"""Advance a row's status by prompt_id (running/done/failed/cancelled). Returns
|
||||
True if a row was updated. Idempotent-safe; terminal rows are left untouched."""
|
||||
@@ -112,6 +142,28 @@ def set_status(org: str, pid: str, status: str, note: str | None = None) -> bool
|
||||
return hit
|
||||
|
||||
|
||||
def add_event(org: str, output_prefix: str, s: str, note: str | None = None) -> bool:
|
||||
"""Append one event to the row(s) carrying this landed output — an annotation,
|
||||
not a lifecycle move. The render completed (its status is terminal); a later
|
||||
judgment like the quality gate flagging it records an event without touching
|
||||
that status. Keyed by output like ``has_output``. Returns True if a row was
|
||||
annotated."""
|
||||
if not output_prefix:
|
||||
return False
|
||||
rows = load(org)
|
||||
now = int(time.time())
|
||||
hit = False
|
||||
for r in rows:
|
||||
if r.get("output_prefix") == output_prefix or r.get("output") == output_prefix:
|
||||
r.setdefault("events", []).append(
|
||||
{"s": s, "ts": now, **({"note": note} if note else {})})
|
||||
r["updatedAt"] = now
|
||||
hit = True
|
||||
if hit:
|
||||
_save(org, rows)
|
||||
return hit
|
||||
|
||||
|
||||
def mark_started(org: str, pid: str, at: int) -> bool:
|
||||
"""Stamp when a job first begins rendering (first seen at the head of its lane).
|
||||
Idempotent: only the first observation sets started_at. Feeds duration = finished
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "hanzo-studio"
|
||||
version = "0.17.3"
|
||||
version = "0.17.23"
|
||||
readme = "README.md"
|
||||
license = { file = "LICENSE" }
|
||||
requires-python = ">=3.10"
|
||||
|
||||
+122
-12
@@ -29,13 +29,44 @@ Pure stdlib — no PIL, no deps — so it runs in the studio image or anywhere.
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError: # non-POSIX (never in the studio image); lock is a no-op
|
||||
fcntl = None
|
||||
|
||||
LOCK_NAME = ".library.lock"
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def library_lock(root: str):
|
||||
"""Serialize the read-modify-write of library.json across processes — the
|
||||
build-manifest sidecar and the studio app both rebuild-then-replace it, so an
|
||||
unlocked interleave loses status/caption/tags updates. Same lock file name is
|
||||
used by the app writer (studio_home._save_library)."""
|
||||
if fcntl is None:
|
||||
yield
|
||||
return
|
||||
lock_path = os.path.join(root, LOCK_NAME)
|
||||
fh = open(lock_path, "w")
|
||||
try:
|
||||
fcntl.flock(fh, fcntl.LOCK_EX)
|
||||
yield
|
||||
finally:
|
||||
with contextlib.suppress(OSError):
|
||||
fcntl.flock(fh, fcntl.LOCK_UN)
|
||||
fh.close()
|
||||
|
||||
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp"}
|
||||
VIDEO_EXTS = {".mp4", ".webm"}
|
||||
MODEL3D_EXTS = {".glb"}
|
||||
# Written marketing content (blog posts, campaign briefs, social ad copy) is
|
||||
# first-class in the ONE index too — a .md under marketing/ is a queue asset
|
||||
# whose body is the post/brief and whose caption/hashtags(=tags) are set with
|
||||
@@ -54,6 +85,19 @@ def iso(ts: float) -> str:
|
||||
return datetime.fromtimestamp(ts, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
|
||||
|
||||
def _sha_of(path: str) -> str | None:
|
||||
"""Streamed content hash (a 40MB PNG must not balloon memory). None on error."""
|
||||
import hashlib
|
||||
try:
|
||||
h = hashlib.md5()
|
||||
with open(path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(1 << 20), b""):
|
||||
h.update(chunk)
|
||||
return h.hexdigest()
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
|
||||
def classify(relpath: str, slugs: set[str]) -> dict | None:
|
||||
"""Map a path relative to the library root -> {design, kind, role}.
|
||||
|
||||
@@ -91,15 +135,23 @@ def classify(relpath: str, slugs: set[str]) -> dict | None:
|
||||
def scan(root: str) -> list[str]:
|
||||
"""All manifestable image paths (relative to root), sorted for stable diffs."""
|
||||
found: list[str] = []
|
||||
for base, _dirs, files in os.walk(root):
|
||||
for base, dirs, files in os.walk(root):
|
||||
# Hidden segments are never assets — one rule covering AppleDouble
|
||||
# forks (._foo.png) AND cache trees (.thumbs/ thumbnails indexed as
|
||||
# library rows sent thumb paths into the fix dispatcher).
|
||||
dirs[:] = [d for d in dirs if not d.startswith(".")]
|
||||
for f in files:
|
||||
if f.startswith("."):
|
||||
continue
|
||||
ext = os.path.splitext(f)[1].lower()
|
||||
full = os.path.join(base, f)
|
||||
rel = os.path.relpath(full, root)
|
||||
is_image = ext in IMAGE_EXTS
|
||||
is_video = ext in VIDEO_EXTS
|
||||
is_model3d = ext in MODEL3D_EXTS
|
||||
is_mktg_text = (ext in MARKETING_TEXT_EXTS
|
||||
and rel.replace(os.sep, "/").startswith("marketing/"))
|
||||
if not (is_image or is_mktg_text):
|
||||
if not (is_image or is_video or is_model3d or is_mktg_text):
|
||||
continue
|
||||
try:
|
||||
if os.path.getsize(full) == 0: # skip a render caught mid-write
|
||||
@@ -133,9 +185,25 @@ def build(root: str) -> dict:
|
||||
if norm == MANIFEST_NAME:
|
||||
continue
|
||||
info = classify(norm, slugs)
|
||||
if info is None:
|
||||
continue
|
||||
prev = prior.get(norm, {})
|
||||
if info is None:
|
||||
# Taxonomy enriches, never gates: a file outside the naming scheme is
|
||||
# still an asset. Dropping it here erased every ingest-indexed render on
|
||||
# each sweep cycle. A video keeps kind="video" so it survives every sweep
|
||||
# as the motion-judged, video-card kind.
|
||||
_e = os.path.splitext(norm)[1].lower()
|
||||
kind = prev.get("kind") or (
|
||||
"model3d" if _e in MODEL3D_EXTS else "video" if _e in VIDEO_EXTS else "render")
|
||||
info = {"design": prev.get("design"), "kind": kind, "role": prev.get("role")}
|
||||
abspath = os.path.join(root, rel)
|
||||
mtime = os.path.getmtime(abspath)
|
||||
# Content hash powers the dedup invariant below. Reuse the prior sha when the
|
||||
# file is byte-stable (same mtime), so a sweep hashes only NEW/changed files —
|
||||
# O(new), not O(all), even on a large gallery. Only images are hashed (docs are
|
||||
# not the byte-duplicate-clutter concern).
|
||||
is_img = norm.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))
|
||||
sha = prev.get("sha") if (prev.get("mtime") == mtime and prev.get("sha")) \
|
||||
else (_sha_of(abspath) if is_img else None)
|
||||
entry = {
|
||||
"path": norm,
|
||||
"design": info["design"],
|
||||
@@ -143,13 +211,39 @@ def build(root: str) -> dict:
|
||||
"role": info["role"],
|
||||
"status": prev.get("status", "draft"),
|
||||
"tags": prev.get("tags", []),
|
||||
"updatedAt": iso(os.path.getmtime(os.path.join(root, rel))),
|
||||
"mtime": mtime,
|
||||
"sha": sha,
|
||||
"updatedAt": iso(mtime),
|
||||
}
|
||||
caption = prev.get("caption")
|
||||
if caption:
|
||||
entry["caption"] = caption
|
||||
assets.append(entry)
|
||||
|
||||
# ── Structural invariant: AT MOST ONE active asset per content hash. A re-rendered
|
||||
# byte-identical duplicate (a retried / double-dispatched render) is auto-marked
|
||||
# deleted here — so the gallery can NEVER accumulate duplicate images, whatever bug
|
||||
# fires upstream. Keeper = a curated asset if one exists, else the earliest by
|
||||
# mtime; distinct renders (different seed → different bytes → different sha) are
|
||||
# untouched. Idempotent (a prior-swept dup stays deleted) and recoverable like any
|
||||
# soft-delete. This is the last line of defense that makes duplicate clutter
|
||||
# impossible-by-construction, independent of dispatch/retry correctness. ──
|
||||
from collections import defaultdict
|
||||
_CURATED = {"approved", "published", "queued", "flagged"}
|
||||
_by_sha: dict = defaultdict(list)
|
||||
for a in assets:
|
||||
if a.get("sha") and a["status"] != "deleted":
|
||||
_by_sha[a["sha"]].append(a)
|
||||
deduped = 0
|
||||
for group in _by_sha.values():
|
||||
if len(group) < 2:
|
||||
continue
|
||||
group.sort(key=lambda a: (a["status"] not in _CURATED, a.get("mtime", 0)))
|
||||
for dup in group[1:]:
|
||||
dup["status"] = "deleted"
|
||||
dup["dedup"] = True
|
||||
deduped += 1
|
||||
|
||||
by_status: dict[str, int] = {}
|
||||
by_kind: dict[str, int] = {}
|
||||
for a in assets:
|
||||
@@ -172,12 +266,24 @@ def build(root: str) -> dict:
|
||||
|
||||
|
||||
def write_atomic(root: str, doc: dict) -> str:
|
||||
# A UNIQUE tmp per writer: a fixed ".library.json.tmp" is shared by every
|
||||
# build-manifest sidecar, so two pods overlapping during a roll interleave into
|
||||
# one inode -> a torn library.json -> load_prior()/reader gets JSONDecodeError
|
||||
# -> {} -> every asset resets to status="draft" (the curation wipe). mkstemp in
|
||||
# the SAME dir keeps the rename atomic and same-filesystem.
|
||||
manifest_path = os.path.join(root, MANIFEST_NAME)
|
||||
tmp = os.path.join(root, ".library.json.tmp")
|
||||
with open(tmp, "w") as fh:
|
||||
json.dump(doc, fh, indent=2, ensure_ascii=False)
|
||||
fh.write("\n")
|
||||
os.replace(tmp, manifest_path)
|
||||
fd, tmp = tempfile.mkstemp(prefix=".library.json.", suffix=".tmp", dir=root)
|
||||
try:
|
||||
with os.fdopen(fd, "w") as fh:
|
||||
json.dump(doc, fh, indent=2, ensure_ascii=False)
|
||||
fh.write("\n")
|
||||
os.replace(tmp, manifest_path)
|
||||
except BaseException:
|
||||
try:
|
||||
os.unlink(tmp)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
return manifest_path
|
||||
|
||||
|
||||
@@ -189,8 +295,12 @@ def run_once(root: str, quiet: bool = False, strict: bool = True) -> dict | None
|
||||
if not quiet:
|
||||
print(msg, file=sys.stderr)
|
||||
return None # watch mode: tolerate a tenant with no renders yet
|
||||
doc = build(root)
|
||||
path = write_atomic(root, doc)
|
||||
# Hold the lock across build (which reads the prior manifest to preserve
|
||||
# status/caption/tags) AND the write, so a concurrent app-side status change
|
||||
# can't be read-then-clobbered.
|
||||
with library_lock(root):
|
||||
doc = build(root)
|
||||
path = write_atomic(root, doc)
|
||||
if not quiet:
|
||||
m = doc["_meta"]
|
||||
print(f"library_manifest: {m['count']} assets -> {path} {m['byStatus']}")
|
||||
|
||||
@@ -382,11 +382,14 @@ class PromptServer():
|
||||
|
||||
@routes.get("/ready")
|
||||
async def readiness_check(request):
|
||||
# Check that the prompt worker is alive and base directories exist
|
||||
ready = (
|
||||
os.path.isdir(folder_paths.get_input_directory())
|
||||
and self.prompt_queue is not None
|
||||
)
|
||||
# The render worker must actually be ALIVE — not merely that the queue
|
||||
# object exists (prompt_queue is set once at boot and never cleared, so
|
||||
# it is constant-true and says nothing about the executor). A dead worker
|
||||
# daemon thread accepts /prompt but runs nothing; readiness must fail so
|
||||
# the pod leaves the Service and a roll never promotes a wedged pod.
|
||||
wt = getattr(self, "prompt_worker_thread", None)
|
||||
worker_alive = wt.is_alive() if wt is not None else (self.prompt_queue is not None)
|
||||
ready = os.path.isdir(folder_paths.get_input_directory()) and worker_alive
|
||||
if ready:
|
||||
return web.json_response({"status": "ready"})
|
||||
return web.json_response({"status": "not_ready"}, status=503)
|
||||
@@ -435,11 +438,11 @@ class PromptServer():
|
||||
},
|
||||
)
|
||||
|
||||
response = web.FileResponse(os.path.join(self.web_root, "index.html"))
|
||||
response.headers['Cache-Control'] = 'no-cache'
|
||||
response.headers["Pragma"] = "no-cache"
|
||||
response.headers["Expires"] = "0"
|
||||
return response
|
||||
# Serve the editor SPA with the shared Studio shell injected (the ONE
|
||||
# chrome — Studio/Editor toggle, wallet, GPUs, chat sidebar — that the
|
||||
# Studio home also includes). No-cache, same as before.
|
||||
from middleware.studio_home import editor_index
|
||||
return editor_index(self.web_root)
|
||||
|
||||
@routes.get("/embeddings")
|
||||
def get_embeddings(request):
|
||||
@@ -577,7 +580,19 @@ class PromptServer():
|
||||
# Force type=output regardless of what the client sent.
|
||||
forced = dict(post)
|
||||
forced["type"] = "output"
|
||||
return image_upload(forced, request=request)
|
||||
resp = image_upload(forced, request=request)
|
||||
# Fix outputs: restore the render's global colour to its source, undoing
|
||||
# the full-regen (denoise 1.0) white-balance/exposure drift and its
|
||||
# fix-of-a-fix compounding. Pod-side + fail-safe, so it can never affect a
|
||||
# render's success; a render that didn't drift matches to itself untouched.
|
||||
try:
|
||||
saved = json.loads(resp.body) if getattr(resp, "body", None) else {}
|
||||
if saved.get("name"):
|
||||
from middleware.studio_home import colormatch_fix_ingest
|
||||
colormatch_fix_ingest(request, saved.get("subfolder", ""), saved["name"])
|
||||
except Exception:
|
||||
pass
|
||||
return resp
|
||||
|
||||
@routes.post("/upload/mask")
|
||||
async def upload_mask(request):
|
||||
@@ -1028,6 +1043,15 @@ class PromptServer():
|
||||
async def post_prompt(request):
|
||||
logging.info("got prompt")
|
||||
|
||||
# Execution-engine (worker) mode: this process IS a BYO-GPU render backend,
|
||||
# not a front. A direct /prompt POST without the coordinator token is refused
|
||||
# so no render reaches a GPU without appearing in the org's visible queue.
|
||||
# One way onto a GPU: enqueue -> claim -> execute. (Policy: worker_client.)
|
||||
from middleware.worker_client import reject_untrusted_worker_submit
|
||||
_gate = reject_untrusted_worker_submit(request, args.worker_mode)
|
||||
if _gate is not None:
|
||||
return _gate
|
||||
|
||||
# Billing: check balance before accepting prompt
|
||||
if args.enable_billing and args.billing_check_balance:
|
||||
iam_user = request.get("iam_user")
|
||||
@@ -1047,7 +1071,11 @@ class PromptServer():
|
||||
json_data = self.trigger_on_prompt(json_data)
|
||||
|
||||
# --- Prompt routing: check if this should go to a GPU worker ---
|
||||
if not args.worker_mode:
|
||||
# Legacy in-cluster PUSH federation (prompt_router -> compute_config
|
||||
# workers) is OFF by default: the gpu-jobs enqueue below (dispatch_if_worker)
|
||||
# is the ONE submit path. Set STUDIO_LEGACY_PUSH_ROUTER=1 only to resurrect
|
||||
# the old push seam. Keeps exactly one way onto a GPU in prod.
|
||||
if not args.worker_mode and os.environ.get("STUDIO_LEGACY_PUSH_ROUTER") == "1":
|
||||
org_id = self._get_org_id(request)
|
||||
try:
|
||||
route_result = await prompt_router.route_prompt(org_id, json_data)
|
||||
@@ -1100,11 +1128,22 @@ class PromptServer():
|
||||
# every checkpoint-referencing graph ("not in []").
|
||||
if not args.worker_mode:
|
||||
try:
|
||||
from middleware.gpu_dispatch import dispatch_if_worker
|
||||
from middleware.gpu_dispatch import dispatch_if_worker, has_models_for
|
||||
except ImportError: # package-style checkout
|
||||
from .middleware.gpu_dispatch import dispatch_if_worker
|
||||
from .middleware.gpu_dispatch import dispatch_if_worker, has_models_for
|
||||
if dispatch_if_worker(request, org_id, prompt_id, prompt):
|
||||
return web.json_response({"prompt_id": prompt_id, "number": number, "node_errors": {}})
|
||||
# Dispatch didn't happen. If THIS box lacks the models THIS graph
|
||||
# needs (a coordinator pod, even one with a stray SD1.5), local
|
||||
# validation is GUARANTEED to fail with a cryptic "not in []" — the
|
||||
# render belongs on the org's GPU worker, momentarily unreachable.
|
||||
# Return a clear, retryable error instead of the doomed validation.
|
||||
# A box that CAN load the graph falls through and renders it.
|
||||
if not has_models_for(prompt):
|
||||
return web.json_response({"error": {
|
||||
"type": "gpu_worker_unavailable",
|
||||
"message": "GPU render worker is momentarily unavailable — please retry.",
|
||||
"details": "", "extra_info": {}}, "node_errors": {}}, status=503)
|
||||
|
||||
valid = await execution.validate_prompt(prompt_id, prompt, partial_execution_targets)
|
||||
extra_data = {}
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
"""library.json write integrity — the curation-wipe class.
|
||||
|
||||
A fixed shared tmp let two build-manifest sidecars (one per pod during a roll)
|
||||
interleave into one inode -> torn JSON -> readers get {} -> every asset resets to
|
||||
draft. These pin: unique tmp per write, no fixed-tmp residue, concurrent writers
|
||||
never produce a torn/invalid file, and the shared lock file coordinates."""
|
||||
|
||||
import concurrent.futures
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "scripts"))
|
||||
|
||||
import library_manifest as lm
|
||||
|
||||
|
||||
def _doc(n):
|
||||
return {"_meta": {"count": n}, "assets": [{"path": f"a{i}.png", "status": "approved"} for i in range(n)]}
|
||||
|
||||
|
||||
def test_write_atomic_uses_unique_tmp_no_fixed_residue(tmp_path):
|
||||
root = str(tmp_path)
|
||||
lm.write_atomic(root, _doc(3))
|
||||
# the manifest is valid and complete
|
||||
doc = json.loads((tmp_path / "library.json").read_text())
|
||||
assert len(doc["assets"]) == 3
|
||||
# the old fixed shared tmp must never be the write path
|
||||
assert not (tmp_path / ".library.json.tmp").exists()
|
||||
|
||||
|
||||
def test_concurrent_writes_never_tear(tmp_path):
|
||||
root = str(tmp_path)
|
||||
# many writers hammering the same manifest — under the fixed-tmp bug this tore;
|
||||
# with unique tmp + atomic rename every reader sees a COMPLETE valid doc.
|
||||
def w(n):
|
||||
lm.write_atomic(root, _doc(n))
|
||||
return json.loads((tmp_path / "library.json").read_text()) # read must always parse
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as ex:
|
||||
results = list(ex.map(w, [5, 10, 15, 20, 25, 30, 35, 40] * 3))
|
||||
for doc in results:
|
||||
assert "assets" in doc and isinstance(doc["assets"], list) # never {} from a torn read
|
||||
final = json.loads((tmp_path / "library.json").read_text())
|
||||
assert len(final["assets"]) in (5, 10, 15, 20, 25, 30, 35, 40) # one whole writer won, not a splice
|
||||
|
||||
|
||||
def test_library_lock_serializes(tmp_path):
|
||||
root = str(tmp_path)
|
||||
with lm.library_lock(root):
|
||||
assert (tmp_path / lm.LOCK_NAME).exists() # lock file materializes and is held
|
||||
@@ -89,6 +89,15 @@ def test_count_outputs_handles_empty_and_none():
|
||||
assert bm.count_outputs({"outputs": {"1": {"text": "no media here"}}}) == 0
|
||||
|
||||
|
||||
def test_count_outputs_counts_video_extension_agnostic():
|
||||
"""Metering is extension-agnostic: a SaveVideo history entry (its output reported
|
||||
through `images` by ui.PreviewVideo) meters exactly like an image render — the RFC's
|
||||
shared-untouched claim."""
|
||||
hist = {"outputs": {"12": {"images": [
|
||||
{"filename": "video/x_00001_.mp4", "subfolder": "video", "type": "output"}]}}}
|
||||
assert bm.count_outputs(hist) == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# record_render — the fire-and-forget metered event
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -31,3 +31,223 @@ def test_no_online_gpu_when_all_offline(monkeypatch):
|
||||
])
|
||||
assert g._online_gpu_nodes("tok") == []
|
||||
assert g._has_online_gpu("tok") is False
|
||||
|
||||
|
||||
# ── The cloud pod NEVER renders locally ──────────────────────────────────────────
|
||||
|
||||
class _Req:
|
||||
"""Minimal stand-in for an aiohttp request that _bearer reads."""
|
||||
def __init__(self, auth="Bearer u"):
|
||||
self.headers = {"Authorization": auth} if auth else {}
|
||||
self.cookies = {}
|
||||
|
||||
|
||||
class _Resp:
|
||||
def read(self):
|
||||
return b"{}"
|
||||
|
||||
|
||||
def _capture_enqueue(monkeypatch):
|
||||
"""Mock the gpu-jobs POST + input collection; return a dict the enqueue fills."""
|
||||
sent: dict = {}
|
||||
|
||||
def fake_urlopen(req, timeout=0):
|
||||
sent["url"] = req.full_url
|
||||
sent["auth"] = req.headers.get("Authorization")
|
||||
return _Resp()
|
||||
|
||||
monkeypatch.setattr(g.urllib.request, "urlopen", fake_urlopen)
|
||||
monkeypatch.setattr(g, "_collect_input_images", lambda prompt, org: [])
|
||||
monkeypatch.setattr(g, "_track_dispatched",
|
||||
lambda org, pid, prompt, node="gpu": sent.update(node=node))
|
||||
return sent
|
||||
|
||||
|
||||
def test_modelless_pod_enqueues_even_when_org_sees_no_gpu(monkeypatch):
|
||||
"""The karma case: the org-scoped fleet read shows no GPU (workers live under a
|
||||
sibling org), but a MODEL-LESS pod must still enqueue to gpu-jobs and wait — never
|
||||
self-POST the GPU-less pod (the instant 'failed validation')."""
|
||||
monkeypatch.setattr(g, "_online_gpu_nodes", lambda tok: []) # org-blind
|
||||
monkeypatch.setattr(g, "has_models_for", lambda prompt: False) # cloud coordinator
|
||||
sent = _capture_enqueue(monkeypatch)
|
||||
ok = g.dispatch_if_worker(_Req("Bearer karma"), "karma", "pid1", {})
|
||||
assert ok is True # enqueued to the worker lane, NOT local
|
||||
assert "gpu-jobs" in sent["url"]
|
||||
assert sent["auth"] == "Bearer karma" # tenancy still flows from the user token
|
||||
assert sent["node"] == "gpu" # claimer decided at pick-up
|
||||
|
||||
|
||||
def test_modelless_pod_never_returns_local(monkeypatch):
|
||||
"""Even if the enqueue POST fails, a model-less pod returns False so server.py
|
||||
yields a retryable 503 — it must NEVER fall through to local validation."""
|
||||
monkeypatch.setattr(g, "_online_gpu_nodes", lambda tok: [])
|
||||
monkeypatch.setattr(g, "has_models_for", lambda prompt: False)
|
||||
monkeypatch.setattr(g, "_collect_input_images", lambda prompt, org: [])
|
||||
|
||||
def boom(req, timeout=0):
|
||||
raise OSError("gpu-jobs unreachable")
|
||||
|
||||
monkeypatch.setattr(g.urllib.request, "urlopen", boom)
|
||||
# False here means "not enqueued"; server.py then 503s on a model-less pod. The
|
||||
# point: dispatch never claims a local render happened.
|
||||
assert g.dispatch_if_worker(_Req("Bearer karma"), "karma", "pid", {}) is False
|
||||
|
||||
|
||||
def test_box_with_models_renders_local_when_no_worker(monkeypatch):
|
||||
"""A box that HAS its own models (local dev / worker node) still renders locally
|
||||
when it sees no online worker — unchanged behavior."""
|
||||
monkeypatch.setattr(g, "_online_gpu_nodes", lambda tok: [])
|
||||
monkeypatch.setattr(g, "has_models_for", lambda prompt: True)
|
||||
touched = {"urlopen": False}
|
||||
monkeypatch.setattr(g.urllib.request, "urlopen",
|
||||
lambda *a, **k: touched.update(urlopen=True) or _Resp())
|
||||
assert g.dispatch_if_worker(_Req(), "acme", "pid", {}) is False
|
||||
assert touched["urlopen"] is False # never touched the gpu-jobs API
|
||||
|
||||
|
||||
def test_online_worker_enqueues_with_its_label(monkeypatch):
|
||||
monkeypatch.setattr(g, "_online_gpu_nodes",
|
||||
lambda tok: [{"name": "spark", "status": "online", "gpu": True}])
|
||||
monkeypatch.setattr(g, "has_models_for", lambda prompt: True) # an online worker wins
|
||||
sent = _capture_enqueue(monkeypatch)
|
||||
assert g.dispatch_if_worker(_Req(), "acme", "pid", {}) is True
|
||||
assert sent["node"] == "spark"
|
||||
|
||||
|
||||
def test_no_token_is_not_a_local_render_claim(monkeypatch):
|
||||
assert g.dispatch_if_worker(_Req(auth=""), "acme", "pid", {}) is False
|
||||
|
||||
|
||||
# ── The capacity gate is advisory + service-identity ─────────────────────────────
|
||||
|
||||
def test_service_token_prefers_worker_then_commerce(monkeypatch):
|
||||
monkeypatch.setenv("STUDIO_WORKER_TOKEN", "w")
|
||||
monkeypatch.setenv("STUDIO_COMMERCE_TOKEN", "c")
|
||||
assert g._service_token() == "Bearer w"
|
||||
monkeypatch.delenv("STUDIO_WORKER_TOKEN")
|
||||
assert g._service_token() == "Bearer c"
|
||||
monkeypatch.delenv("STUDIO_COMMERCE_TOKEN")
|
||||
assert g._service_token() == ""
|
||||
|
||||
|
||||
def test_advisory_gate_reads_fleet_with_service_identity(monkeypatch):
|
||||
"""karma's user token sees no GPU; the SERVICE token sees the fleet's worker — the
|
||||
gate reports online, so every org's UI reflects true availability."""
|
||||
seen = []
|
||||
|
||||
def fleet(tok):
|
||||
seen.append(tok)
|
||||
return [{"name": "spark", "status": "online", "gpu": True}] if tok == "Bearer svc" else []
|
||||
|
||||
monkeypatch.setattr(g, "_fleet_machines", fleet)
|
||||
monkeypatch.setenv("STUDIO_WORKER_TOKEN", "svc")
|
||||
monkeypatch.delenv("STUDIO_COMMERCE_TOKEN", raising=False)
|
||||
assert g.has_online_gpu_advisory("Bearer karma") is True
|
||||
assert seen[0] == "Bearer svc" # service identity consulted FIRST
|
||||
|
||||
|
||||
def test_advisory_gate_falls_back_to_user_token(monkeypatch):
|
||||
monkeypatch.delenv("STUDIO_WORKER_TOKEN", raising=False)
|
||||
monkeypatch.delenv("STUDIO_COMMERCE_TOKEN", raising=False)
|
||||
monkeypatch.setattr(g, "_fleet_machines",
|
||||
lambda tok: [{"name": "spark", "status": "online", "gpu": True}]
|
||||
if tok == "Bearer user" else [])
|
||||
assert g.has_online_gpu_advisory("Bearer user") is True
|
||||
|
||||
|
||||
def test_advisory_gate_false_when_nothing_online(monkeypatch):
|
||||
monkeypatch.delenv("STUDIO_WORKER_TOKEN", raising=False)
|
||||
monkeypatch.delenv("STUDIO_COMMERCE_TOKEN", raising=False)
|
||||
monkeypatch.setattr(g, "_fleet_machines", lambda tok: [])
|
||||
assert g.has_online_gpu_advisory("Bearer user") is False
|
||||
|
||||
|
||||
# ── Per-GPU targeting: X-Target-GPU → taskQueue "gpu:<identity>" ──────────────────
|
||||
|
||||
class _ReqH(_Req):
|
||||
"""_Req plus an arbitrary header map (so a test can set X-Target-GPU)."""
|
||||
def __init__(self, auth="Bearer u", headers=None):
|
||||
super().__init__(auth)
|
||||
if headers:
|
||||
self.headers.update(headers)
|
||||
|
||||
|
||||
def _capture_body(monkeypatch):
|
||||
"""Capture the enqueue POST body so a test can assert taskQueue + the URL namespace."""
|
||||
import json as _json
|
||||
sent: dict = {}
|
||||
|
||||
def fake_urlopen(req, timeout=0):
|
||||
sent["url"] = req.full_url
|
||||
sent["body"] = _json.loads(req.data.decode()) if req.data else {}
|
||||
return _Resp()
|
||||
|
||||
monkeypatch.setattr(g.urllib.request, "urlopen", fake_urlopen)
|
||||
monkeypatch.setattr(g, "_collect_input_images", lambda prompt, org: [])
|
||||
monkeypatch.setattr(g, "_track_dispatched",
|
||||
lambda org, pid, prompt, node="gpu": sent.update(node=node))
|
||||
return sent
|
||||
|
||||
|
||||
def test_untargeted_uses_shared_gpu_jobs_lane(monkeypatch):
|
||||
"""No X-Target-GPU → the shared 'gpu-jobs' lane any worker drains; namespace unchanged."""
|
||||
monkeypatch.setattr(g, "_online_gpu_nodes", lambda tok: [])
|
||||
monkeypatch.setattr(g, "has_models_for", lambda prompt: False)
|
||||
sent = _capture_body(monkeypatch)
|
||||
assert g.dispatch_if_worker(_Req("Bearer acme"), "acme", "pid", {}) is True
|
||||
assert sent["body"]["taskQueue"] == "gpu-jobs" # shared lane
|
||||
assert "namespaces/gpu-jobs/activities" in sent["url"] # namespace is always gpu-jobs
|
||||
assert sent["node"] == "gpu"
|
||||
|
||||
|
||||
def test_x_target_gpu_routes_to_that_machine_lane(monkeypatch):
|
||||
"""A targeted render — where the target IS one of the caller-org's own online GPUs —
|
||||
enqueues on the machine's OWN lane 'gpu:<identity>' (namespace still gpu-jobs) and is
|
||||
tracked to that node, so only spark claims it. An in-fleet target ALWAYS enqueues,
|
||||
even on a box that could render the graph locally."""
|
||||
monkeypatch.setattr(g, "_online_gpu_nodes",
|
||||
lambda tok: [{"name": "spark", "status": "online", "gpu": True}])
|
||||
monkeypatch.setattr(g, "has_models_for", lambda prompt: True) # would render local…
|
||||
sent = _capture_body(monkeypatch)
|
||||
ok = g.dispatch_if_worker(_ReqH("Bearer acme", {"X-Target-GPU": "spark"}), "acme", "pid", {})
|
||||
assert ok is True # …but an in-fleet target forces the queue
|
||||
assert sent["body"]["taskQueue"] == "gpu:spark" # the machine's own lane
|
||||
assert "namespaces/gpu-jobs/activities" in sent["url"] # namespace unchanged
|
||||
assert sent["node"] == "spark" # tracked to the targeted GPU
|
||||
|
||||
|
||||
def test_out_of_fleet_target_falls_back_to_shared(monkeypatch):
|
||||
"""A target that is NOT one of the caller-org's own online GPUs is ignored — the job
|
||||
rides the shared lane, never a 'gpu:<bogus>' lane no worker claims (a 30-min hang),
|
||||
and an injected label (XSS payload) never becomes a taskQueue / queue-row node."""
|
||||
monkeypatch.setattr(g, "_online_gpu_nodes",
|
||||
lambda tok: [{"name": "spark", "status": "online", "gpu": True}])
|
||||
monkeypatch.setattr(g, "has_models_for", lambda prompt: False)
|
||||
sent = _capture_body(monkeypatch)
|
||||
ok = g.dispatch_if_worker(
|
||||
_ReqH("Bearer acme", {"X-Target-GPU": "<img src=x onerror=alert(1)>"}), "acme", "pid", {})
|
||||
assert ok is True
|
||||
assert sent["body"]["taskQueue"] == "gpu-jobs" # bogus target → shared lane, not gpu:<payload>
|
||||
assert sent["node"] == "spark" # node decided by the real fleet, not the header
|
||||
|
||||
|
||||
def test_blank_target_header_falls_back_to_shared(monkeypatch):
|
||||
"""An empty/whitespace X-Target-GPU is not a target — shared lane."""
|
||||
monkeypatch.setattr(g, "_online_gpu_nodes", lambda tok: [])
|
||||
monkeypatch.setattr(g, "has_models_for", lambda prompt: False)
|
||||
sent = _capture_body(monkeypatch)
|
||||
assert g.dispatch_if_worker(_ReqH("Bearer acme", {"X-Target-GPU": " "}), "acme", "pid", {}) is True
|
||||
assert sent["body"]["taskQueue"] == "gpu-jobs"
|
||||
|
||||
|
||||
# ── A BYO-GPU video job records its REAL prefix, not the "render" fallback ─────────
|
||||
def test_track_savevideo(tmp_path, monkeypatch):
|
||||
import json
|
||||
monkeypatch.setattr(g.folder_paths, "get_org_output_directory", lambda o=None: str(tmp_path))
|
||||
graph = {"7": {"class_type": "LoadImage", "inputs": {"image": "ref.png"}},
|
||||
"12": {"class_type": "SaveVideo", "inputs": {"filename_prefix": "video/beach_ab12cd"}}}
|
||||
g._track_dispatched("acme", "pid-v", graph, node="spark")
|
||||
jobs = json.loads((tmp_path / "render_jobs.json").read_text())
|
||||
assert jobs[-1]["outPrefix"] == "video/beach_ab12cd" # the video prefix, not "render"
|
||||
assert jobs[-1]["prefix"] == "beach_ab12cd"
|
||||
assert jobs[-1]["refs"] == ["ref.png"] and jobs[-1]["node"] == "spark"
|
||||
|
||||
@@ -4,10 +4,14 @@ The JWT path is exercised end to end with a locally generated RSA keypair:
|
||||
sign an access token like Hanzo IAM would, then verify it through the same
|
||||
code the middleware uses. Covers signature, exp, iss, aud and the org claim.
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
from urllib.parse import parse_qs, urlencode, urlsplit
|
||||
|
||||
import jwt
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
from aiohttp.test_utils import make_mocked_request
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
|
||||
from middleware import iam_auth_middleware as m
|
||||
@@ -144,6 +148,14 @@ def test_protected_paths(path):
|
||||
assert m._is_public_path(path) is False
|
||||
|
||||
|
||||
def test_video_never_public():
|
||||
"""A tenant video is only ever served by /v1/library/file?path=… (the URL path never
|
||||
ends in a media suffix), so the suffix bypass cannot apply. ".mp4" must NEVER join the
|
||||
static-suffix allowlist, and the org-scoped file route stays auth-gated — pinned here."""
|
||||
assert m._is_public_path("/a/b.mp4") is False
|
||||
assert m._is_public_path("/v1/library/file") is False
|
||||
|
||||
|
||||
# --- server-owned login entry point (/login → OIDC code flow) ---
|
||||
|
||||
def test_login_path_is_public():
|
||||
@@ -196,3 +208,124 @@ def test_handle_login_returns_to_app_root_not_the_login_path():
|
||||
# State is HMAC-signed with the client_id key (see _sign_state usage).
|
||||
payload = m._verify_state(cookie, "hanzo-studio")
|
||||
assert payload is not None and payload["t"] == "/"
|
||||
|
||||
|
||||
# --- browser-vs-API gate: anon navigation → 302 /login, XHR/JSON → 401 ---
|
||||
|
||||
def _gate(headers):
|
||||
"""Run the auth middleware for an unauthenticated GET /studio with the given
|
||||
request headers; return the response. localhost_bypass is off so the token
|
||||
gate actually runs (prod sits behind the ingress, never loopback)."""
|
||||
mw = m.create_iam_auth_middleware("https://hanzo.id", localhost_bypass=False)
|
||||
|
||||
async def _ok(_request):
|
||||
return web.Response(text="ok")
|
||||
|
||||
req = make_mocked_request("GET", "/studio", headers=headers)
|
||||
return asyncio.run(mw(req, _ok))
|
||||
|
||||
|
||||
def test_anon_browser_navigation_redirects_to_login():
|
||||
"""A top-level browser navigation (Accept: text/html) to a protected page is
|
||||
bounced to the server-owned /login entry point, carrying ?next so the OIDC
|
||||
round trip returns to /studio — never a raw 401."""
|
||||
resp = _gate({"Accept": "text/html,application/xhtml+xml"})
|
||||
assert resp.status == 302
|
||||
loc = resp.headers["Location"]
|
||||
assert loc.startswith("/login?")
|
||||
assert parse_qs(urlsplit(loc).query)["next"] == ["/studio"]
|
||||
|
||||
|
||||
def test_anon_plain_curl_redirects_to_login():
|
||||
"""The confirmed repro: a bare client (Accept: */*, no fetch-metadata) is
|
||||
treated as a navigation and redirected, not 401'd."""
|
||||
resp = _gate({"Accept": "*/*"})
|
||||
assert resp.status == 302
|
||||
assert resp.headers["Location"].startswith("/login?next=")
|
||||
|
||||
|
||||
def test_anon_no_accept_header_redirects_to_login():
|
||||
"""No Accept header at all is still ambiguous, not an API call → redirect."""
|
||||
resp = _gate({})
|
||||
assert resp.status == 302
|
||||
assert "/login" in resp.headers["Location"]
|
||||
|
||||
|
||||
def test_anon_sec_fetch_navigate_redirects():
|
||||
"""Modern browsers stamp Sec-Fetch-Mode: navigate on top-level loads."""
|
||||
resp = _gate({"Accept": "*/*", "Sec-Fetch-Mode": "navigate"})
|
||||
assert resp.status == 302
|
||||
assert "/login" in resp.headers["Location"]
|
||||
|
||||
|
||||
def test_anon_xhr_gets_401():
|
||||
"""An XHR (X-Requested-With) must get a JSON 401 so the SPA can react, not a
|
||||
redirect it would follow into HTML."""
|
||||
resp = _gate({"Accept": "*/*", "X-Requested-With": "XMLHttpRequest"})
|
||||
assert resp.status == 401
|
||||
|
||||
|
||||
def test_anon_json_accept_gets_401():
|
||||
"""An explicit JSON client (Accept: application/json) is an API caller."""
|
||||
resp = _gate({"Accept": "application/json"})
|
||||
assert resp.status == 401
|
||||
|
||||
|
||||
def test_anon_fetch_cors_gets_401():
|
||||
"""A programmatic fetch() carries Sec-Fetch-Mode: cors — 401, not redirect."""
|
||||
resp = _gate({"Accept": "*/*", "Sec-Fetch-Mode": "cors"})
|
||||
assert resp.status == 401
|
||||
|
||||
|
||||
def test_is_api_request_predicate():
|
||||
def api(headers):
|
||||
return m._is_api_request(make_mocked_request("GET", "/studio", headers=headers))
|
||||
|
||||
assert api({"Accept": "*/*"}) is False # plain curl → nav
|
||||
assert api({"Accept": "text/html"}) is False # browser → nav
|
||||
assert api({}) is False # unknown → nav
|
||||
assert api({"Sec-Fetch-Mode": "navigate"}) is False # top-level nav
|
||||
assert api({"Sec-Fetch-Mode": "cors"}) is True # fetch/XHR
|
||||
assert api({"Sec-Fetch-Mode": "same-origin"}) is True # fetch/XHR
|
||||
assert api({"X-Requested-With": "XMLHttpRequest"}) is True # legacy AJAX
|
||||
assert api({"Accept": "application/json"}) is True # JSON client
|
||||
|
||||
|
||||
def test_safe_next_predicate():
|
||||
assert m._safe_next("/studio") == "/studio"
|
||||
assert m._safe_next("/studio?org=acme&run=1") == "/studio?org=acme&run=1"
|
||||
assert m._safe_next(None) == "/"
|
||||
assert m._safe_next("") == "/"
|
||||
assert m._safe_next("//evil.example/x") == "/" # protocol-relative
|
||||
assert m._safe_next("https://evil.example") == "/" # absolute
|
||||
assert m._safe_next("/\\evil.example") == "/" # backslash-relative
|
||||
assert m._safe_next("/login") == "/" # would loop
|
||||
assert m._safe_next("/login?next=/x") == "/" # would loop
|
||||
|
||||
|
||||
def test_handle_login_honors_next_param():
|
||||
"""A redirect from the auth gate carries ?next=/studio; login must return the
|
||||
user there after the OIDC round trip (the signed state proves the target)."""
|
||||
mw = m.create_iam_auth_middleware("https://hanzo.id")
|
||||
req = make_mocked_request(
|
||||
"GET", "/login?" + urlencode({"next": "/studio"}),
|
||||
headers={"X-Forwarded-Proto": "https", "X-Forwarded-Host": "studio.hanzo.ai"},
|
||||
)
|
||||
resp = asyncio.run(mw.handle_login(req))
|
||||
assert resp.status == 302
|
||||
payload = m._verify_state(resp.cookies[m._STATE_COOKIE].value, "hanzo-studio")
|
||||
assert payload["t"] == "/studio"
|
||||
|
||||
|
||||
def test_handle_login_rejects_open_redirect_next():
|
||||
"""An off-origin, protocol-relative or self-referential next is dropped to
|
||||
'/', closing the open-redirect and login-loop vectors."""
|
||||
mw = m.create_iam_auth_middleware("https://hanzo.id")
|
||||
for bad in ("//evil.example/x", "https://evil.example", "/login"):
|
||||
req = make_mocked_request(
|
||||
"GET", "/login?" + urlencode({"next": bad}),
|
||||
headers={"X-Forwarded-Proto": "https", "X-Forwarded-Host": "studio.hanzo.ai"},
|
||||
)
|
||||
resp = asyncio.run(mw.handle_login(req))
|
||||
payload = m._verify_state(resp.cookies[m._STATE_COOKIE].value, "hanzo-studio")
|
||||
assert payload["t"] == "/", f"next={bad!r} should be rejected"
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
"""Quality gate — the verdict must survive REASONING judges.
|
||||
|
||||
zen-vl became a reasoning model: with a starved max_tokens the whole budget went
|
||||
to reasoning_content, content came back null, the strict parse saw nothing and
|
||||
the gate failed open on EVERY render (flagged count reached zero while corrupted
|
||||
outputs sat in the grid). These pin the fixed parsing."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
|
||||
from middleware import quality_gate as qg
|
||||
|
||||
|
||||
def _completion(content=None, reasoning=None):
|
||||
msg = {"role": "assistant", "content": content}
|
||||
if reasoning is not None:
|
||||
msg["reasoning_content"] = reasoning
|
||||
return {"choices": [{"message": msg}]}
|
||||
|
||||
|
||||
def test_reply_reads_plain_content():
|
||||
assert qg._reply_text(_completion(content="FAIL — mottled skin")) == "FAIL — mottled skin"
|
||||
|
||||
|
||||
def test_reply_reads_parts_array():
|
||||
assert qg._reply_text(_completion(content=[{"type": "text", "text": "PASS"}])) == "PASS"
|
||||
|
||||
|
||||
def test_reply_falls_back_to_reasoning_content():
|
||||
got = qg._reply_text(_completion(content=None, reasoning="the skin shows mosaic damage. FAIL"))
|
||||
assert "FAIL" in got
|
||||
|
||||
|
||||
def test_reply_empty_when_nothing():
|
||||
assert qg._reply_text(_completion(content=None)) == ""
|
||||
|
||||
|
||||
def test_verdict_from_reasoning_reply_flags():
|
||||
reply = qg._reply_text(_completion(content=None, reasoning="Looking closely... verdict: FAIL"))
|
||||
assert qg.parse_verdict(reply) == "fail"
|
||||
|
||||
|
||||
def test_judge_payload_not_starved():
|
||||
import inspect
|
||||
src = inspect.getsource(qg.judge)
|
||||
assert '"max_tokens": 8,' not in src
|
||||
|
||||
|
||||
def test_verdict_last_token_wins_over_deliberation():
|
||||
assert qg.parse_verdict("I must answer PASS or FAIL. The skin is mottled: FAIL") == "fail"
|
||||
assert qg.parse_verdict("Could be a FAIL, but the artifacts are compression. PASS") == "pass"
|
||||
assert qg.parse_verdict("no verdict words here") == "inconclusive"
|
||||
@@ -0,0 +1,157 @@
|
||||
"""Quality gate — the VIDEO judge (middleware/quality_gate.py).
|
||||
|
||||
The motion gate samples frames across a clip and asks the MOTION judge; a clear FAIL
|
||||
flags the row exactly like the image gate. A video whose frames can't be sampled is
|
||||
marked visibly UNGATED (gate='skipped' + an 'ungated' work-log event) — NEVER a silent
|
||||
pass. The image path still fails open, byte-for-byte unchanged.
|
||||
"""
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
|
||||
from middleware import quality_gate as qg
|
||||
from middleware import worklog as wl
|
||||
|
||||
av = pytest.importorskip("av")
|
||||
np = pytest.importorskip("numpy")
|
||||
|
||||
|
||||
def _write_mp4(path: Path, frames: int = 6) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
c = av.open(str(path), mode="w")
|
||||
st = c.add_stream("libx264", rate=24)
|
||||
st.width, st.height, st.pix_fmt = 64, 64, "yuv420p"
|
||||
for i in range(frames):
|
||||
arr = np.full((64, 64, 3), (i * 40) % 255, dtype=np.uint8)
|
||||
fr = av.VideoFrame.from_ndarray(arr, format="rgb24")
|
||||
for pkt in st.encode(fr):
|
||||
c.mux(pkt)
|
||||
for pkt in st.encode():
|
||||
c.mux(pkt)
|
||||
c.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def orgdir(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(wl.folder_paths, "get_org_output_directory", lambda o=None: str(tmp_path))
|
||||
return tmp_path
|
||||
|
||||
|
||||
# ── frame sampling ────────────────────────────────────────────────────────────────
|
||||
def test_frames_samples_video(tmp_path):
|
||||
mp4 = tmp_path / "clip.mp4"
|
||||
_write_mp4(mp4)
|
||||
frames = qg._frames(mp4)
|
||||
assert frames is not None and 1 <= len(frames) <= 3
|
||||
assert all(isinstance(b, str) and b for b in frames) # base64 JPEG strings
|
||||
|
||||
# a garbage container yields None (→ ungated, never a silent pass)
|
||||
bad = tmp_path / "bad.mp4"
|
||||
bad.write_bytes(b"\x00\x00\x00 ftyp" + b"garbage")
|
||||
assert qg._frames(bad) is None
|
||||
|
||||
|
||||
# ── the payload carries N frames + the MOTION prompt ──────────────────────────────
|
||||
@pytest.mark.asyncio
|
||||
async def test_video_judge_payload(monkeypatch):
|
||||
seen = {}
|
||||
|
||||
class _Resp:
|
||||
status = 200
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *a):
|
||||
return False
|
||||
|
||||
async def json(self):
|
||||
return {"choices": [{"message": {"content": "PASS"}}]}
|
||||
|
||||
class _Session:
|
||||
def __init__(self, *a, **k):
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *a):
|
||||
return False
|
||||
|
||||
def post(self, url, json=None, headers=None):
|
||||
seen["payload"] = json
|
||||
return _Resp()
|
||||
|
||||
monkeypatch.setattr(qg.aiohttp, "ClientSession", _Session)
|
||||
reply = await qg.judge(["b64one", "b64two"], "Bearer u", qg.MOTION)
|
||||
assert reply == "PASS"
|
||||
content = seen["payload"]["messages"][0]["content"]
|
||||
texts = [c for c in content if c["type"] == "text"]
|
||||
images = [c for c in content if c["type"] == "image_url"]
|
||||
assert len(texts) == 1 and texts[0]["text"] == qg.MOTION
|
||||
assert len(images) == 2 # one part per frame
|
||||
|
||||
|
||||
# ── a clear FAIL flags the row ────────────────────────────────────────────────────
|
||||
@pytest.mark.asyncio
|
||||
async def test_video_fail_flags(orgdir, monkeypatch):
|
||||
root = orgdir
|
||||
_write_mp4(root / "video" / "c.mp4")
|
||||
(root / "library.json").write_text(json.dumps(
|
||||
{"assets": [{"path": "video/c.mp4", "status": "draft"}]}))
|
||||
|
||||
async def fake_judge(images, bearer, prompt):
|
||||
assert prompt == qg.MOTION # the motion contract, not PROMPT
|
||||
return "FAIL — hair morphs between frames"
|
||||
|
||||
monkeypatch.setattr(qg, "judge", fake_judge)
|
||||
await qg.run("acme", root, "video/c.mp4", "")
|
||||
lib = json.loads((root / "library.json").read_text())
|
||||
assert lib["assets"][0]["status"] == "flagged"
|
||||
|
||||
|
||||
# ── an unsamplable video is UNGATED — badge + event, never silent ─────────────────
|
||||
@pytest.mark.asyncio
|
||||
async def test_unreadable_video_is_ungated_not_silent(orgdir, monkeypatch):
|
||||
root = orgdir
|
||||
(root / "video").mkdir(parents=True, exist_ok=True)
|
||||
(root / "video" / "broken.mp4").write_bytes(b"\x00\x00\x00 ftyp" + b"not decodable")
|
||||
(root / "library.json").write_text(json.dumps(
|
||||
{"assets": [{"path": "video/broken.mp4", "status": "draft"}]}))
|
||||
# a work-log row for this output so the 'ungated' event has somewhere to land
|
||||
wl.record("acme", kind="video", prompt="", refs=[], uploads=[], parents=[],
|
||||
output_prefix="video/broken.mp4", pid="up-x", status="done")
|
||||
|
||||
async def never(*a, **k):
|
||||
raise AssertionError("judge must not be called for an unsamplable video")
|
||||
|
||||
monkeypatch.setattr(qg, "judge", never)
|
||||
await qg.run("acme", root, "video/broken.mp4", "")
|
||||
|
||||
lib = json.loads((root / "library.json").read_text())
|
||||
assert lib["assets"][0].get("gate") == "skipped" # the visible badge field
|
||||
assert lib["assets"][0]["status"] == "draft" # status untouched
|
||||
events = [e["s"] for r in wl.load("acme") for e in r.get("events", [])]
|
||||
assert "ungated" in events # the work-log trail, not silent
|
||||
|
||||
|
||||
# ── the image path still fails OPEN, unchanged ────────────────────────────────────
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_still_fails_open(orgdir, monkeypatch):
|
||||
root = orgdir
|
||||
(root / "x.png").write_bytes(b"not a real png") # unreadable → _frames None
|
||||
(root / "library.json").write_text(json.dumps(
|
||||
{"assets": [{"path": "x.png", "status": "draft"}]}))
|
||||
|
||||
async def never(*a, **k):
|
||||
raise AssertionError("judge must not run on an unreadable image")
|
||||
|
||||
monkeypatch.setattr(qg, "judge", never)
|
||||
await qg.run("acme", root, "x.png", "")
|
||||
lib = json.loads((root / "library.json").read_text())
|
||||
assert lib["assets"][0]["status"] == "draft" # fail open: not flagged
|
||||
assert "gate" not in lib["assets"][0] # and NOT marked skipped (images fail open)
|
||||
@@ -0,0 +1,361 @@
|
||||
"""Unit tests for the fix pipeline (middleware/studio_home.py):
|
||||
|
||||
* the fix/compose graphs name only models that exist on the worker AND carry none
|
||||
of the Flux-Kontext reference-method wiring that stippled Qwen edits;
|
||||
* the honesty guard flips an implausibly-fast completion to FAILED;
|
||||
* iteration hygiene falls the fix base back to the nearest clean ancestor.
|
||||
|
||||
Pure/filesystem only — no torch, no GPU, no network.
|
||||
"""
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
|
||||
from middleware import studio_home as sh
|
||||
|
||||
# What the render worker actually has on disk (spark: models/{diffusion_models,
|
||||
# text_encoders,vae}). The graph must never name anything outside this.
|
||||
AVAILABLE = {
|
||||
"diffusion_models": {"qwen-image-edit-2511.safetensors", "flux1-fill-dev.safetensors",
|
||||
"flux-2-klein-4b.safetensors", "wan2.2_ti2v_5B_fp16.safetensors"},
|
||||
"text_encoders": {"qwen2.5-vl-7b.safetensors", "clip_l.safetensors",
|
||||
"t5xxl_fp16.safetensors", "flux2-text-encoder.safetensors"},
|
||||
"vae": {"qwen-image-vae.safetensors", "flux-ae.safetensors", "flux2-vae.safetensors"},
|
||||
}
|
||||
_LOADERS = {"UNETLoader": ("diffusion_models", "unet_name"),
|
||||
"CheckpointLoaderSimple": ("checkpoints", "ckpt_name"),
|
||||
"CLIPLoader": ("text_encoders", "clip_name"),
|
||||
"VAELoader": ("vae", "vae_name")}
|
||||
|
||||
|
||||
def _assert_models_exist(graph):
|
||||
checked = 0
|
||||
for node in graph.values():
|
||||
spec = _LOADERS.get(node.get("class_type"))
|
||||
if not spec:
|
||||
continue
|
||||
folder, key = spec
|
||||
name = node["inputs"][key]
|
||||
assert name in AVAILABLE[folder], f"{name!r} is not on the worker ({folder})"
|
||||
checked += 1
|
||||
return checked
|
||||
|
||||
|
||||
# ── graph construction ───────────────────────────────────────────────────────────
|
||||
|
||||
def test_fix_graph_names_only_existing_models():
|
||||
g = sh._fix_graph("ref.png", "brighten the skin tone", "fixes/x")
|
||||
assert _assert_models_exist(g) == 3 # unet + clip + vae all verified
|
||||
|
||||
|
||||
def test_fix_graph_is_the_canonical_qwen_edit():
|
||||
"""ComfyUI's own Qwen-Image-Edit template: NO Flux-Kontext reference-method
|
||||
override (the stipple cause), and the SAME FluxKontextImageScale'd reference feeds
|
||||
BOTH the edit conditioning and the VAEEncode sampling latent — matched size, so no
|
||||
frame-within-frame border."""
|
||||
g = sh._fix_graph("ref.png", "recolor the top", "fixes/x")
|
||||
types = [n["class_type"] for n in g.values()]
|
||||
assert "FluxKontextMultiReferenceLatentMethod" not in types # the stipple cause
|
||||
assert types.count("TextEncodeQwenImageEditPlus") == 2 # positive + negative
|
||||
ks = next(n for n in g.values() if n["class_type"] == "KSampler")
|
||||
# LOW-denoise edit (default 0.25) of the SOURCE latent — never 1.0 (that is Regenerate).
|
||||
assert ks["inputs"]["denoise"] == sh.EDIT_DENOISE_DEFAULT == 0.25
|
||||
latent = g[ks["inputs"]["latent_image"][0]]
|
||||
assert latent["class_type"] == "VAEEncode"
|
||||
scaled = latent["inputs"]["pixels"] # the sampling latent's source
|
||||
assert g[scaled[0]]["class_type"] == "FluxKontextImageScale" # the scaled reference
|
||||
assert g["6"]["inputs"]["image1"] == scaled # conditioning uses the SAME scaled ref
|
||||
|
||||
|
||||
def test_compose_graph_clean_family_and_models():
|
||||
g = sh._compose_graph(["a.png", "b.png"], "merge the two looks", "composes/x")
|
||||
types = [n["class_type"] for n in g.values()]
|
||||
assert "FluxKontextMultiReferenceLatentMethod" not in types
|
||||
assert _assert_models_exist(g) == 3
|
||||
|
||||
|
||||
# ── the honesty guard ─────────────────────────────────────────────────────────────
|
||||
|
||||
def test_implausibly_fast_only_flags_unsampled_renders():
|
||||
now = 1_000_000
|
||||
r30 = {"ts": now, "params": {"steps": 30}}
|
||||
assert sh._implausibly_fast(r30, now + 1) is True # 1s for 30 steps: impossible
|
||||
assert sh._implausibly_fast(r30, now + 4) is True # <5s: impossible
|
||||
assert sh._implausibly_fast(r30, now + 350) is False # a real 30-step render
|
||||
assert sh._implausibly_fast({"ts": now, "params": {"steps": 4}}, now + 1) is False # exempt
|
||||
assert sh._implausibly_fast({"ts": now, "params": {}}, now + 1) is False # unknown steps
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def org_root(tmp_path, monkeypatch):
|
||||
"""A per-org output dir that is BOTH the library root and the worklog dir, as in
|
||||
the pod layout (orgs/<org>/output)."""
|
||||
def out_dir(o=None):
|
||||
d = tmp_path / "orgs" / (o or "default") / "output"
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return str(d)
|
||||
monkeypatch.setattr(sh.folder_paths, "get_org_output_directory", out_dir)
|
||||
return out_dir
|
||||
|
||||
|
||||
def _land(root: Path, rel: str, mtime: int):
|
||||
p = root / rel
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
p.write_bytes(b"\x89PNG\r\n")
|
||||
os.utime(p, (mtime, mtime))
|
||||
|
||||
|
||||
def test_finalize_landed_done_failed_and_stale(org_root):
|
||||
org = "acme"
|
||||
root = Path(org_root(org))
|
||||
for pid, prefix in (("fast", "fixes/fast"), ("slow", "fixes/slow"), ("stale", "fixes/stale")):
|
||||
sh.worklog.record(org, kind="fix", prompt="p", refs=[], uploads=[], parents=[],
|
||||
output_prefix=prefix, params={"steps": 30}, pid=pid, status="queued")
|
||||
ts = {r["id"]: r["ts"] for r in sh.worklog.load(org)}
|
||||
_land(root, "fixes/fast_00001_.png", ts["fast"] + 1) # landed 1s after dispatch
|
||||
_land(root, "fixes/slow_00001_.png", ts["slow"] + 350) # a real render
|
||||
_land(root, "fixes/stale_00001_.png", ts["stale"] - 100) # predates dispatch: sibling
|
||||
|
||||
rows = sh.worklog.load(org)
|
||||
sh._finalize_landed(org, root, rows, int(time.time()) + 500)
|
||||
|
||||
status = {r["id"]: r["status"] for r in sh.worklog.load(org)}
|
||||
assert status["fast"] == "failed" # honesty guard: engine did not sample
|
||||
assert status["slow"] == "done"
|
||||
assert status["stale"] == "queued" # pre-existing sibling, never delivered as this job
|
||||
failed = next(r for r in sh.worklog.load(org) if r["id"] == "fast")
|
||||
assert any("did not sample" in (e.get("note") or "") for e in failed.get("events", []))
|
||||
|
||||
|
||||
# ── iteration hygiene: clean-ancestor fallback ────────────────────────────────────
|
||||
|
||||
def test_clean_base_falls_back_to_nearest_clean_ancestor(org_root):
|
||||
org = "acme"
|
||||
root = Path(org_root(org))
|
||||
# lineage: root.png (approved) → fix1 (flagged) → fix2 (flagged, the chosen base)
|
||||
for rel in ("root.png", "fixes/root_00001_.png", "fixes/root_00001__00001_.png"):
|
||||
(root / rel).parent.mkdir(parents=True, exist_ok=True)
|
||||
(root / rel).write_bytes(b"\x89PNG\r\n")
|
||||
(root / "library.json").write_text(json.dumps({"assets": [
|
||||
{"path": "root.png", "status": "approved"},
|
||||
{"path": "fixes/root_00001_.png", "status": "flagged"},
|
||||
{"path": "fixes/root_00001__00001_.png", "status": "flagged"},
|
||||
]}))
|
||||
sh.worklog.record(org, kind="fix", prompt="a", refs=[], uploads=[], parents=["root.png"],
|
||||
output_prefix="fixes/root", pid="p1", status="done")
|
||||
sh.worklog.record(org, kind="fix", prompt="b", refs=[], uploads=[],
|
||||
parents=["fixes/root_00001_.png"],
|
||||
output_prefix="fixes/root_00001_", pid="p2", status="done")
|
||||
|
||||
# the flagged fix2 falls back to the approved root (skipping the flagged fix1)
|
||||
base, substituted = sh._clean_base(org, root, "fixes/root_00001__00001_.png")
|
||||
assert base == "root.png" and substituted is True
|
||||
# a clean base is returned unchanged
|
||||
assert sh._clean_base(org, root, "root.png") == ("root.png", False)
|
||||
assert sh._asset_status(root, "fixes/root_00001_.png") == "flagged"
|
||||
|
||||
|
||||
# ── Job identity: every dispatch owns a unique output namespace ──────────────────
|
||||
def test_job_tag_is_short_hex_and_unique():
|
||||
a, b = sh._job_tag(), sh._job_tag()
|
||||
assert len(a) == 6 and int(a, 16) >= 0
|
||||
assert a != b
|
||||
|
||||
|
||||
# ── View/pose-change routing: repose asks use a free latent, anchored fixes keep
|
||||
# the encoded-base latent — and the prompt tail never contradicts the ask ─────
|
||||
def test_fix_graph_anchored_default_keeps_pose():
|
||||
g = sh._fix_graph("x.png", "fix the coloring", "fixes/x")
|
||||
assert g["10"]["class_type"] == "VAEEncode"
|
||||
p = g["6"]["inputs"]["prompt"]
|
||||
assert "pose, framing" in p # preservation directive
|
||||
assert "apply only the requested change" in p.lower()
|
||||
|
||||
|
||||
def test_fix_graph_never_uses_empty_latent():
|
||||
"""Edit is ALWAYS a low-denoise img2img of the source — a pose/view change is now
|
||||
REGENERATE, not Edit, so the fix graph NEVER falls back to EmptyLatentImage
|
||||
(the old repose auto-switch, which drifted the whole image, is gone)."""
|
||||
for instruction in ("fix the coloring", "same outfit but from the side only",
|
||||
"make her standing and looking away"):
|
||||
g = sh._fix_graph("x.png", instruction, "fixes/x")
|
||||
assert not any("Empty" in n.get("class_type", "") for n in g.values()), instruction
|
||||
ks = next(n for n in g.values() if n["class_type"] == "KSampler")["inputs"]
|
||||
assert g[ks["latent_image"][0]]["class_type"] == "VAEEncode"
|
||||
assert sh._assert_edit_graph(g) == ""
|
||||
|
||||
|
||||
def test_repose_lexicon_matches_view_changes_only():
|
||||
hits = ["show me the same model but from the side only", "turn around",
|
||||
"back view please", "make her facing left", "different angle"]
|
||||
misses = ["fix the coloring", "brighten the skin slightly", "remove the second strap"]
|
||||
for t in hits:
|
||||
assert sh._REPOSE.search(t), t
|
||||
for t in misses:
|
||||
assert not sh._REPOSE.search(t), t
|
||||
|
||||
|
||||
# ── Flywheel retrieval: pointed-out mistakes condition the NEXT generation ──────
|
||||
def test_mistakes_selects_corrections_not_praise():
|
||||
ratings = {
|
||||
"a.png": {"vote": -1, "notes": "sleeves way too wide"},
|
||||
"b.png": {"stars": 1, "notes": "color washed out"},
|
||||
"c.png": {"vote": 1, "notes": "love this lighting"}, # praise — excluded
|
||||
"d.png": {"notes": "strap should be thinner"}, # bare note — a correction
|
||||
"e.png": {"vote": -1}, # no text — nothing to inject
|
||||
}
|
||||
out = sh._mistakes(ratings, ["a.png", "b.png", "c.png", "d.png", "e.png", "missing.png"])
|
||||
assert "sleeves way too wide" in out and "color washed out" in out
|
||||
assert "strap should be thinner" in out
|
||||
assert all("love this lighting" != m for m in out)
|
||||
assert len(out) <= 4
|
||||
|
||||
|
||||
def test_fix_graph_injects_lessons_into_both_encoders():
|
||||
g = sh._fix_graph("x.png", "redo the sleeves", "fixes/x",
|
||||
lessons=["sleeves way too wide", "color washed out"])
|
||||
pos, neg = g["6"]["inputs"]["prompt"], g["7"]["inputs"]["prompt"]
|
||||
assert "do not repeat" in pos and "sleeves way too wide" in pos
|
||||
assert "sleeves way too wide" in neg and "color washed out" in neg
|
||||
|
||||
|
||||
def test_compose_graph_injects_lessons():
|
||||
g = sh._compose_graph(["x.png", "y.png"], "combine them", "composes/x",
|
||||
lessons=["strap should be thinner"])
|
||||
assert "strap should be thinner" in g["6"]["inputs"]["prompt"]
|
||||
assert "strap should be thinner" in g["7"]["inputs"]["prompt"]
|
||||
|
||||
|
||||
# ── Ingest upsert identity: a counter-suffixed landed file belongs to its
|
||||
# dispatch row — mirroring must never add a duplicate ingest row ────────────
|
||||
def test_has_output_matches_counter_suffixed_files(tmp_path, monkeypatch):
|
||||
import middleware.worklog as wl
|
||||
monkeypatch.setattr(wl.folder_paths, "get_org_output_directory", lambda org: str(tmp_path))
|
||||
wl.record("acme", kind="fix", prompt="p", refs=[], uploads=[], parents=[],
|
||||
output_prefix="fixes/x_ab12cd")
|
||||
assert wl.has_output("acme", "fixes/x_ab12cd_00001_.png") # the row owns its landed file
|
||||
assert wl.has_output("acme", "fixes/x_ab12cd") # exact prefix still matches
|
||||
assert not wl.has_output("acme", "fixes/other_00001_.png") # unrelated file → new row ok
|
||||
|
||||
|
||||
# ── Job identity for re-run and template-render (the phantom-image class) ────────
|
||||
def test_retag_points_every_saveimage_at_unique_prefix():
|
||||
g = {"1": {"class_type": "KSampler", "inputs": {"seed": 1}},
|
||||
"9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "fixes/original"}},
|
||||
"10": {"class_type": "SaveImage", "inputs": {"filename_prefix": "fixes/original"}}}
|
||||
sh._retag(g, "reruns/original_ab12cd")
|
||||
assert g["9"]["inputs"]["filename_prefix"] == "reruns/original_ab12cd"
|
||||
assert g["10"]["inputs"]["filename_prefix"] == "reruns/original_ab12cd"
|
||||
# a re-run graph never keeps the source's SaveImage prefix (that was the collision)
|
||||
assert all(n.get("inputs", {}).get("filename_prefix") != "fixes/original"
|
||||
for n in g.values() if n.get("class_type") == "SaveImage")
|
||||
|
||||
|
||||
# ── State isolation & job construction ────────────────────────────────────────────
|
||||
# An Edit is built from the SELECTED source alone; two concurrent edits of DIFFERENT
|
||||
# sources can never share a staged file, a LoadImage ref, a positive prompt, a
|
||||
# sampling latent or a parent. Content-addressed staging (name == sha256 of the
|
||||
# bytes) is the binding proof: the worker's LoadImage resolves the exact selected
|
||||
# image, and different sources are cryptographically un-collidable on a shared
|
||||
# worker input dir — the "unrelated random picture / wrong image" class of bug.
|
||||
from middleware import gpu_dispatch as gd
|
||||
|
||||
|
||||
def _pos_prompt(g):
|
||||
"""The positive prompt actually wired into KSampler (not read by node id)."""
|
||||
ks = next(n for n in g.values() if n["class_type"] == "KSampler")
|
||||
return g[ks["inputs"]["positive"][0]]["inputs"]["prompt"]
|
||||
|
||||
|
||||
def _load_ref(g):
|
||||
return next(n for n in g.values() if n["class_type"] == "LoadImage")["inputs"]["image"]
|
||||
|
||||
|
||||
def _save_prefix(g):
|
||||
return next(n for n in g.values() if n["class_type"] in sh._SAVERS)["inputs"]["filename_prefix"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def in_dir(tmp_path, monkeypatch):
|
||||
"""Per-org input dir (orgs/<org>/input) where sources are staged for the worker —
|
||||
patched on BOTH modules that touch it: studio_home stages, gpu_dispatch inlines."""
|
||||
def _d(o=None):
|
||||
d = tmp_path / "orgs" / (o or "default") / "input"
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return str(d)
|
||||
monkeypatch.setattr(sh.folder_paths, "get_org_input_directory", _d)
|
||||
monkeypatch.setattr(gd.folder_paths, "get_org_input_directory", _d)
|
||||
return _d
|
||||
|
||||
|
||||
def test_stage_source_is_content_addressed(in_dir, tmp_path):
|
||||
d = Path(in_dir("acme"))
|
||||
a = tmp_path / "a.png"; a.write_bytes(b"AAAA-dress-bytes")
|
||||
b = tmp_path / "b.png"; b.write_bytes(b"BBBB-clasp-bytes")
|
||||
na, nb = sh._stage_source(d, a, "fix"), sh._stage_source(d, b, "fix")
|
||||
# different bytes -> different name (cryptographic, not a salted-hash coincidence)
|
||||
assert na != nb
|
||||
# the name IS the sha256 of the bytes: a verifiable binding to the EXACT source
|
||||
assert na == f"fix_src_{hashlib.sha256(a.read_bytes()).hexdigest()[:16]}.png"
|
||||
assert (d / na).read_bytes() == a.read_bytes() # staged bytes are the source
|
||||
assert sh._stage_source(d, a, "fix") == na # idempotent: same source, same name
|
||||
a.write_bytes(b"AAAA-dress-bytez") # one byte changes the address
|
||||
assert sh._stage_source(d, a, "fix") != na
|
||||
|
||||
|
||||
def test_two_concurrent_edits_are_isolated(in_dir, tmp_path):
|
||||
"""The reported scenario: Job A (design A, 'make the dress red') and Job B
|
||||
(design B, 'add a silver clasp') constructed back-to-back share NOTHING."""
|
||||
d = Path(in_dir("acme"))
|
||||
srcA = tmp_path / "A.png"; srcA.write_bytes(b"\x89PNG-design-A" + b"\x00" * 64)
|
||||
srcB = tmp_path / "B.png"; srcB.write_bytes(b"\x89PNG-design-B" + b"\x00" * 64)
|
||||
refA, refB = sh._stage_source(d, srcA, "fix"), sh._stage_source(d, srcB, "fix")
|
||||
gA = sh._fix_graph(refA, "make the dress red", "fixes/A_" + sh._job_tag())
|
||||
gB = sh._fix_graph(refB, "add a silver clasp", "fixes/B_" + sh._job_tag())
|
||||
# 1) distinct staged sources, each graph LoadImages its OWN source
|
||||
assert refA != refB
|
||||
assert _load_ref(gA) == refA and _load_ref(gB) == refB
|
||||
# 2) prompts do not bleed across jobs
|
||||
assert "make the dress red" in _pos_prompt(gA) and "clasp" not in _pos_prompt(gA)
|
||||
assert "add a silver clasp" in _pos_prompt(gB) and "dress red" not in _pos_prompt(gB)
|
||||
# 3) both are true low-denoise edits of their own source — no text-to-image fallback
|
||||
assert sh._assert_edit_graph(gA) == "" and sh._assert_edit_graph(gB) == ""
|
||||
for g in (gA, gB):
|
||||
ks = next(n for n in g.values() if n["class_type"] == "KSampler")
|
||||
assert g[ks["inputs"]["latent_image"][0]]["class_type"] == "VAEEncode"
|
||||
# 4) distinct output namespaces (no phantom-image cross-claim)
|
||||
assert _save_prefix(gA) != _save_prefix(gB)
|
||||
|
||||
|
||||
def test_collect_input_images_ships_exact_source_and_is_org_scoped(in_dir, tmp_path):
|
||||
"""gpu_dispatch inlines the bytes the worker will LoadImage. It must ship the
|
||||
SELECTED org's copy exactly — never a same-named file from another org."""
|
||||
dA, dB = Path(in_dir("acme")), Path(in_dir("other"))
|
||||
src = tmp_path / "s.png"; src.write_bytes(b"\x89PNG-real-source" + b"\x01" * 40)
|
||||
ref = sh._stage_source(dA, src, "fix")
|
||||
(dB / ref).write_bytes(b"\x89PNG-WRONG-org" + b"\x02" * 40) # decoy, same basename
|
||||
shipped = gd._collect_input_images(sh._fix_graph(ref, "brighten", "fixes/x"), "acme")
|
||||
got = {s["name"]: base64.b64decode(s["data"]) for s in shipped}
|
||||
assert ref in got and got[ref] == src.read_bytes() # acme's exact bytes
|
||||
assert b"WRONG-org" not in got[ref]
|
||||
|
||||
|
||||
def test_assert_edit_graph_rejects_text_to_image_fallback():
|
||||
"""The hard invariant: an Edit can NEVER be silently served by a text-to-image
|
||||
(EmptyLatentImage) graph. If source resolution ever produced one, the assert
|
||||
REJECTS it — dispatch raises, never renders an unrelated picture."""
|
||||
t2i = {
|
||||
"1": {"class_type": "EmptySD3LatentImage", "inputs": {"width": 1024, "height": 1024}},
|
||||
"2": {"class_type": "KSampler", "inputs": {"latent_image": ["1", 0], "denoise": 1.0}},
|
||||
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0], "filename_prefix": "x"}},
|
||||
}
|
||||
reason = sh._assert_edit_graph(t2i)
|
||||
assert reason and "EmptyLatentImage" in reason # rejected with a clear cause
|
||||
@@ -0,0 +1,129 @@
|
||||
"""Studio 3D lane — Hunyuan3D-2.1 image-to-mesh (.glb) over the shared library/queue/
|
||||
flywheel. Mirrors the video-lane tests: the proven graph builder, the dispatch route
|
||||
(org from the token, unique job tag, one required reference), and .glb ingest/serve."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
from middleware import studio_home as sh
|
||||
from middleware import worklog as wl
|
||||
|
||||
|
||||
class _Req:
|
||||
def __init__(self, *, body=None, query=None, org="acme"):
|
||||
self._body = body
|
||||
self.query = query or {}
|
||||
self.headers = {}
|
||||
self._user = {"org_id": org, "sub": f"u@{org}"} if org else None
|
||||
self.cookies = {}
|
||||
self.content_length = None
|
||||
def get(self, key, default=None):
|
||||
return self._user if key == "iam_user" else default
|
||||
async def json(self):
|
||||
return self._body if self._body is not None else {}
|
||||
|
||||
|
||||
class _FakeQueue:
|
||||
def get_tasks_remaining(self):
|
||||
return 0
|
||||
|
||||
|
||||
class _FakeServer:
|
||||
port = 1
|
||||
def __init__(self):
|
||||
self.prompt_queue = _FakeQueue()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def orgaware(tmp_path, monkeypatch):
|
||||
def outd(o="default"):
|
||||
d = tmp_path / "orgs" / o / "output"
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return str(d)
|
||||
def ind(o="default"):
|
||||
d = tmp_path / "orgs" / o / "input"
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return str(d)
|
||||
monkeypatch.setattr(sh.folder_paths, "get_org_output_directory", outd)
|
||||
monkeypatch.setattr(sh.folder_paths, "get_org_input_directory", ind)
|
||||
monkeypatch.setattr(sh.folder_paths, "get_output_directory", lambda: str(tmp_path))
|
||||
monkeypatch.setattr(wl.folder_paths, "get_org_output_directory", outd)
|
||||
return outd, ind
|
||||
|
||||
|
||||
def _handler(server, method, path, monkeypatch):
|
||||
monkeypatch.setattr(sh.stacks_store, "add_stacks_routes", lambda *a, **k: None)
|
||||
routes = web.RouteTableDef()
|
||||
sh.add_studio_home_routes(routes, server)
|
||||
for r in routes:
|
||||
if getattr(r, "method", None) == method and getattr(r, "path", None) == path:
|
||||
return r.handler
|
||||
raise AssertionError(f"no route {method} {path}")
|
||||
|
||||
|
||||
# ── 1. the proven Hunyuan3D graph ──────────────────────────────────────────────────
|
||||
def test_model3d_graph_is_the_proven_hunyuan3d():
|
||||
g = sh._model3d_graph("ref.png", "model3d/x_ab12cd", seed=7)
|
||||
cls = {k: v["class_type"] for k, v in g.items()}
|
||||
# the ONE bundled checkpoint (not separate loaders — that produced noise)
|
||||
assert g["1"]["class_type"] == "ImageOnlyCheckpointLoader"
|
||||
assert g["1"]["inputs"]["ckpt_name"] == sh._HY3D_CKPT
|
||||
# the flow-matching wrapper WITHOUT which the sampler yields fragments
|
||||
assert "ModelSamplingAuraFlow" in cls.values()
|
||||
assert g["5"]["inputs"]["shift"] == 1.0 and g["5"]["inputs"]["model"] == ["1", 0]
|
||||
# image → clip-vision → conditioning; latent is EmptyLatentHunyuan3Dv2(4096)
|
||||
assert g["2"]["inputs"]["image"] == "ref.png"
|
||||
assert g["6"]["inputs"]["resolution"] == 4096
|
||||
assert g["7"]["inputs"]["seed"] == 7 and g["7"]["inputs"]["cfg"] == 5.0
|
||||
# ends at SaveGLB with the unique prefix
|
||||
assert g["10"]["class_type"] == "SaveGLB"
|
||||
assert g["10"]["inputs"]["filename_prefix"] == "model3d/x_ab12cd"
|
||||
assert "VAEDecodeHunyuan3D" in cls.values() and "VoxelToMesh" in cls.values()
|
||||
|
||||
|
||||
# ── 2. dispatch + route (org from token, unique tag, one required reference) ────────
|
||||
@pytest.mark.asyncio
|
||||
async def test_dispatch_model3d_route(orgaware, monkeypatch):
|
||||
outd, _ = orgaware
|
||||
root = Path(outd("acme"))
|
||||
root.joinpath("library.json").write_text('{"assets":[]}')
|
||||
(root / "designs").mkdir(parents=True, exist_ok=True)
|
||||
from PIL import Image
|
||||
Image.new("RGB", (8, 8)).save(root / "designs" / "hero.png", "PNG")
|
||||
captured = {}
|
||||
async def fake_queue(request, server, graph):
|
||||
captured["graph"] = graph
|
||||
return "pid-3d"
|
||||
monkeypatch.setattr(sh, "_queue_prompt", fake_queue)
|
||||
model3d = _handler(_FakeServer(), "POST", "/v1/library/model3d", monkeypatch)
|
||||
|
||||
# a library image → 3D; 'org' in the body is ignored (token org wins)
|
||||
resp = await model3d(_Req(body={"path": "designs/hero.png", "org": "evil"}, org="acme"))
|
||||
assert resp.status == 200
|
||||
import re
|
||||
out = json.loads(resp.text)
|
||||
assert re.match(r"^model3d/.+_[0-9a-f]{6}$", out["output_prefix"]) # unique job tag
|
||||
rows = wl.load("acme")
|
||||
assert len(rows) == 1 and rows[0]["kind"] == "model3d"
|
||||
assert wl.load("evil") == []
|
||||
assert any(n.get("class_type") == "SaveGLB" for n in captured["graph"].values())
|
||||
|
||||
# NO reference → 400 (there is no text-to-3D)
|
||||
r2 = await model3d(_Req(body={}, org="acme"))
|
||||
assert r2.status == 400
|
||||
# both a path and an upload → 400
|
||||
r3 = await model3d(_Req(body={"path": "a.png", "upload": "u.png"}, org="acme"))
|
||||
assert r3.status == 400
|
||||
|
||||
|
||||
# ── 3. .glb ingest + serve ─────────────────────────────────────────────────────────
|
||||
def test_glb_is_recognized_as_model3d():
|
||||
assert ".glb" in sh._UPLOAD_EXT
|
||||
assert sh._CTYPE[".glb"] == "model/gltf-binary"
|
||||
assert ".glb" in sh._MODEL3D_EXT and ".glb" in sh._SAVE_EXT
|
||||
@@ -0,0 +1,493 @@
|
||||
"""Unit tests for the Studio video lane (middleware/studio_home.py):
|
||||
|
||||
* the Wan 2.2 TI2V-5B graph builder (t2v omits the reference node; clamps + snaps);
|
||||
* dispatch + the /v1/library/video route (org from the token, never the body);
|
||||
* landed-output + queue detection across .mp4/.webm (the honesty guard covers video);
|
||||
* the ONE indexed ingest door accepts video, defaults kind, judges, colour-matches;
|
||||
* the manifest sweeper keeps a video as kind="video";
|
||||
* serving + the first-frame poster (404 on decode failure, never the raw mp4);
|
||||
* embedded-metadata read parity (Versions / Re-run / provenance) and job re-tagging.
|
||||
|
||||
Filesystem + av only — no torch, no GPU, no network. av/numpy are hard studio deps.
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from middleware import studio_home as sh
|
||||
from middleware import worklog as wl
|
||||
|
||||
av = pytest.importorskip("av")
|
||||
np = pytest.importorskip("numpy")
|
||||
|
||||
|
||||
# ── media fixtures: real, tiny, in-test-encoded ───────────────────────────────────
|
||||
def _write_mp4(path: Path, frames: int = 6, meta: dict | None = None) -> None:
|
||||
"""A tiny real mp4 (ftyp at offset 4). ``meta`` embeds container metadata the way
|
||||
SaveVideo does — via movflags=use_metadata_tags, so av reads it back."""
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
opts = {"movflags": "use_metadata_tags"} if meta else {}
|
||||
c = av.open(str(path), mode="w", options=opts)
|
||||
for k, v in (meta or {}).items():
|
||||
c.metadata[k] = v
|
||||
st = c.add_stream("libx264", rate=24)
|
||||
st.width, st.height, st.pix_fmt = 64, 64, "yuv420p"
|
||||
for i in range(frames):
|
||||
arr = np.full((64, 64, 3), (i * 40) % 255, dtype=np.uint8)
|
||||
fr = av.VideoFrame.from_ndarray(arr, format="rgb24")
|
||||
for pkt in st.encode(fr):
|
||||
c.mux(pkt)
|
||||
for pkt in st.encode():
|
||||
c.mux(pkt)
|
||||
c.close()
|
||||
|
||||
|
||||
def _write_webm(path: Path, frames: int = 6) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
c = av.open(str(path), mode="w")
|
||||
st = c.add_stream("libvpx-vp9", rate=24)
|
||||
st.width, st.height, st.pix_fmt = 64, 64, "yuv420p"
|
||||
for i in range(frames):
|
||||
arr = np.full((64, 64, 3), (i * 40) % 255, dtype=np.uint8)
|
||||
fr = av.VideoFrame.from_ndarray(arr, format="rgb24")
|
||||
for pkt in st.encode(fr):
|
||||
c.mux(pkt)
|
||||
for pkt in st.encode():
|
||||
c.mux(pkt)
|
||||
c.close()
|
||||
|
||||
|
||||
# ── request + server doubles ──────────────────────────────────────────────────────
|
||||
class _Req:
|
||||
"""A minimal stand-in for an aiohttp request the route handlers read: JSON body,
|
||||
query, headers, cookies, and the iam_user the middleware would have attached."""
|
||||
def __init__(self, *, body=None, raw=b"", query=None, headers=None,
|
||||
org="acme", cookies=None, content_length=None):
|
||||
self._body = body
|
||||
self._raw = raw
|
||||
self.query = query or {}
|
||||
self.headers = headers or {}
|
||||
self._user = {"org_id": org, "sub": f"u@{org}"} if org else None
|
||||
self.cookies = cookies or {}
|
||||
self.content_length = content_length if content_length is not None else (len(raw) or None)
|
||||
|
||||
def get(self, key, default=None):
|
||||
return self._user if key == "iam_user" else default
|
||||
|
||||
async def json(self):
|
||||
return self._body if self._body is not None else {}
|
||||
|
||||
async def read(self):
|
||||
return self._raw
|
||||
|
||||
|
||||
class _FakeQueue:
|
||||
def get_current_queue_volatile(self):
|
||||
return ([], [])
|
||||
|
||||
|
||||
class _FakeServer:
|
||||
port = 1
|
||||
|
||||
def __init__(self):
|
||||
self.prompt_queue = _FakeQueue()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def orgaware(tmp_path, monkeypatch):
|
||||
def outd(o=None):
|
||||
d = tmp_path / "orgs" / (o or "default") / "output"
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return str(d)
|
||||
|
||||
def ind(o=None):
|
||||
d = tmp_path / "orgs" / (o or "default") / "input"
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return str(d)
|
||||
|
||||
monkeypatch.setattr(sh.folder_paths, "get_org_output_directory", outd)
|
||||
monkeypatch.setattr(sh.folder_paths, "get_org_input_directory", ind)
|
||||
monkeypatch.setattr(sh.folder_paths, "get_output_directory", lambda: str(tmp_path))
|
||||
monkeypatch.setattr(wl.folder_paths, "get_org_output_directory", outd)
|
||||
return outd, ind
|
||||
|
||||
|
||||
def _handler(server, method, path, monkeypatch):
|
||||
"""Extract ONE registered route handler, without the stacks sub-router or a live app."""
|
||||
monkeypatch.setattr(sh.stacks_store, "add_stacks_routes", lambda *a, **k: None)
|
||||
routes = web.RouteTableDef()
|
||||
sh.add_studio_home_routes(routes, server)
|
||||
for r in routes:
|
||||
if getattr(r, "method", None) == method and getattr(r, "path", None) == path:
|
||||
return r.handler
|
||||
raise AssertionError(f"no route {method} {path}")
|
||||
|
||||
|
||||
def _lib(root: Path, assets):
|
||||
(root / "library.json").write_text(json.dumps({"assets": assets}))
|
||||
|
||||
|
||||
# ── 1. the graph builder ──────────────────────────────────────────────────────────
|
||||
def test_video_graph():
|
||||
# text-to-video: NO LoadImage node, and NO start_image key on the latent
|
||||
t2v = sh._video_graph("a woman walks on a sunlit beach", "video/x", length=73)
|
||||
assert "7" not in t2v
|
||||
assert "start_image" not in t2v["8"]["inputs"]
|
||||
assert t2v["8"]["inputs"]["length"] == 73 # 4n+1, unchanged
|
||||
|
||||
# image-to-video: LoadImage → start_image; clamp/snap of the dims + length
|
||||
g = sh._video_graph("beach scene", "video/y_ab12cd", ref_name="ref.png",
|
||||
width=650, height=641, length=72)
|
||||
assert g["7"]["class_type"] == "LoadImage" and g["7"]["inputs"]["image"] == "ref.png"
|
||||
assert g["8"]["inputs"]["start_image"] == ["7", 0]
|
||||
assert (g["8"]["inputs"]["width"], g["8"]["inputs"]["height"]) == (640, 640) # step of 32
|
||||
assert g["8"]["inputs"]["length"] == 69 # 72 → nearest 4n+1 below
|
||||
|
||||
ks = g["9"]["inputs"]
|
||||
assert ks["steps"] == 20 and ks["cfg"] == 5.0
|
||||
assert ks["sampler_name"] == "uni_pc" and ks["scheduler"] == "simple"
|
||||
assert g["4"]["class_type"] == "ModelSamplingSD3" and g["4"]["inputs"]["shift"] == 8
|
||||
sv = g["12"]
|
||||
assert sv["class_type"] == "SaveVideo"
|
||||
assert sv["inputs"] == {"video": ["11", 0], "filename_prefix": "video/y_ab12cd",
|
||||
"format": "mp4", "codec": "h264"}
|
||||
assert g["6"]["inputs"]["text"] == sh._WAN_NEG
|
||||
assert sh._WAN_NEG.startswith("色调艳丽") and "倒着走" in sh._WAN_NEG
|
||||
|
||||
# lessons condition BOTH encoders, exactly like _fix_graph
|
||||
gl = sh._video_graph("beach", "video/z", lessons=["hair morphs between frames"])
|
||||
assert "do not repeat" in gl["5"]["inputs"]["text"]
|
||||
assert "hair morphs between frames" in gl["5"]["inputs"]["text"]
|
||||
assert "hair morphs between frames" in gl["6"]["inputs"]["text"]
|
||||
|
||||
|
||||
# ── 2. dispatch + the route (org from the token, not the body) ─────────────────────
|
||||
@pytest.mark.asyncio
|
||||
async def test_dispatch_route(orgaware, monkeypatch):
|
||||
outd, _ = orgaware
|
||||
captured = {}
|
||||
|
||||
async def fake_queue(request, server, graph):
|
||||
captured["graph"] = graph
|
||||
return "pid-v"
|
||||
|
||||
monkeypatch.setattr(sh, "_queue_prompt", fake_queue)
|
||||
video = _handler(_FakeServer(), "POST", "/v1/library/video", monkeypatch)
|
||||
|
||||
# an 'org' field in the body is IGNORED — the row lands under the TOKEN org
|
||||
resp = await video(_Req(body={"prompt": "a woman walking on a beach at sunset",
|
||||
"org": "evil"}, org="acme"))
|
||||
assert resp.status == 200
|
||||
out = json.loads(resp.text)
|
||||
import re
|
||||
assert re.match(r"^video/.+_[0-9a-f]{6}$", out["output_prefix"])
|
||||
rows = wl.load("acme")
|
||||
assert len(rows) == 1 and rows[0]["kind"] == "video"
|
||||
assert rows[0]["lane"] in ("local", "gpu")
|
||||
assert wl.load("evil") == [] # never under the body's org
|
||||
# the graph that queued is the Wan graph (t2v: no LoadImage)
|
||||
assert any(n.get("class_type") == "SaveVideo" for n in captured["graph"].values())
|
||||
assert not any(n.get("class_type") == "LoadImage" for n in captured["graph"].values())
|
||||
|
||||
# prompt < 3 chars → 400
|
||||
r2 = await video(_Req(body={"prompt": "hi"}, org="acme"))
|
||||
assert r2.status == 400
|
||||
|
||||
# two references (a library path AND an upload) → 400
|
||||
r3 = await video(_Req(body={"prompt": "a valid video prompt", "path": "a.png",
|
||||
"upload": "u.png"}, org="acme"))
|
||||
assert r3.status == 400
|
||||
|
||||
|
||||
# ── 3. landed-output + the honesty guard, over video ──────────────────────────────
|
||||
def test_landed_video(orgaware):
|
||||
outd, _ = orgaware
|
||||
root = Path(outd("acme"))
|
||||
(root / "video").mkdir(parents=True, exist_ok=True)
|
||||
(root / "video" / "x_ab12cd_00001_.mp4").write_bytes(b"\x00\x00\x00 ftyp")
|
||||
assert sh._landed_output(root, "video/x_ab12cd") == "video/x_ab12cd_00001_.mp4"
|
||||
(root / "video" / "w_00001_.webm").write_bytes(b"\x1a\x45\xdf\xa3")
|
||||
assert sh._landed_output(root, "video/w") == "video/w_00001_.webm"
|
||||
|
||||
# child-prefix must NOT be swallowed: "a" never claims "a_00001__00001_.mp4"
|
||||
(root / "video" / "a_00001__00001_.mp4").write_bytes(b"\x00\x00\x00 ftyp")
|
||||
assert sh._landed_output(root, "video/a") is None
|
||||
|
||||
# png unchanged
|
||||
(root / "fixes").mkdir(parents=True, exist_ok=True)
|
||||
(root / "fixes" / "y_00001_.png").write_bytes(b"\x89PNG\r\n")
|
||||
assert sh._landed_output(root, "fixes/y") == "fixes/y_00001_.png"
|
||||
|
||||
# _finalize_landed marks the video row done at the file mtime …
|
||||
wl.record("acme", kind="video", prompt="p", refs=[], uploads=[], parents=[],
|
||||
output_prefix="video/done", params={"steps": 20}, pid="dv", status="queued")
|
||||
ts = next(r["ts"] for r in wl.load("acme") if r["id"] == "dv")
|
||||
dp = root / "video" / "done_00001_.mp4"
|
||||
dp.write_bytes(b"\x00\x00\x00 ftyp")
|
||||
os.utime(dp, (ts + 200, ts + 200))
|
||||
rows = wl.load("acme")
|
||||
sh._finalize_landed("acme", root, rows, int(time.time()) + 500)
|
||||
assert next(r for r in wl.load("acme") if r["id"] == "dv")["status"] == "done"
|
||||
|
||||
# … but a <5s landing for a 20-step render is impossible → failed
|
||||
wl.record("acme", kind="video", prompt="p", refs=[], uploads=[], parents=[],
|
||||
output_prefix="video/fast", params={"steps": 20}, pid="fv", status="queued")
|
||||
tsf = next(r["ts"] for r in wl.load("acme") if r["id"] == "fv")
|
||||
fp = root / "video" / "fast_00001_.mp4"
|
||||
fp.write_bytes(b"\x00\x00\x00 ftyp")
|
||||
os.utime(fp, (tsf + 2, tsf + 2))
|
||||
rows = wl.load("acme")
|
||||
sh._finalize_landed("acme", root, rows, int(time.time()) + 500)
|
||||
failed = next(r for r in wl.load("acme") if r["id"] == "fv")
|
||||
assert failed["status"] == "failed"
|
||||
assert any("did not sample" in (e.get("note") or "") for e in failed.get("events", []))
|
||||
|
||||
|
||||
# ── 4. the live BYO-GPU queue drops a video job when its mp4 lands ─────────────────
|
||||
@pytest.mark.asyncio
|
||||
async def test_render_queue_video(orgaware, monkeypatch):
|
||||
"""The render-queue is AUTHORITATIVE: it reads the org's gpu-jobs activities from the
|
||||
tasks engine and shows the REAL claiming GPU (activity.identity), dropping terminal
|
||||
jobs. A STARTED studio.render is live; a COMPLETED one falls off."""
|
||||
outd, _ = orgaware
|
||||
root = Path(outd("acme"))
|
||||
_lib(root, [])
|
||||
render_queue = _handler(_FakeServer(), "GET", "/v1/render-queue", monkeypatch)
|
||||
base = {
|
||||
"execution": {"runId": "p", "workflowId": "p"},
|
||||
"type": {"name": "studio.render"},
|
||||
"taskQueue": "gpu:spark", "identity": "spark",
|
||||
"input": {"prompt": {"12": {"class_type": "SaveVideo",
|
||||
"inputs": {"filename_prefix": "video/x"}}}},
|
||||
}
|
||||
|
||||
async def _started(_req):
|
||||
return [{**base, "status": "STARTED"}]
|
||||
monkeypatch.setattr(sh, "_list_gpu_jobs", _started)
|
||||
resp = await render_queue(_Req(org="acme"))
|
||||
jobs = json.loads(resp.text)["jobs"]
|
||||
assert len(jobs) == 1 and jobs[0]["node"] == "spark" # real claimer shown
|
||||
|
||||
async def _completed(_req):
|
||||
return [{**base, "status": "COMPLETED"}]
|
||||
monkeypatch.setattr(sh, "_list_gpu_jobs", _completed)
|
||||
resp = await render_queue(_Req(org="acme"))
|
||||
assert json.loads(resp.text)["jobs"] == [] # engine says done → dropped
|
||||
|
||||
|
||||
# ── 5. worklog params + per-kind ETA bucketing ────────────────────────────────────
|
||||
def test_graph_params_video(orgaware):
|
||||
outd, _ = orgaware
|
||||
g = sh._video_graph("a beach", "video/x", seed=7, length=73)
|
||||
p = sh._graph_params(g)
|
||||
for k in ("steps", "cfg", "sampler", "scheduler", "seed", "width", "height", "length", "fps"):
|
||||
assert k in p, k
|
||||
assert p["fps"] == 24 and p["length"] == 73 and p["seed"] == 7
|
||||
|
||||
wl.record("acme", kind="video", prompt="", refs=[], uploads=[], parents=[],
|
||||
output_prefix="video/a", pid="v1", status="queued")
|
||||
wl.mark_started("acme", "v1", 1000)
|
||||
wl.mark_finished("acme", "v1", 1120)
|
||||
wl.record("acme", kind="fix", prompt="", refs=[], uploads=[], parents=[],
|
||||
output_prefix="fixes/a", pid="f1", status="queued")
|
||||
wl.mark_started("acme", "f1", 1000)
|
||||
wl.mark_finished("acme", "f1", 1030)
|
||||
medians = wl.median_durations(wl.load("acme"))
|
||||
assert medians.get("video") == 120 and medians.get("fix") == 30 # bucketed apart
|
||||
|
||||
|
||||
# ── 6. the ONE indexed ingest door accepts video ──────────────────────────────────
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_video(orgaware, tmp_path, monkeypatch):
|
||||
outd, _ = orgaware
|
||||
root = Path(outd("acme"))
|
||||
_lib(root, [])
|
||||
scheduled = {}
|
||||
cm = {}
|
||||
monkeypatch.setattr(sh.quality_gate, "schedule",
|
||||
lambda request, org, r, rel: scheduled.update(rel=rel, org=org))
|
||||
monkeypatch.setattr(sh, "colormatch_fix_ingest",
|
||||
lambda request, sub, name: cm.update(sub=sub, name=name))
|
||||
monkeypatch.setattr(sh.stacks_store, "absorb", lambda *a, **k: None) # no stacks DB in this unit
|
||||
upload = _handler(_FakeServer(), "POST", "/v1/library/upload", monkeypatch)
|
||||
|
||||
src = tmp_path / "clip.mp4"
|
||||
_write_mp4(src)
|
||||
mp4 = src.read_bytes()
|
||||
|
||||
resp = await upload(_Req(raw=mp4, query={"name": "clip.mp4"}, org="acme"))
|
||||
body = json.loads(resp.text)
|
||||
assert resp.status == 200 and body["path"] == "renders/clip.mp4"
|
||||
assert (root / "renders" / "clip.mp4").is_file()
|
||||
# kind defaults to "video" by suffix
|
||||
lib = json.loads((root / "library.json").read_text())
|
||||
row = next(a for a in lib["assets"] if a["path"] == "renders/clip.mp4")
|
||||
assert row["kind"] == "video"
|
||||
# the judge fired on the created row; the colour-restore hook saw (sub, name)
|
||||
assert scheduled.get("rel") == "renders/clip.mp4"
|
||||
assert cm == {"sub": "renders", "name": "clip.mp4"}
|
||||
|
||||
# webm (EBML) accepted too
|
||||
wsrc = tmp_path / "clip.webm"
|
||||
_write_webm(wsrc)
|
||||
r_web = await upload(_Req(raw=wsrc.read_bytes(), query={"name": "clip.webm"}, org="acme"))
|
||||
assert r_web.status == 200
|
||||
|
||||
# garbage named .mp4 → rejected; ftyp at OFFSET 0 (not 4) → rejected (regression pin)
|
||||
r_bad = await upload(_Req(raw=b"totally not a video, just text bytes here", query={"name": "x.mp4"}, org="acme"))
|
||||
assert r_bad.status == 400
|
||||
r_off0 = await upload(_Req(raw=b"ftyp" + b"\x00" * 40, query={"name": "y.mp4"}, org="acme"))
|
||||
assert r_off0.status == 400
|
||||
|
||||
# over the 256MB ceiling → 413 (checked via content_length, no giant allocation)
|
||||
r_big = await upload(_Req(raw=mp4, query={"name": "big.mp4"}, org="acme",
|
||||
content_length=sh._MAX_UPLOAD + 1))
|
||||
assert r_big.status == 413
|
||||
|
||||
|
||||
# ── 7. the manifest sweeper keeps a video as kind="video" ─────────────────────────
|
||||
def test_manifest_video(tmp_path):
|
||||
from scripts import library_manifest as lm
|
||||
root = tmp_path
|
||||
(root / "video").mkdir()
|
||||
(root / "video" / "x.mp4").write_bytes(b"\x00\x00\x00 ftyp")
|
||||
(root / "video" / "empty.mp4").write_bytes(b"") # zero-byte → skipped
|
||||
assert "video/x.mp4" in lm.scan(str(root))
|
||||
assert "video/empty.mp4" not in lm.scan(str(root))
|
||||
|
||||
doc = lm.build(str(root))
|
||||
row = next(a for a in doc["assets"] if a["path"] == "video/x.mp4")
|
||||
assert row["kind"] == "video" and row["status"] == "draft"
|
||||
|
||||
# a curated status + tags survive the next sweep (the erase-every-cycle bug stays dead)
|
||||
lm.write_atomic(str(root), {"assets": [
|
||||
{"path": "video/x.mp4", "kind": "video", "status": "approved", "tags": ["hero"]}]})
|
||||
doc2 = lm.build(str(root))
|
||||
row2 = next(a for a in doc2["assets"] if a["path"] == "video/x.mp4")
|
||||
assert row2["kind"] == "video" and row2["status"] == "approved" and row2["tags"] == ["hero"]
|
||||
|
||||
|
||||
# ── 8. serving + the first-frame poster ───────────────────────────────────────────
|
||||
@pytest.mark.asyncio
|
||||
async def test_serve_and_poster(orgaware, monkeypatch):
|
||||
outd, _ = orgaware
|
||||
root = Path(outd("acme"))
|
||||
_write_mp4(root / "video" / "clip.mp4")
|
||||
_lib(root, [{"path": "video/clip.mp4", "kind": "video", "status": "draft"}])
|
||||
get_file = _handler(_FakeServer(), "GET", "/v1/library/file", monkeypatch)
|
||||
|
||||
# full file served as video/mp4
|
||||
resp = await get_file(_Req(query={"path": "video/clip.mp4"}, org="acme"))
|
||||
assert resp.content_type == "video/mp4"
|
||||
|
||||
# ?thumb=1 → a cached JPEG poster in .thumbs/<key>.jpg
|
||||
tresp = await get_file(_Req(query={"path": "video/clip.mp4", "thumb": "1"}, org="acme"))
|
||||
assert isinstance(tresp, web.FileResponse)
|
||||
key = __import__("re").sub(r"[^a-zA-Z0-9]", "_", str((root / "video" / "clip.mp4").resolve()))[-120:]
|
||||
assert (root / ".thumbs" / f"{key}.jpg").is_file()
|
||||
|
||||
# a poster that cannot decode → 404, NEVER the raw mp4 bytes (use an un-cached path)
|
||||
_write_mp4(root / "video" / "clip2.mp4")
|
||||
_lib(root, [{"path": "video/clip.mp4"}, {"path": "video/clip2.mp4"}])
|
||||
|
||||
def boom(*a, **k):
|
||||
raise RuntimeError("cannot decode")
|
||||
|
||||
monkeypatch.setattr(av, "open", boom)
|
||||
r404 = await get_file(_Req(query={"path": "video/clip2.mp4", "thumb": "1"}, org="acme"))
|
||||
assert isinstance(r404, web.Response) and r404.status == 404
|
||||
assert not isinstance(r404, web.FileResponse)
|
||||
|
||||
|
||||
# ── 9. embedded-metadata read parity: provenance, Open-in-Editor, Re-run ──────────
|
||||
@pytest.mark.asyncio
|
||||
async def test_embedded_and_rerun(orgaware, monkeypatch):
|
||||
outd, _ = orgaware
|
||||
root = Path(outd("acme"))
|
||||
graph = sh._video_graph("a woman on a beach", "video/x_ab12cd", seed=42, length=73)
|
||||
mp4 = root / "video" / "x_ab12cd_00001_.mp4"
|
||||
_write_mp4(mp4, meta={"prompt": json.dumps(graph)})
|
||||
_lib(root, [{"path": "video/x_ab12cd_00001_.mp4", "kind": "video", "status": "draft"}])
|
||||
|
||||
# _embedded reads the container metadata back
|
||||
info = sh._embedded(mp4)
|
||||
assert "prompt" in info and json.loads(info["prompt"])["12"]["class_type"] == "SaveVideo"
|
||||
|
||||
# _provenance yields the video params
|
||||
prov = sh._provenance(mp4)
|
||||
assert prov["workflow"] == "video/x_ab12cd" and prov["steps"] == 20
|
||||
assert prov["cfg"] == 5.0 and prov["seed"] == 42 and prov["length"] == 73 and prov["fps"] == 24
|
||||
|
||||
# GET /v1/workflow serves format 'api' for the done video row. The dispatch ts must
|
||||
# predate the file mtime (else finalize reads the mp4 as a pre-existing sibling).
|
||||
ts = int(mp4.stat().st_mtime) - 100
|
||||
wl.record("acme", kind="video", prompt="p", refs=[], uploads=[], parents=[],
|
||||
output_prefix="video/x_ab12cd", params={"steps": 20}, pid="pv", status="queued")
|
||||
rows = wl.load("acme")
|
||||
for x in rows:
|
||||
if x["id"] == "pv":
|
||||
x["ts"] = ts
|
||||
wl._save("acme", rows)
|
||||
get_wf = _handler(_FakeServer(), "GET", "/v1/workflow", monkeypatch)
|
||||
wfresp = await get_wf(_Req(query={"id": "pv"}, org="acme"))
|
||||
wf = json.loads(wfresp.text)
|
||||
assert wf["format"] == "api" and wf["graph"]["12"]["class_type"] == "SaveVideo"
|
||||
|
||||
# POST /v1/library/rerun dispatches kind='rerun' with a FRESH SaveVideo tag
|
||||
captured = {}
|
||||
|
||||
async def fake_queue(request, server, g):
|
||||
captured["graph"] = g
|
||||
return "pid-rerun"
|
||||
|
||||
monkeypatch.setattr(sh, "_queue_prompt", fake_queue)
|
||||
rerun = _handler(_FakeServer(), "POST", "/v1/library/rerun", monkeypatch)
|
||||
rr = await rerun(_Req(body={"path": "video/x_ab12cd_00001_.mp4"}, org="acme"))
|
||||
assert rr.status == 200
|
||||
sv = next(n for n in captured["graph"].values() if n["class_type"] == "SaveVideo")
|
||||
import re
|
||||
assert re.match(r"^video/x_[0-9a-f]{6}$", sv["inputs"]["filename_prefix"])
|
||||
assert sv["inputs"]["filename_prefix"] != "video/x_ab12cd" # a fresh namespace
|
||||
assert any(r2["kind"] == "rerun" for r2 in wl.load("acme"))
|
||||
|
||||
|
||||
# ── 10. template re-tag: video AND image savers get a fresh job identity ──────────
|
||||
@pytest.mark.asyncio
|
||||
async def test_template_retag(orgaware, monkeypatch):
|
||||
import re
|
||||
outd, _ = orgaware
|
||||
tdir = Path(outd("acme")).parent / "templates"
|
||||
tdir.mkdir(parents=True, exist_ok=True)
|
||||
captured = {}
|
||||
|
||||
async def fake_queue(request, server, g):
|
||||
captured["graph"] = g
|
||||
return "pid-t"
|
||||
|
||||
monkeypatch.setattr(sh, "_queue_prompt", fake_queue)
|
||||
render = _handler(_FakeServer(), "POST", "/v1/templates/render", monkeypatch)
|
||||
|
||||
# a template whose graph saves via SaveVideo
|
||||
vgraph = sh._video_graph("a beach", "video/tmpl", length=73)
|
||||
(tdir / "clip.json").write_text(json.dumps({"name": "Clip", "slug": "clip", "graph": vgraph}))
|
||||
resp = await render(_Req(body={"name": "clip"}, org="acme"))
|
||||
assert resp.status == 200
|
||||
sv = next(n for n in captured["graph"].values() if n["class_type"] == "SaveVideo")
|
||||
assert re.match(r"^video/tmpl_[0-9a-f]{6}$", sv["inputs"]["filename_prefix"])
|
||||
|
||||
# a SaveImage template is re-tagged too (the unique-job-tag rule holds for images)
|
||||
igraph = {"9": {"class_type": "KSampler", "inputs": {"seed": 1, "steps": 30}},
|
||||
"15": {"class_type": "SaveImage", "inputs": {"filename_prefix": "templates/pic"}}}
|
||||
(tdir / "pic.json").write_text(json.dumps({"name": "Pic", "slug": "pic", "graph": igraph}))
|
||||
resp2 = await render(_Req(body={"name": "pic"}, org="acme"))
|
||||
assert resp2.status == 200
|
||||
si = next(n for n in captured["graph"].values() if n["class_type"] == "SaveImage")
|
||||
assert re.match(r"^templates/pic_[0-9a-f]{6}$", si["inputs"]["filename_prefix"])
|
||||
@@ -0,0 +1,164 @@
|
||||
"""worker-mode /prompt gate + the /v1/worker/execute claim seam.
|
||||
|
||||
Covers:
|
||||
- verify_worker_token / reject_untrusted_worker_submit — the ONE policy gate that makes
|
||||
a direct /prompt POST on a worker box refuse un-tokened submits (403), so no render
|
||||
reaches a GPU without appearing in the org's gpu-jobs queue.
|
||||
- POST /v1/worker/execute — validates X-Worker-Token and RUNS the claimed graph on the
|
||||
local prompt_queue, returning {prompt_id}, scoped to the job's org.
|
||||
"""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
|
||||
# CPU CI runners have no NVIDIA driver: force studio's CPU mode BEFORE importing
|
||||
# worker_client (which pulls execution -> studio.model_management, whose import-time
|
||||
# device probe would otherwise call torch.cuda and crash on a GPU-less box). Same
|
||||
# switch the studio pod boots with (--cpu); harmless on a GPU box.
|
||||
import studio.cli_args
|
||||
studio.cli_args.args.cpu = True
|
||||
|
||||
from middleware import worker_client as wc
|
||||
|
||||
|
||||
class _Req:
|
||||
def __init__(self, headers=None):
|
||||
self.headers = headers or {}
|
||||
|
||||
|
||||
# ── The gate policy: worker-mode /prompt refuses un-tokened submits ───────────────
|
||||
|
||||
def test_empty_token_dev_bypass_ONLY_outside_worker_mode(monkeypatch):
|
||||
"""FAIL CLOSED: an empty/unset STUDIO_WORKER_TOKEN is a valid single-trust-domain dev
|
||||
signal ONLY when NOT in worker mode. On a worker box (worker_mode=True) an empty token
|
||||
must NOT open /prompt or /v1/worker/execute — that was the fail-open hole (KMS /v1
|
||||
migration has emptied the secret before)."""
|
||||
monkeypatch.setattr(wc, "WORKER_TOKEN", "")
|
||||
monkeypatch.setattr(wc.args, "worker_mode", False)
|
||||
assert wc.verify_worker_token(_Req()) is True # local dev, single trust domain
|
||||
monkeypatch.setattr(wc.args, "worker_mode", True)
|
||||
assert wc.verify_worker_token(_Req()) is False # worker box + empty token → refuse
|
||||
|
||||
|
||||
def test_verify_worker_token_enforced(monkeypatch):
|
||||
monkeypatch.setattr(wc, "WORKER_TOKEN", "s3cret")
|
||||
assert wc.verify_worker_token(_Req({"X-Worker-Token": "s3cret"})) is True
|
||||
assert wc.verify_worker_token(_Req({"X-Worker-Token": "nope"})) is False
|
||||
assert wc.verify_worker_token(_Req()) is False
|
||||
# constant-time path (hmac.compare_digest) accepts only the exact secret
|
||||
assert wc.verify_worker_token(_Req({"X-Worker-Token": "s3cret "})) is False
|
||||
|
||||
|
||||
def test_gate_allows_when_not_worker_mode(monkeypatch):
|
||||
monkeypatch.setattr(wc, "WORKER_TOKEN", "s3cret")
|
||||
assert wc.reject_untrusted_worker_submit(_Req(), worker_mode=False) is None
|
||||
|
||||
|
||||
def test_gate_403_in_worker_mode_without_token(monkeypatch):
|
||||
monkeypatch.setattr(wc, "WORKER_TOKEN", "s3cret")
|
||||
resp = wc.reject_untrusted_worker_submit(_Req(), worker_mode=True)
|
||||
assert resp is not None and resp.status == 403 # direct /prompt refused
|
||||
|
||||
|
||||
def test_gate_403_in_worker_mode_with_EMPTY_token(monkeypatch):
|
||||
"""Fail closed even when the secret is empty (KMS-emptied): worker mode + empty token
|
||||
still refuses a direct /prompt (403) — runtime belt to main.py's refuse-to-start guard."""
|
||||
monkeypatch.setattr(wc, "WORKER_TOKEN", "")
|
||||
monkeypatch.setattr(wc.args, "worker_mode", True)
|
||||
resp = wc.reject_untrusted_worker_submit(_Req(), worker_mode=True)
|
||||
assert resp is not None and resp.status == 403
|
||||
|
||||
|
||||
def test_gate_allows_worker_mode_with_token(monkeypatch):
|
||||
monkeypatch.setattr(wc, "WORKER_TOKEN", "s3cret")
|
||||
ok = _Req({"X-Worker-Token": "s3cret"})
|
||||
assert wc.reject_untrusted_worker_submit(ok, worker_mode=True) is None
|
||||
|
||||
|
||||
# ── The claim seam: POST /v1/worker/execute validates token + runs locally ────────
|
||||
|
||||
class _Queue:
|
||||
def __init__(self):
|
||||
self.items = []
|
||||
|
||||
def put(self, item):
|
||||
self.items.append(item)
|
||||
|
||||
|
||||
class _NRM:
|
||||
def apply_replacements(self, prompt):
|
||||
pass
|
||||
|
||||
|
||||
class _Server:
|
||||
def __init__(self):
|
||||
self.prompt_queue = _Queue()
|
||||
self.node_replace_manager = _NRM()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def worker_app(monkeypatch):
|
||||
import execution
|
||||
|
||||
async def fake_validate(pid, prompt, targets=None):
|
||||
return (True, None, list(prompt.keys()), {})
|
||||
|
||||
monkeypatch.setattr(execution, "validate_prompt", fake_validate)
|
||||
monkeypatch.setattr(wc, "WORKER_TOKEN", "s3cret")
|
||||
server = _Server()
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
wc.add_worker_routes(routes, server)
|
||||
app.add_routes(routes)
|
||||
app["srv"] = server
|
||||
return app
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_worker_execute_rejects_without_token(worker_app, aiohttp_client):
|
||||
client = await aiohttp_client(worker_app)
|
||||
r = await client.post("/v1/worker/execute",
|
||||
json={"prompt": {"9": {"class_type": "SaveImage", "inputs": {}}}})
|
||||
assert r.status == 401
|
||||
assert worker_app["srv"].prompt_queue.items == [] # nothing ran
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_worker_execute_runs_claimed_job_with_token(worker_app, aiohttp_client):
|
||||
client = await aiohttp_client(worker_app)
|
||||
graph = {"9": {"class_type": "SaveImage", "inputs": {}}}
|
||||
r = await client.post("/v1/worker/execute",
|
||||
headers={"X-Worker-Token": "s3cret"},
|
||||
json={"prompt": graph, "prompt_id": "pid-1", "org": "acme"})
|
||||
assert r.status == 200
|
||||
body = await r.json()
|
||||
assert body["prompt_id"] == "pid-1"
|
||||
items = worker_app["srv"].prompt_queue.items
|
||||
assert len(items) == 1 # queued for local render
|
||||
_number, pid, _prompt, extra, _outputs, _sensitive = items[0]
|
||||
assert pid == "pid-1"
|
||||
assert extra["org_id"] == "acme" # worker-LOCAL scope label (gallery org re-derived from IAM)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_worker_execute_rejects_non_dict_prompt(worker_app, aiohttp_client):
|
||||
"""A non-object prompt (e.g. an injected string) is a clean 400, not a 500."""
|
||||
client = await aiohttp_client(worker_app)
|
||||
r = await client.post("/v1/worker/execute", headers={"X-Worker-Token": "s3cret"},
|
||||
json={"prompt": "<img src=x onerror=alert(1)>"})
|
||||
assert r.status == 400
|
||||
assert worker_app["srv"].prompt_queue.items == [] # nothing ran
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_worker_execute_rejects_non_numeric_number(worker_app, aiohttp_client):
|
||||
"""A non-numeric `number` is a clean 400, not a float() 500."""
|
||||
client = await aiohttp_client(worker_app)
|
||||
r = await client.post("/v1/worker/execute", headers={"X-Worker-Token": "s3cret"},
|
||||
json={"prompt": {"9": {"class_type": "SaveImage", "inputs": {}}},
|
||||
"number": "not-a-number"})
|
||||
assert r.status == 400
|
||||
@@ -0,0 +1,157 @@
|
||||
"""Unit tests for the render quality gate (middleware/quality_gate.py)."""
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
|
||||
from middleware import quality_gate as qg
|
||||
from middleware import studio_home as sh
|
||||
from middleware import worklog
|
||||
|
||||
|
||||
def test_parse_verdict_strict():
|
||||
assert qg.parse_verdict("PASS") == "pass"
|
||||
assert qg.parse_verdict("pass") == "pass"
|
||||
assert qg.parse_verdict("FAIL") == "fail"
|
||||
assert qg.parse_verdict("fail.") == "fail"
|
||||
assert qg.parse_verdict("FAIL — mottled, speckled skin") == "fail"
|
||||
assert qg.parse_verdict("The skin is clean. PASS") == "pass"
|
||||
# garbage / ambiguous / empty → inconclusive (the gate fails open on these)
|
||||
assert qg.parse_verdict("") == "inconclusive"
|
||||
assert qg.parse_verdict(None) == "inconclusive"
|
||||
assert qg.parse_verdict("MAYBE") == "inconclusive"
|
||||
assert qg.parse_verdict("PASS FAIL") == "fail" # both tokens → the LAST is the verdict (reasoning judges deliberate before concluding)
|
||||
assert qg.parse_verdict("failure to load") == "inconclusive" # not a whole-token FAIL
|
||||
|
||||
|
||||
def test_enabled_off_switch(monkeypatch):
|
||||
for off in ("off", "OFF", "0", "false", "no"):
|
||||
monkeypatch.setenv("STUDIO_GATE", off)
|
||||
assert qg.enabled() is False
|
||||
for on in ("", "on", "1", "true", "yes"):
|
||||
monkeypatch.setenv("STUDIO_GATE", on)
|
||||
assert qg.enabled() is True
|
||||
monkeypatch.delenv("STUDIO_GATE", raising=False)
|
||||
assert qg.enabled() is True # default ON
|
||||
|
||||
|
||||
def test_config_defaults_and_overrides(monkeypatch):
|
||||
monkeypatch.delenv("STUDIO_GATE_MODEL", raising=False)
|
||||
monkeypatch.delenv("STUDIO_GATE_URL", raising=False)
|
||||
assert qg._model() == "zen-vl" # the cloud's VL model
|
||||
assert qg._url().endswith("/v1/chat/completions")
|
||||
monkeypatch.setenv("STUDIO_GATE_URL", "http://127.0.0.1:1234/v1/chat/completions")
|
||||
assert qg._url() == "http://127.0.0.1:1234/v1/chat/completions" # local-engine override
|
||||
monkeypatch.setenv("STUDIO_GATE_MODEL", "zen-vl-local")
|
||||
assert qg._model() == "zen-vl-local"
|
||||
|
||||
|
||||
def _const(v):
|
||||
async def j(images, bearer, prompt):
|
||||
return v
|
||||
return j
|
||||
|
||||
|
||||
def _seed(tmp_path, status="draft", org="acme"):
|
||||
from PIL import Image
|
||||
root = tmp_path / "orgs" / org / "output"
|
||||
(root / "renders").mkdir(parents=True)
|
||||
Image.new("RGB", (8, 8), (200, 150, 120)).save(root / "renders" / "r.png", "PNG")
|
||||
(root / "library.json").write_text(json.dumps(
|
||||
{"assets": [{"path": "renders/r.png", "status": status}]}))
|
||||
return root
|
||||
|
||||
|
||||
def _status(root, rel="renders/r.png"):
|
||||
lib = json.loads((root / "library.json").read_text())
|
||||
return next(a["status"] for a in lib["assets"] if a["path"] == rel)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_flags_on_fail_and_annotates_worklog(tmp_path, monkeypatch):
|
||||
root = _seed(tmp_path)
|
||||
monkeypatch.setattr(sh.folder_paths, "get_org_output_directory", lambda o=None: str(root))
|
||||
worklog.record("acme", kind="render", prompt="", refs=[], uploads=[], parents=[],
|
||||
output_prefix="renders/r.png", node="spark", status="done")
|
||||
seen = {}
|
||||
|
||||
async def fake_judge(images, bearer, prompt):
|
||||
seen["b64"], seen["bearer"], seen["prompt"] = images, bearer, prompt
|
||||
return "FAIL - mottled, speckled skin all over the arms and shoulders here"
|
||||
|
||||
monkeypatch.setattr(qg, "judge", fake_judge)
|
||||
await qg.run("acme", root, "renders/r.png", "Bearer t")
|
||||
|
||||
assert seen["b64"] and seen["bearer"] == "Bearer t" # downscaled + same bearer
|
||||
assert _status(root) == "flagged" # asset hidden from default grid
|
||||
row = next(r for r in worklog.load("acme") if r.get("output_prefix") == "renders/r.png")
|
||||
ev = [e for e in row.get("events", []) if e.get("s") == "flagged"]
|
||||
assert len(ev) == 1 and ev[0]["note"].startswith("quality gate: FAIL")
|
||||
assert len(ev[0]["note"]) <= len("quality gate: ") + 80 # first 80 chars of the reply
|
||||
assert row["status"] == "done" # worklog lifecycle untouched — annotation only
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_fails_open_on_pass_inconclusive_and_error(tmp_path, monkeypatch):
|
||||
root = _seed(tmp_path)
|
||||
monkeypatch.setattr(sh.folder_paths, "get_org_output_directory", lambda o=None: str(root))
|
||||
for verdict in ("PASS", "who knows", None): # pass · inconclusive · judge error (None)
|
||||
monkeypatch.setattr(qg, "judge", _const(verdict))
|
||||
await qg.run("acme", root, "renders/r.png", "Bearer t")
|
||||
assert _status(root) == "draft"
|
||||
|
||||
async def boom(images, bearer, prompt):
|
||||
raise RuntimeError("429 rate limited") # 402/429/timeout/unreachable class
|
||||
|
||||
monkeypatch.setattr(qg, "judge", boom)
|
||||
await qg.run("acme", root, "renders/r.png", "Bearer t") # must not raise
|
||||
assert _status(root) == "draft" # good work is never hidden by accident
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_never_clobbers_curated_status(tmp_path, monkeypatch):
|
||||
root = _seed(tmp_path, status="approved") # a human already approved it
|
||||
monkeypatch.setattr(sh.folder_paths, "get_org_output_directory", lambda o=None: str(root))
|
||||
monkeypatch.setattr(qg, "judge", _const("FAIL"))
|
||||
await qg.run("acme", root, "renders/r.png", "Bearer t")
|
||||
assert _status(root) == "approved" # a late FAIL must not yank curated work
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_missing_or_unreadable_is_a_no_op(tmp_path, monkeypatch):
|
||||
root = _seed(tmp_path)
|
||||
monkeypatch.setattr(sh.folder_paths, "get_org_output_directory", lambda o=None: str(root))
|
||||
called = {"n": 0}
|
||||
|
||||
async def counting_judge(images, bearer, prompt):
|
||||
called["n"] += 1
|
||||
return "FAIL"
|
||||
|
||||
monkeypatch.setattr(qg, "judge", counting_judge)
|
||||
await qg.run("acme", root, "renders/gone.png", "Bearer t") # no such file
|
||||
assert called["n"] == 0 and _status(root) == "draft" # never judged, never flagged
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_is_tenant_scoped(tmp_path, monkeypatch):
|
||||
a_root = _seed(tmp_path, org="acme")
|
||||
b_root = _seed(tmp_path, org="globex")
|
||||
roots = {"acme": a_root, "globex": b_root}
|
||||
monkeypatch.setattr(sh.folder_paths, "get_org_output_directory",
|
||||
lambda o=None: str(roots.get(o, a_root)))
|
||||
monkeypatch.setattr(qg, "judge", _const("FAIL"))
|
||||
await qg.run("acme", a_root, "renders/r.png", "Bearer t")
|
||||
assert _status(a_root) == "flagged"
|
||||
assert _status(b_root) == "draft" # the other tenant is untouched
|
||||
|
||||
|
||||
def test_schedule_off_switch_never_fires(monkeypatch):
|
||||
"""STUDIO_GATE=off short-circuits before touching the loop — nothing scheduled."""
|
||||
monkeypatch.setenv("STUDIO_GATE", "off")
|
||||
fired = {"n": 0}
|
||||
monkeypatch.setattr(qg.asyncio, "ensure_future", lambda coro: fired.__setitem__("n", fired["n"] + 1))
|
||||
qg.schedule(object(), "acme", Path("/tmp"), "renders/r.png")
|
||||
assert fired["n"] == 0
|
||||
@@ -0,0 +1,358 @@
|
||||
"""Stacks — the Procreate-style gallery folders.
|
||||
|
||||
Covers the per-org SQLite store (middleware/stacks_store.py) and its wiring into
|
||||
the Studio home routes: every REST route, transactional membership, tenancy
|
||||
isolation, one-stack-per-asset, order persistence, and delete/unstack semantics
|
||||
(remove is never delete; delete defaults to preserving; delete-images soft-deletes
|
||||
in the library), plus the generation-integration routing and auto-stack setting.
|
||||
"""
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
|
||||
from middleware import stacks_store as st
|
||||
from middleware import studio_home as sh
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def org(tmp_path, monkeypatch):
|
||||
"""A tmp per-org tree + a fresh connection cache — the SAME folder_paths seam
|
||||
production tenancy uses, so a Stack DB lands at ``orgs/<org>/stacks.db``."""
|
||||
def outd(o=None):
|
||||
d = tmp_path / "orgs" / (o or "default") / "output"
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return str(d)
|
||||
|
||||
def ind(o=None):
|
||||
d = tmp_path / "orgs" / (o or "default") / "input"
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return str(d)
|
||||
|
||||
monkeypatch.setattr(sh.folder_paths, "get_org_output_directory", outd)
|
||||
monkeypatch.setattr(sh.folder_paths, "get_org_input_directory", ind)
|
||||
monkeypatch.setattr(sh.folder_paths, "get_output_directory", lambda: str(tmp_path))
|
||||
for c in st._conns.values():
|
||||
try:
|
||||
c.close()
|
||||
except Exception:
|
||||
pass
|
||||
st._conns.clear()
|
||||
st._locks.clear()
|
||||
yield "acme"
|
||||
|
||||
|
||||
# ── store: the file lives where tenancy expects it ──────────────────────────────
|
||||
def test_db_path_is_org_rooted(org, tmp_path):
|
||||
st.create(org, "x", [])
|
||||
assert (tmp_path / "orgs" / "acme" / "stacks.db").is_file()
|
||||
|
||||
|
||||
# ── create: drag A onto B = a Stack with BOTH, nothing copied/merged/deleted ─────
|
||||
def test_create_from_two_assets_holds_both_cover_first(org):
|
||||
s = st.create(org, "Fashion", ["a.png", "b.png"])
|
||||
assert s["name"] == "Fashion" and s["count"] == 2
|
||||
assert [m["assetId"] for m in s["assets"]] == ["a.png", "b.png"]
|
||||
assert s["coverAssetId"] == "a.png" # default cover = first image
|
||||
# reloading from the DB yields the same membership (nothing copied/merged)
|
||||
got = st.get(org, s["id"])
|
||||
assert [m["assetId"] for m in got["assets"]] == ["a.png", "b.png"]
|
||||
|
||||
|
||||
def test_create_defaults_untitled_and_dedupes(org):
|
||||
s = st.create(org, " ", ["a.png", "a.png", "b.png"])
|
||||
assert s["name"] == "Untitled Stack"
|
||||
assert [m["assetId"] for m in s["assets"]] == ["a.png", "b.png"]
|
||||
|
||||
|
||||
def test_single_image_stack_allowed(org):
|
||||
s = st.create(org, "solo", ["only.png"])
|
||||
assert s["count"] == 1 and st.get(org, s["id"])["count"] == 1
|
||||
|
||||
|
||||
def test_empty_stack_allowed_for_generation_destination(org):
|
||||
s = st.create(org, "New", [])
|
||||
assert s["count"] == 0 and s["coverAssetId"] is None
|
||||
|
||||
|
||||
# ── one stack per asset: adding to a second Stack MOVES it ───────────────────────
|
||||
def test_one_stack_per_asset_move_on_readd(org):
|
||||
a = st.create(org, "A", ["x.png", "y.png"])
|
||||
b = st.create(org, "B", ["z.png"])
|
||||
st.add_assets(org, b["id"], ["x.png"]) # steal x from A
|
||||
assert [m["assetId"] for m in st.get(org, a["id"])["assets"]] == ["y.png"]
|
||||
assert [m["assetId"] for m in st.get(org, b["id"])["assets"]] == ["z.png", "x.png"]
|
||||
assert st.stack_of(org, "x.png") == b["id"]
|
||||
|
||||
|
||||
def test_moving_the_cover_out_repoints_source_cover(org):
|
||||
a = st.create(org, "A", ["cover.png", "second.png"])
|
||||
assert a["coverAssetId"] == "cover.png"
|
||||
b = st.create(org, "B", [])
|
||||
st.add_assets(org, b["id"], ["cover.png"]) # cover leaves A
|
||||
assert st.get(org, a["id"])["coverAssetId"] == "second.png"
|
||||
|
||||
|
||||
# ── add / remove: remove returns to gallery, never deletes ───────────────────────
|
||||
def test_remove_returns_asset_and_repoints_cover(org):
|
||||
s = st.create(org, "S", ["a.png", "b.png", "c.png"])
|
||||
st.remove_assets(org, s["id"], ["a.png"]) # a.png was the cover
|
||||
got = st.get(org, s["id"])
|
||||
assert [m["assetId"] for m in got["assets"]] == ["b.png", "c.png"]
|
||||
assert [m["position"] for m in got["assets"]] == [0, 1] # positions compact
|
||||
assert got["coverAssetId"] == "b.png"
|
||||
assert st.stack_of(org, "a.png") is None # free again, not deleted
|
||||
|
||||
|
||||
# ── reorder: order persists and is validated against membership ──────────────────
|
||||
def test_reorder_persists_and_survives_reload(org):
|
||||
s = st.create(org, "S", ["a.png", "b.png", "c.png"])
|
||||
st.reorder(org, s["id"], ["c.png", "a.png", "b.png"])
|
||||
assert [m["assetId"] for m in st.get(org, s["id"])["assets"]] == ["c.png", "a.png", "b.png"]
|
||||
|
||||
|
||||
def test_reorder_rejects_wrong_set(org):
|
||||
s = st.create(org, "S", ["a.png", "b.png"])
|
||||
with pytest.raises(ValueError):
|
||||
st.reorder(org, s["id"], ["a.png"]) # missing a member
|
||||
with pytest.raises(ValueError):
|
||||
st.reorder(org, s["id"], ["a.png", "b.png", "c.png"]) # foreign member
|
||||
|
||||
|
||||
# ── rename / describe / cover / archive ─────────────────────────────────────────
|
||||
def test_patch_name_description_cover(org):
|
||||
s = st.create(org, "old", ["a.png", "b.png"])
|
||||
p = st.patch(org, s["id"], {"name": "new", "description": "desc", "cover": "b.png"})
|
||||
assert p["name"] == "new" and p["description"] == "desc" and p["coverAssetId"] == "b.png"
|
||||
|
||||
|
||||
def test_cover_must_be_a_member(org):
|
||||
s = st.create(org, "S", ["a.png"])
|
||||
with pytest.raises(ValueError):
|
||||
st.patch(org, s["id"], {"cover": "not-in-stack.png"})
|
||||
|
||||
|
||||
def test_archive_hides_from_default_list(org):
|
||||
s = st.create(org, "S", ["a.png"])
|
||||
st.patch(org, s["id"], {"archived": True})
|
||||
assert st.get(org, s["id"])["archivedAt"] is not None
|
||||
assert s["id"] not in [x["id"] for x in st.list_stacks(org, include_archived=False)]
|
||||
assert s["id"] in [x["id"] for x in st.list_stacks(org, include_archived=True)]
|
||||
|
||||
|
||||
# ── unstack: container gone, EVERY asset preserved ──────────────────────────────
|
||||
def test_unstack_preserves_all_assets(org):
|
||||
s = st.create(org, "S", ["a.png", "b.png", "c.png"])
|
||||
freed = st.unstack(org, s["id"])
|
||||
assert freed == ["a.png", "b.png", "c.png"]
|
||||
assert st.get(org, s["id"]) is None # container dissolved
|
||||
assert all(st.stack_of(org, a) is None for a in freed)
|
||||
|
||||
|
||||
# ── delete: defaults to preserving images; images-mode soft-deletes them ─────────
|
||||
def test_delete_only_returns_images(org):
|
||||
s = st.create(org, "S", ["a.png", "b.png"])
|
||||
out = st.delete(org, s["id"], "only")
|
||||
assert out["mode"] == "only" and out["assetIds"] == ["a.png", "b.png"]
|
||||
assert st.get(org, s["id"]) is None
|
||||
assert st.stack_of(org, "a.png") is None
|
||||
|
||||
|
||||
def test_delete_images_reports_assets_for_soft_delete(org):
|
||||
s = st.create(org, "S", ["a.png", "b.png"])
|
||||
out = st.delete(org, s["id"], "images")
|
||||
assert out["mode"] == "images" and set(out["assetIds"]) == {"a.png", "b.png"}
|
||||
|
||||
|
||||
# ── generation integration: route a prefix, absorb the landed outputs ───────────
|
||||
def test_route_and_absorb_assigns_all_outputs(org):
|
||||
s = st.create(org, "Batch", [])
|
||||
st.route_prefix(org, "renders/foo", s["id"])
|
||||
# SaveImage counter form + a second output of the same job both join
|
||||
assert st.absorb(org, "renders/foo_00001_.png") == s["id"]
|
||||
assert st.absorb(org, "renders/foo_00002_.png") == s["id"]
|
||||
got = st.get(org, s["id"])
|
||||
assert [m["assetId"] for m in got["assets"]] == ["renders/foo_00001_.png", "renders/foo_00002_.png"]
|
||||
assert got["coverAssetId"] == "renders/foo_00001_.png"
|
||||
|
||||
|
||||
def test_absorb_ingested_path_and_miss(org):
|
||||
s = st.create(org, "S", [])
|
||||
st.route_prefix(org, "composes/bar", s["id"])
|
||||
assert st.absorb(org, "composes/bar") == s["id"] # ingested-path (no counter)
|
||||
assert st.absorb(org, "renders/unrelated_00001_.png") is None
|
||||
|
||||
|
||||
# ── per-org settings (auto-stack batch generations) ─────────────────────────────
|
||||
def test_settings_roundtrip(org):
|
||||
assert st.get_setting(org, "auto_stack_batches", False) is False
|
||||
st.set_setting(org, "auto_stack_batches", True)
|
||||
assert st.get_setting(org, "auto_stack_batches", False) is True
|
||||
|
||||
|
||||
def test_dest_stack_auto_creates_for_batch_when_enabled(org):
|
||||
# explicit id wins
|
||||
assert sh._dest_stack(org, "sid-123", prompt="x", multi=True, owner="", workspace=None) == "sid-123"
|
||||
# off → None even for a batch
|
||||
assert sh._dest_stack(org, None, prompt="a red hat", multi=True, owner="", workspace=None) is None
|
||||
st.set_setting(org, "auto_stack_batches", True)
|
||||
# on + single output → still None (auto-stack is for batches)
|
||||
assert sh._dest_stack(org, None, prompt="x", multi=False, owner="", workspace=None) is None
|
||||
# on + batch → a new '<short prompt> — <date>' Stack
|
||||
sid = sh._dest_stack(org, None, prompt="a red hat on a model", multi=True, owner="u1", workspace="proj")
|
||||
got = st.get(org, sid)
|
||||
assert got is not None and got["name"].startswith("a red hat on a") and got["ownerId"] == "u1"
|
||||
assert got["workspaceId"] == "proj"
|
||||
|
||||
|
||||
def test_index_upload_absorbs_landed_output_into_routed_stack(org, tmp_path):
|
||||
"""The real ingest point: a render routed to a Stack, then its output landing
|
||||
via _index_upload, joins that Stack — the ONE membership-assignment hook (and it
|
||||
still returns `created` for the quality gate)."""
|
||||
root = Path(sh.folder_paths.get_org_output_directory(org))
|
||||
(root / "library.json").write_text(json.dumps({"assets": []}))
|
||||
s = st.create(org, "Gen", [])
|
||||
st.route_prefix(org, "renders/gen", s["id"])
|
||||
created = sh._index_upload(org, root, "renders/gen_00001_.png",
|
||||
design=None, kind="render", role=None, node="upload")
|
||||
assert created is True # gate contract preserved
|
||||
assert [m["assetId"] for m in st.get(org, s["id"])["assets"]] == ["renders/gen_00001_.png"]
|
||||
|
||||
|
||||
def test_soft_delete_assets_marks_library_deleted(org, tmp_path):
|
||||
root = Path(sh.folder_paths.get_org_output_directory(org))
|
||||
(root / "library.json").write_text(json.dumps({"assets": [
|
||||
{"path": "a.png", "status": "draft"}, {"path": "b.png", "status": "approved"}]}))
|
||||
sh._soft_delete_assets(org, ["a.png"])
|
||||
lib = json.loads((root / "library.json").read_text())
|
||||
by = {a["path"]: a["status"] for a in lib["assets"]}
|
||||
assert by == {"a.png": "deleted", "b.png": "approved"} # only the named one, files never unlinked
|
||||
assert (root / "library.json").is_file()
|
||||
|
||||
|
||||
# ══ route surface ═══════════════════════════════════════════════════════════════
|
||||
class _FakeServer:
|
||||
port = 1
|
||||
|
||||
|
||||
async def _client(monkeypatch, sub="u1"):
|
||||
from aiohttp import web
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
@web.middleware
|
||||
async def fake_auth(request, handler):
|
||||
o = request.headers.get("X-Org")
|
||||
if o:
|
||||
request["iam_user"] = {"org_id": o, "orgs": [o], "sub": sub}
|
||||
return await handler(request)
|
||||
|
||||
app = web.Application(middlewares=[fake_auth])
|
||||
routes = web.RouteTableDef()
|
||||
sh.add_studio_home_routes(routes, _FakeServer())
|
||||
app.add_routes(routes)
|
||||
ts = TestServer(app)
|
||||
await ts.start_server()
|
||||
client = TestClient(ts)
|
||||
await client.start_server()
|
||||
return client
|
||||
|
||||
|
||||
async def _mk(client, org, name, ids):
|
||||
r = await client.post("/v1/stacks", json={"name": name, "assetIds": ids}, headers={"X-Org": org})
|
||||
assert r.status == 200
|
||||
return await r.json()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routes_full_lifecycle(org, monkeypatch):
|
||||
client = await _client(monkeypatch)
|
||||
# create
|
||||
s = await _mk(client, "acme", "Fashion", ["a.png", "b.png"])
|
||||
assert s["count"] == 2 and s["ownerId"] == "u1"
|
||||
sid = s["id"]
|
||||
# list
|
||||
lst = await (await client.get("/v1/stacks", headers={"X-Org": "acme"})).json()
|
||||
assert lst["org"] == "acme" and [x["id"] for x in lst["stacks"]] == [sid]
|
||||
# get
|
||||
assert (await (await client.get(f"/v1/stacks/{sid}", headers={"X-Org": "acme"})).json())["name"] == "Fashion"
|
||||
# add
|
||||
r = await client.post(f"/v1/stacks/{sid}/assets", json={"assetIds": ["c.png"]}, headers={"X-Org": "acme"})
|
||||
assert [m["assetId"] for m in (await r.json())["assets"]] == ["a.png", "b.png", "c.png"]
|
||||
# reorder (persists)
|
||||
r = await client.patch(f"/v1/stacks/{sid}/order", json={"assetIds": ["c.png", "b.png", "a.png"]}, headers={"X-Org": "acme"})
|
||||
assert [m["assetId"] for m in (await r.json())["assets"]] == ["c.png", "b.png", "a.png"]
|
||||
reloaded = await (await client.get(f"/v1/stacks/{sid}", headers={"X-Org": "acme"})).json()
|
||||
assert [m["assetId"] for m in reloaded["assets"]] == ["c.png", "b.png", "a.png"]
|
||||
# rename + cover
|
||||
r = await client.patch(f"/v1/stacks/{sid}", json={"name": "Looks", "cover": "b.png"}, headers={"X-Org": "acme"})
|
||||
body = await r.json()
|
||||
assert body["name"] == "Looks" and body["coverAssetId"] == "b.png"
|
||||
# remove (returns to gallery)
|
||||
r = await client.delete(f"/v1/stacks/{sid}/assets", json={"assetIds": ["c.png"]}, headers={"X-Org": "acme"})
|
||||
assert [m["assetId"] for m in (await r.json())["assets"]] == ["b.png", "a.png"]
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routes_unstack_and_delete_default_preserves(org, monkeypatch):
|
||||
client = await _client(monkeypatch)
|
||||
s = await _mk(client, "acme", "S", ["a.png", "b.png"])
|
||||
r = await client.post(f"/v1/stacks/{s['id']}/unstack", headers={"X-Org": "acme"})
|
||||
assert (await r.json())["assetIds"] == ["a.png", "b.png"]
|
||||
assert (await client.get(f"/v1/stacks/{s['id']}", headers={"X-Org": "acme"})).status == 404
|
||||
# delete default (mode=only) preserves images
|
||||
s2 = await _mk(client, "acme", "S2", ["x.png"])
|
||||
r = await client.delete(f"/v1/stacks/{s2['id']}", headers={"X-Org": "acme"})
|
||||
j = await r.json()
|
||||
assert j["mode"] == "only" and j["assetIds"] == ["x.png"]
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routes_delete_images_soft_deletes_library(org, monkeypatch):
|
||||
root = Path(sh.folder_paths.get_org_output_directory("acme"))
|
||||
(root / "library.json").write_text(json.dumps({"assets": [{"path": "a.png", "status": "draft"}]}))
|
||||
client = await _client(monkeypatch)
|
||||
s = await _mk(client, "acme", "S", ["a.png"])
|
||||
r = await client.delete(f"/v1/stacks/{s['id']}?mode=images", headers={"X-Org": "acme"})
|
||||
assert (await r.json())["mode"] == "images"
|
||||
lib = json.loads((root / "library.json").read_text())
|
||||
assert lib["assets"][0]["status"] == "deleted" # soft-deleted, file kept
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routes_tenancy_cross_org_404(org, monkeypatch):
|
||||
client = await _client(monkeypatch)
|
||||
s = await _mk(client, "acme", "Secret", ["a.png"])
|
||||
sid = s["id"]
|
||||
# globex sees an empty list and cannot read/patch/delete acme's Stack
|
||||
assert (await (await client.get("/v1/stacks", headers={"X-Org": "globex"})).json())["stacks"] == []
|
||||
assert (await client.get(f"/v1/stacks/{sid}", headers={"X-Org": "globex"})).status == 404
|
||||
assert (await client.patch(f"/v1/stacks/{sid}", json={"name": "x"}, headers={"X-Org": "globex"})).status == 404
|
||||
assert (await client.post(f"/v1/stacks/{sid}/assets", json={"assetIds": ["z.png"]}, headers={"X-Org": "globex"})).status == 404
|
||||
assert (await client.delete(f"/v1/stacks/{sid}", headers={"X-Org": "globex"})).status == 404
|
||||
assert (await client.post(f"/v1/stacks/{sid}/unstack", headers={"X-Org": "globex"})).status == 404
|
||||
# acme still intact
|
||||
assert (await client.get(f"/v1/stacks/{sid}", headers={"X-Org": "acme"})).status == 200
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routes_validation_and_settings(org, monkeypatch):
|
||||
client = await _client(monkeypatch)
|
||||
# bad cover → 400
|
||||
s = await _mk(client, "acme", "S", ["a.png"])
|
||||
assert (await client.patch(f"/v1/stacks/{s['id']}", json={"cover": "ghost.png"}, headers={"X-Org": "acme"})).status == 400
|
||||
# bad reorder → 400
|
||||
assert (await client.patch(f"/v1/stacks/{s['id']}/order", json={"assetIds": ["a.png", "b.png"]}, headers={"X-Org": "acme"})).status == 400
|
||||
# settings roundtrip
|
||||
assert (await (await client.get("/v1/stacks/settings", headers={"X-Org": "acme"})).json())["autoStack"] is False
|
||||
r = await client.patch("/v1/stacks/settings", json={"autoStack": True}, headers={"X-Org": "acme"})
|
||||
assert (await r.json())["autoStack"] is True
|
||||
# "settings" is NOT read as a stack id
|
||||
assert (await client.get("/v1/stacks/settings", headers={"X-Org": "acme"})).status == 200
|
||||
await client.close()
|
||||
@@ -27,12 +27,92 @@ def test_fix_graph_keeps_recipe_invariants():
|
||||
ks = g["13"]["inputs"]
|
||||
assert ks["cfg"] == 4.0 and ks["steps"] == 30 # proven floor, not the weak 3.0
|
||||
assert g["12"]["class_type"] == "CFGNorm"
|
||||
assert g["8"]["inputs"]["reference_latents_method"] == "index_timestep_zero"
|
||||
assert g["6"]["inputs"]["image1"] == ["2", 0] # scaled ref, never raw
|
||||
# The canonical Qwen edit: NO Flux-Kontext reference-method override (that
|
||||
# `index_timestep_zero` graft onto Qwen conditioning was the noise-stipple cause);
|
||||
# the SAME scaled reference feeds conditioning AND the VAEEncode latent. It is now a
|
||||
# LOW-denoise edit (default 0.25) — never 1.0 (that is Regenerate).
|
||||
assert all(n["class_type"] != "FluxKontextMultiReferenceLatentMethod" for n in g.values())
|
||||
assert ks["denoise"] == sh.EDIT_DENOISE_DEFAULT == 0.25
|
||||
assert g[ks["latent_image"][0]]["class_type"] == "VAEEncode"
|
||||
assert g["6"]["inputs"]["image1"] == ["2", 0] # the scaled reference conditions the edit
|
||||
assert g["10"]["inputs"]["pixels"] == ["2", 0] # the SAME scaled reference sizes the latent
|
||||
assert g["2"]["class_type"] == "FluxKontextImageScale"
|
||||
assert "straps thicker" in g["6"]["inputs"]["prompt"]
|
||||
assert g["15"]["inputs"]["filename_prefix"] == "fixes/test"
|
||||
|
||||
|
||||
# ── Generated-image Edit path: the spec invariants (edit ≠ regenerate) ───────────
|
||||
|
||||
def test_edit_uses_source_latent_never_empty():
|
||||
"""Edit encodes the SOURCE (LoadImage → VAEEncode) as the KSampler latent — never
|
||||
an EmptyLatentImage (that would be a text-to-image regenerate)."""
|
||||
g = sh._fix_graph("fix_src.png", "brighten very slightly", "fixes/e")
|
||||
ks = g["13"]["inputs"]
|
||||
assert g["1"]["class_type"] == "LoadImage"
|
||||
assert g[ks["latent_image"][0]]["class_type"] == "VAEEncode"
|
||||
assert not any("Empty" in n.get("class_type", "") for n in g.values())
|
||||
assert sh._assert_edit_graph(g) == ""
|
||||
|
||||
|
||||
def test_edit_denoise_defaults_near_025_and_maps_strength():
|
||||
assert sh._edit_denoise(None) == 0.25
|
||||
assert sh._edit_denoise(25) == 0.25 # percent
|
||||
assert sh._edit_denoise(0.4) == 0.4 # fraction
|
||||
assert sh._edit_denoise(10) == 0.1
|
||||
assert sh._edit_denoise(100) == 0.95 # 100% clamps below a true 1.0 regen
|
||||
assert sh._edit_denoise("bad") == 0.25
|
||||
assert sh._fix_graph("s.png", "x", "p", denoise=0.4)["13"]["inputs"]["denoise"] == 0.4
|
||||
|
||||
|
||||
def test_regenerate_empty_latent_is_rejected_as_an_edit():
|
||||
"""A graph with EmptyLatentImage / denoise 1.0 must be refused by the edit guard —
|
||||
an Edit can NEVER silently become a Regenerate."""
|
||||
regen = {
|
||||
"10": {"class_type": "EmptySD3LatentImage", "inputs": {}},
|
||||
"13": {"class_type": "KSampler", "inputs": {"latent_image": ["10", 0], "denoise": 1.0}},
|
||||
"15": {"class_type": "SaveImage", "inputs": {}},
|
||||
}
|
||||
assert "EmptyLatentImage" in sh._assert_edit_graph(regen)
|
||||
# denoise 1.0 over a real source latent is still refused (that band is Regenerate).
|
||||
hot = sh._fix_graph("s.png", "x", "p")
|
||||
hot["13"]["inputs"]["denoise"] = 1.0
|
||||
assert "denoise" in sh._assert_edit_graph(hot)
|
||||
|
||||
|
||||
def test_edit_seed_reuse_is_deterministic_and_varying_still_edits_source():
|
||||
a = sh._fix_graph("s.png", "x", "p", seed=123)
|
||||
b = sh._fix_graph("s.png", "x", "p", seed=123)
|
||||
assert a["13"]["inputs"]["seed"] == b["13"]["inputs"]["seed"] == 123 # locked seed → deterministic
|
||||
c = sh._fix_graph("s.png", "x", "p", seed=999)
|
||||
assert c["13"]["inputs"]["seed"] == 999
|
||||
# a varied seed STILL edits the source: latent is the VAEEncode, not an empty one.
|
||||
assert c[c["13"]["inputs"]["latent_image"][0]]["class_type"] == "VAEEncode"
|
||||
|
||||
|
||||
def test_edit_prompt_preserves_and_does_not_expand():
|
||||
"""The edit instruction is wrapped with a PRESERVATION directive, not creative
|
||||
scene/style expansion."""
|
||||
g = sh._fix_graph("s.png", "brighten very slightly", "p")
|
||||
pos = g["6"]["inputs"]["prompt"].lower()
|
||||
assert "brighten very slightly" in pos
|
||||
assert "apply only the requested change" in pos and "leave everything else unchanged" in pos
|
||||
|
||||
|
||||
def test_masked_edit_inpaints_and_composites_outside_pixels():
|
||||
"""Edit selected area: SetLatentNoiseMask restricts sampling, ImageCompositeMasked
|
||||
keeps pixels OUTSIDE the mask (the composite's destination is the original)."""
|
||||
g = sh._inpaint_graph("src.png", "mask.png", "recolor the top", "fixes/m")
|
||||
cts = {n["class_type"] for n in g.values()}
|
||||
assert {"SetLatentNoiseMask", "ImageCompositeMasked", "VAEEncode", "LoadImage"} <= cts
|
||||
ks = g["13"]["inputs"]
|
||||
assert g[ks["latent_image"][0]]["class_type"] == "SetLatentNoiseMask"
|
||||
comp = next(n for n in g.values() if n["class_type"] == "ImageCompositeMasked")
|
||||
assert comp["inputs"]["destination"] == ["2", 0] # composite OVER the original scaled source
|
||||
save = next(n for n in g.values() if n["class_type"] in sh._SAVERS)
|
||||
assert save["inputs"]["images"][0] == "18" # the composited image is saved
|
||||
assert sh._assert_edit_graph(g) == ""
|
||||
|
||||
|
||||
def test_safe_asset_blocks_traversal(tmp_path):
|
||||
root = tmp_path / "out"
|
||||
(root / "designs").mkdir(parents=True)
|
||||
@@ -44,7 +124,7 @@ def test_safe_asset_blocks_traversal(tmp_path):
|
||||
|
||||
|
||||
def test_status_vocab_matches_library_lifecycle():
|
||||
assert sh._STATUSES == ("draft", "approved", "queued", "published", "deleted")
|
||||
assert sh._STATUSES == ("draft", "approved", "queued", "published", "deleted", "flagged")
|
||||
|
||||
|
||||
def test_resolve_uploads_keeps_only_safe_existing_images(tmp_path, monkeypatch):
|
||||
@@ -340,3 +420,515 @@ async def test_home_redirect_only_hits_root_html(monkeypatch):
|
||||
req = make_mocked_request("GET", "/", headers={"Accept": "text/html"})
|
||||
resp = await sh.home_redirect(req, handler)
|
||||
assert resp.text == "spa"
|
||||
|
||||
|
||||
# ── Shared shell: editor index injection + templates + workflow + wallet ─────────
|
||||
async def _client(server, org_middleware=True):
|
||||
"""A real aiohttp app with the studio-home routes and a fake-auth middleware that
|
||||
stamps request['iam_user'] from an X-Org header — so route tests exercise the
|
||||
SAME server-side org scoping (_org_of) production uses."""
|
||||
from aiohttp import web
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
@web.middleware
|
||||
async def fake_auth(request, handler):
|
||||
o = request.headers.get("X-Org")
|
||||
if o:
|
||||
request["iam_user"] = {"org_id": o, "orgs": [o]}
|
||||
return await handler(request)
|
||||
|
||||
app = web.Application(middlewares=[fake_auth] if org_middleware else [])
|
||||
routes = web.RouteTableDef()
|
||||
sh.add_studio_home_routes(routes, server)
|
||||
app.add_routes(routes)
|
||||
ts = TestServer(app)
|
||||
await ts.start_server()
|
||||
client = TestClient(ts)
|
||||
await client.start_server()
|
||||
return client
|
||||
|
||||
|
||||
def test_editor_index_injects_shell_once(tmp_path):
|
||||
root = tmp_path / "web"
|
||||
root.mkdir()
|
||||
(root / "index.html").write_text("<html><head></head><body><div id=app></div></body></html>")
|
||||
resp = sh.editor_index(str(root))
|
||||
assert resp.headers["Cache-Control"] == "no-cache"
|
||||
assert '<link rel="stylesheet" href="/studio/shell.css"></head>' in resp.text
|
||||
assert '<script src="/studio/shell.js" defer></script></body>' in resp.text
|
||||
# idempotent: an index that already carries the shell is served unchanged
|
||||
(root / "index.html").write_text(resp.text)
|
||||
again = sh.editor_index(str(root))
|
||||
assert again.text.count("/studio/shell.js") == 1
|
||||
|
||||
|
||||
def test_template_api_prompt_shapes():
|
||||
api = {"9": {"class_type": "KSampler", "inputs": {}}, "15": {"class_type": "SaveImage", "inputs": {}}}
|
||||
assert sh._template_api_prompt({"output": api, "workflow": {"nodes": []}}) == api # editor save
|
||||
assert sh._template_api_prompt(api) == api # bare API prompt
|
||||
assert sh._template_api_prompt({"nodes": [{"id": 1}], "last_node_id": 1}) is None # UI-only, not runnable
|
||||
assert sh._template_api_prompt({}) is None and sh._template_api_prompt("x") is None
|
||||
assert sh._template_slug("Karma — Ghost Mannequin!") == "karma-ghost-mannequin"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_templates_crud_and_tenant_isolation(tmp_path, monkeypatch):
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
client = await _client(_FakeServer())
|
||||
graph = {"output": {"15": {"class_type": "SaveImage", "inputs": {"filename_prefix": "t/x"}}}}
|
||||
r = await client.post("/v1/templates", json={"name": "Ghost Mannequin", "graph": graph}, headers={"X-Org": "acme"})
|
||||
assert r.status == 200 and (await r.json())["slug"] == "ghost-mannequin"
|
||||
# acme sees it (renderable — the graph carries an API prompt), fetches it by slug
|
||||
r = await client.get("/v1/templates", headers={"X-Org": "acme"})
|
||||
body = await r.json()
|
||||
assert body["org"] == "acme" and len(body["templates"]) == 1
|
||||
assert body["templates"][0]["name"] == "Ghost Mannequin" and body["templates"][0]["renderable"] is True
|
||||
r = await client.get("/v1/templates/ghost-mannequin", headers={"X-Org": "acme"})
|
||||
assert (await r.json())["graph"] == graph
|
||||
# tenant isolation: globex's list is empty and it can't read acme's template
|
||||
assert (await (await client.get("/v1/templates", headers={"X-Org": "globex"})).json())["templates"] == []
|
||||
assert (await client.get("/v1/templates/ghost-mannequin", headers={"X-Org": "globex"})).status == 404
|
||||
# validation
|
||||
assert (await client.post("/v1/templates", json={"name": "", "graph": graph}, headers={"X-Org": "acme"})).status == 400
|
||||
assert (await client.post("/v1/templates", json={"name": "x", "graph": {}}, headers={"X-Org": "acme"})).status == 400
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_template_render_dispatches_and_rejects_ui_only(tmp_path, monkeypatch):
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
async def fake_queue(request, server, graph):
|
||||
return "pid-tmpl"
|
||||
monkeypatch.setattr(sh, "_queue_prompt", fake_queue)
|
||||
client = await _client(_FakeServer())
|
||||
runnable = {"output": {"9": {"class_type": "KSampler", "inputs": {"seed": 1}},
|
||||
"15": {"class_type": "SaveImage", "inputs": {"filename_prefix": "t/run"}}}}
|
||||
await client.post("/v1/templates", json={"name": "Run Me", "graph": runnable}, headers={"X-Org": "acme"})
|
||||
r = await client.post("/v1/templates/render", json={"name": "run-me"}, headers={"X-Org": "acme"})
|
||||
assert r.status == 200 and (await r.json())["prompt_id"] == "pid-tmpl"
|
||||
assert sh.worklog.load("acme")[-1]["kind"] == "template" # went through the ONE funnel
|
||||
# UI-only graph can't run headless -> 422 with a reason, not a silent dispatch
|
||||
await client.post("/v1/templates", json={"name": "UI Only", "graph": {"nodes": [], "last_node_id": 0}}, headers={"X-Org": "acme"})
|
||||
assert (await client.post("/v1/templates/render", json={"name": "ui-only"}, headers={"X-Org": "acme"})).status == 422
|
||||
# cross-tenant render: globex has no such template -> 404
|
||||
assert (await client.post("/v1/templates/render", json={"name": "run-me"}, headers={"X-Org": "globex"})).status == 404
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_workflow_reads_embedded_graph_tenant_scoped(tmp_path, monkeypatch):
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
root = Path(sh.folder_paths.get_org_output_directory("acme"))
|
||||
(root / "library.json").write_text('{"assets":[]}')
|
||||
(root / "fixes").mkdir(parents=True, exist_ok=True)
|
||||
info = PngInfo()
|
||||
ui = {"nodes": [{"id": 1, "type": "SaveImage"}], "last_node_id": 1}
|
||||
info.add_text("workflow", json.dumps(ui))
|
||||
info.add_text("prompt", json.dumps({"1": {"class_type": "SaveImage", "inputs": {"filename_prefix": "fixes/x"}}}))
|
||||
Image.new("RGB", (4, 4)).save(root / "fixes" / "x_00001_.png", "PNG", pnginfo=info)
|
||||
sh.worklog.record("acme", kind="fix", prompt="x", refs=[], uploads=[], parents=[],
|
||||
output_prefix="fixes/x", pid="p1")
|
||||
sh.worklog.record("acme", kind="fix", prompt="pending", refs=[], uploads=[], parents=[],
|
||||
output_prefix="fixes/nope", pid="p2")
|
||||
client = await _client(_FakeServer())
|
||||
r = await client.get("/v1/workflow?id=p1", headers={"X-Org": "acme"})
|
||||
body = await r.json()
|
||||
assert r.status == 200 and body["format"] == "ui" and body["graph"] == ui # UI graph preferred
|
||||
# not landed yet -> 404 (no graph to open)
|
||||
assert (await client.get("/v1/workflow?id=p2", headers={"X-Org": "acme"})).status == 404
|
||||
# cross-tenant id is never confirmed -> 404
|
||||
assert (await client.get("/v1/workflow?id=p1", headers={"X-Org": "globex"})).status == 404
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wallet_proxies_and_scopes_org(tmp_path, monkeypatch):
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
async def fake_balance(tok):
|
||||
assert tok == "Bearer tkn"
|
||||
return {"availableCents": 2500, "currency": "usd", "pendingCents": 0}
|
||||
monkeypatch.setattr(sh, "_finance_balance", fake_balance)
|
||||
client = await _client(_FakeServer())
|
||||
r = await client.get("/v1/wallet", headers={"X-Org": "acme", "Authorization": "Bearer tkn"})
|
||||
body = await r.json()
|
||||
assert body == {"org": "acme", "balance": 25.0, "currency": "usd"}
|
||||
# signed-out (no bearer): display-only null balance, never a fabricated number
|
||||
r = await client.get("/v1/wallet", headers={"X-Org": "acme"})
|
||||
assert await r.json() == {"org": "acme", "balance": None, "currency": None}
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_library_upload_roundtrip_dedup_and_tenant_isolation(tmp_path, monkeypatch):
|
||||
import aiohttp
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
# library.json so _library_root resolves for the GET round-trip
|
||||
(Path(sh.folder_paths.get_org_output_directory("acme")) / "library.json").write_text('{"assets":[]}')
|
||||
(Path(sh.folder_paths.get_org_output_directory("globex")) / "library.json").write_text('{"assets":[]}')
|
||||
client = await _client(_FakeServer())
|
||||
png = b"\x89PNG\r\n\x1a\n" + b"acme-bytes"
|
||||
# multipart upload as acme -> renders/r1.png
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("image", png, filename="r1.png", content_type="image/png")
|
||||
r = await client.post("/v1/library/upload", data=form, headers={"X-Org": "acme"})
|
||||
body = await r.json()
|
||||
assert r.status == 200 and body["ok"] and body["path"] == "renders/r1.png" and not body.get("existed")
|
||||
# round-trip through the org-scoped file server
|
||||
g = await client.get("/v1/library/file?path=renders/r1.png", headers={"X-Org": "acme"})
|
||||
assert g.status == 200 and (await g.read()) == png
|
||||
# dedup: identical name+size -> existed, no rewrite
|
||||
form2 = aiohttp.FormData()
|
||||
form2.add_field("image", png, filename="r1.png", content_type="image/png")
|
||||
r2 = await client.post("/v1/library/upload", data=form2, headers={"X-Org": "acme"})
|
||||
assert (await r2.json()) == {"ok": True, "existed": True, "path": "renders/r1.png"}
|
||||
# tenant isolation: org B never sees org A's file; its own upload is a separate tree
|
||||
assert (await client.get("/v1/library/file?path=renders/r1.png", headers={"X-Org": "globex"})).status == 404
|
||||
rg = await client.post("/v1/library/upload?name=r1.png", data=b"\x89PNG\r\n\x1a\nglobex",
|
||||
headers={"X-Org": "globex", "Content-Type": "application/octet-stream"})
|
||||
assert (await rg.json())["path"] == "renders/r1.png"
|
||||
# org A's bytes are untouched by org B's same-named upload
|
||||
assert (await (await client.get("/v1/library/file?path=renders/r1.png", headers={"X-Org": "acme"})).read()) == png
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_library_upload_rejects_traversal_bad_name_and_oversize(tmp_path, monkeypatch):
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
client = await _client(_FakeServer())
|
||||
hdr = {"X-Org": "acme", "Content-Type": "application/octet-stream"}
|
||||
# subpath traversal -> 400
|
||||
assert (await client.post("/v1/library/upload?name=x.png&subpath=../secret", data=b"\x89PNGx", headers=hdr)).status == 400
|
||||
# non-image name -> 400
|
||||
assert (await client.post("/v1/library/upload?name=notes.txt", data=b"hi", headers=hdr)).status == 400
|
||||
# a traversal-y NAME is neutralized to its basename, lands safely under renders/
|
||||
r = await client.post("/v1/library/upload?name=../../evil.png", data=b"\x89PNG\r\n\x1a\nevil", headers=hdr)
|
||||
assert (await r.json())["path"] == "renders/evil.png"
|
||||
assert (Path(sh.folder_paths.get_org_output_directory("acme")) / "renders" / "evil.png").is_file()
|
||||
# size cap -> 413
|
||||
monkeypatch.setattr(sh, "_MAX_UPLOAD", 8)
|
||||
assert (await client.post("/v1/library/upload?name=big.png", data=b"0123456789", headers=hdr)).status == 413
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_library_upload_indexes_into_assets_and_worklog(tmp_path, monkeypatch):
|
||||
import aiohttp
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
client = await _client(_FakeServer())
|
||||
png = b"\x89PNG\r\n\x1a\n" + b"render-bytes"
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("image", png, filename="shot.png", content_type="image/png")
|
||||
r = await client.post("/v1/library/upload?node=spark&design=karma&kind=render",
|
||||
data=form, headers={"X-Org": "acme"})
|
||||
assert (await r.json())["path"] == "renders/shot.png"
|
||||
# instantly findable in Assets (library.json entry, draft)
|
||||
lib = await (await client.get("/v1/library", headers={"X-Org": "acme"})).json()
|
||||
a = next(x for x in lib["assets"] if x["path"] == "renders/shot.png")
|
||||
assert a["status"] == "draft" and a["kind"] == "render" and a["design"] == "karma" and a["tags"] == []
|
||||
# instantly findable in Queue & History with its source node + resolved output
|
||||
wl = await (await client.get("/v1/worklog", headers={"X-Org": "acme"})).json()
|
||||
it = next(x for x in wl["items"] if x["output_prefix"] == "renders/shot.png")
|
||||
assert it["kind"] == "render" and it["node"] == "spark" and it["status"] == "done"
|
||||
assert it["output"] == "renders/shot.png" and it["lane"] == "gpu"
|
||||
# tenant isolation: org B sees neither the asset nor the row
|
||||
lib2 = await (await client.get("/v1/library", headers={"X-Org": "globex"})).json()
|
||||
assert all(x["path"] != "renders/shot.png" for x in lib2.get("assets", []))
|
||||
assert (await (await client.get("/v1/worklog", headers={"X-Org": "globex"})).json())["items"] == []
|
||||
# dedup: re-upload identical bytes adds no second asset or row
|
||||
form2 = aiohttp.FormData()
|
||||
form2.add_field("image", png, filename="shot.png", content_type="image/png")
|
||||
r2 = await client.post("/v1/library/upload?node=spark", data=form2, headers={"X-Org": "acme"})
|
||||
assert (await r2.json()).get("existed") is True
|
||||
lib3 = await (await client.get("/v1/library", headers={"X-Org": "acme"})).json()
|
||||
assert sum(1 for x in lib3["assets"] if x["path"] == "renders/shot.png") == 1
|
||||
wl3 = await (await client.get("/v1/worklog", headers={"X-Org": "acme"})).json()
|
||||
assert sum(1 for x in wl3["items"] if x["output_prefix"] == "renders/shot.png") == 1
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_schedules_quality_gate_only_for_new_renders(tmp_path, monkeypatch):
|
||||
"""The gate fires at the ingest choke point exactly when a NEW asset row is
|
||||
created — a fresh render, or a restored-store file with no row yet — and never
|
||||
again for an identical re-post (dedup). Wiring only; the judgment runs
|
||||
out-of-band and is unit-tested against quality_gate.run."""
|
||||
import aiohttp
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
root = Path(sh.folder_paths.get_org_output_directory("acme"))
|
||||
(root / "library.json").write_text('{"assets":[]}')
|
||||
calls = []
|
||||
monkeypatch.setattr(sh.quality_gate, "schedule",
|
||||
lambda request, org, r, rel: calls.append((org, rel)))
|
||||
client = await _client(_FakeServer())
|
||||
png = b"\x89PNG\r\n\x1a\n" + b"gate-bytes"
|
||||
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("image", png, filename="g.png", content_type="image/png")
|
||||
r = await client.post("/v1/library/upload", data=form, headers={"X-Org": "acme"})
|
||||
assert (await r.json())["path"] == "renders/g.png"
|
||||
assert calls == [("acme", "renders/g.png")] # fresh render → judged once
|
||||
|
||||
form2 = aiohttp.FormData()
|
||||
form2.add_field("image", png, filename="g.png", content_type="image/png")
|
||||
r2 = await client.post("/v1/library/upload", data=form2, headers={"X-Org": "acme"})
|
||||
assert (await r2.json()).get("existed") is True
|
||||
assert calls == [("acme", "renders/g.png")] # identical re-post → NOT re-judged
|
||||
|
||||
# a file already on disk but with no library row (restored store) → a new row → judged
|
||||
png2 = b"\x89PNG\r\n\x1a\n" + b"restored"
|
||||
(root / "renders" / "old.png").write_bytes(png2)
|
||||
form3 = aiohttp.FormData()
|
||||
form3.add_field("image", png2, filename="old.png", content_type="image/png")
|
||||
r3 = await client.post("/v1/library/upload", data=form3, headers={"X-Org": "acme"})
|
||||
assert (await r3.json()).get("existed") is True # bytes already there (dedup write)
|
||||
assert calls == [("acme", "renders/g.png"), ("acme", "renders/old.png")]
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_status_route_accepts_flagged_and_restores(tmp_path, monkeypatch):
|
||||
"""The one lifecycle route gains ``flagged`` (so the gate's status is a first-
|
||||
class one) and drives the UI's Restore back to draft. /v1/library still returns
|
||||
every status — the UI hides flagged; the backend never filters it away."""
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
root = Path(sh.folder_paths.get_org_output_directory("acme"))
|
||||
(root / "a.png").write_bytes(b"\x89PNG")
|
||||
(root / "library.json").write_text(json.dumps({"assets": [{"path": "a.png", "status": "draft"}]}))
|
||||
client = await _client(_FakeServer())
|
||||
r = await client.post("/v1/library/status", json={"path": "a.png", "status": "flagged"}, headers={"X-Org": "acme"})
|
||||
assert r.status == 200 and (await r.json())["status"] == "flagged"
|
||||
lib = await (await client.get("/v1/library", headers={"X-Org": "acme"})).json()
|
||||
assert next(x for x in lib["assets"] if x["path"] == "a.png")["status"] == "flagged"
|
||||
r2 = await client.post("/v1/library/status", json={"path": "a.png", "status": "draft"}, headers={"X-Org": "acme"})
|
||||
assert r2.status == 200 and (await r2.json())["status"] == "draft"
|
||||
assert (await client.post("/v1/library/status", json={"path": "a.png", "status": "bogus"}, headers={"X-Org": "acme"})).status == 400
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_archive_records_negative_training_signal(tmp_path, monkeypatch):
|
||||
"""Curation IS training signal: archiving (status=deleted) or flagging an output
|
||||
records it as a bad response — vote -1 in ratings.json AND an event in feedback.jsonl
|
||||
carrying the deleted/flagged status — and restoring clears that AUTO-negative. The
|
||||
downvote is note-less, so it trains the reward dataset WITHOUT polluting prompt
|
||||
lessons (_mistakes needs note text)."""
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
root = Path(sh.folder_paths.get_org_output_directory("acme"))
|
||||
(root / "a.png").write_bytes(b"\x89PNG")
|
||||
(root / "library.json").write_text(json.dumps({"assets": [{"path": "a.png", "status": "draft"}]}))
|
||||
client = await _client(_FakeServer())
|
||||
|
||||
r = await client.post("/v1/library/status", json={"path": "a.png", "status": "deleted"}, headers={"X-Org": "acme"})
|
||||
assert r.status == 200
|
||||
ratings = json.loads((root / "ratings.json").read_text())
|
||||
assert ratings["a.png"]["vote"] == -1 and ratings["a.png"]["archived"] is True
|
||||
assert not ratings["a.png"].get("notes") # note-less: trains, never injects a lesson
|
||||
fb = [json.loads(x) for x in (root / "feedback.jsonl").read_text().splitlines() if x.strip()]
|
||||
ev = [e for e in fb if e.get("output") == "a.png"][-1]
|
||||
assert ev["vote"] == -1 and ev["status"] == "deleted" # deleted status rides along → negative-only
|
||||
assert sh._mistakes(ratings, {"a.png"}) == [] # a bare downvote never becomes a prompt lesson
|
||||
|
||||
r2 = await client.post("/v1/library/status", json={"path": "a.png", "status": "draft"}, headers={"X-Org": "acme"})
|
||||
assert r2.status == 200
|
||||
ratings2 = json.loads((root / "ratings.json").read_text())
|
||||
assert ratings2["a.png"]["vote"] == 0 and ratings2["a.png"]["archived"] is False # restore clears the auto-negative
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_archive_respects_a_user_judgement(tmp_path, monkeypatch):
|
||||
"""A downvote + note the user typed is user-owned: archiving never overwrites it and
|
||||
restoring never clears it (only the AUTO archive-negative is reversible). Their note
|
||||
stays an injected lesson."""
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
root = Path(sh.folder_paths.get_org_output_directory("acme"))
|
||||
(root / "b.png").write_bytes(b"\x89PNG")
|
||||
(root / "library.json").write_text(json.dumps({"assets": [{"path": "b.png", "status": "draft"}]}))
|
||||
client = await _client(_FakeServer())
|
||||
await client.post("/v1/library/rate", json={"key": "b.png", "path": "b.png", "vote": -1, "notes": "left strap missing"}, headers={"X-Org": "acme"})
|
||||
await client.post("/v1/library/status", json={"path": "b.png", "status": "deleted"}, headers={"X-Org": "acme"})
|
||||
await client.post("/v1/library/status", json={"path": "b.png", "status": "draft"}, headers={"X-Org": "acme"})
|
||||
ratings = json.loads((root / "ratings.json").read_text())
|
||||
assert ratings["b.png"]["vote"] == -1 # survives archive+restore
|
||||
assert ratings["b.png"]["notes"] == "left strap missing"
|
||||
assert sh._mistakes(ratings, {"b.png"}) == ["left strap missing"] # still a real lesson
|
||||
await client.close()
|
||||
|
||||
|
||||
def test_studio_home_hides_flagged_and_offers_review():
|
||||
"""UI contract: flagged is hidden from the normal grid exactly like deleted, and
|
||||
reachable for review + restore. /v1/library returns all statuses, so the grid is
|
||||
the filter — this locks the change I made to it."""
|
||||
html = sh._HTML.read_text()
|
||||
assert "a.status==='deleted'||a.status==='flagged'" in html # hidden from normal views
|
||||
assert "FILTER='flagged'" in html # the Flagged review filter
|
||||
assert "showDel?'Recover':'Restore'" in html # restore affordance
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_worklog_search_matches_paths_and_refs_not_just_prompt(tmp_path, monkeypatch):
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
(Path(sh.folder_paths.get_org_output_directory("acme")) / "library.json").write_text('{"assets":[]}')
|
||||
# an ingested render: empty prompt, findable only by its filename / ref
|
||||
sh.worklog.record("acme", kind="render", prompt="", refs=["srcphoto.png"], uploads=[],
|
||||
parents=[], output_prefix="renders/model_shot.png", pid="r1", status="done")
|
||||
sh.worklog.record("acme", kind="fix", prompt="brighten the sky", refs=[], uploads=[],
|
||||
parents=[], output_prefix="fixes/x", pid="r2")
|
||||
client = await _client(_FakeServer())
|
||||
|
||||
async def ids(q):
|
||||
r = await client.get("/v1/worklog?q=" + q, headers={"X-Org": "acme"})
|
||||
return [it["id"] for it in (await r.json())["items"]]
|
||||
|
||||
assert await ids("model_shot") == ["r1"] # by output filename, empty prompt
|
||||
assert await ids("srcphoto") == ["r1"] # by reference
|
||||
assert await ids("brighten") == ["r2"] # prompt search still works
|
||||
assert await ids("nomatch") == []
|
||||
await client.close()
|
||||
|
||||
|
||||
# ── Queue/History surface: favorites, stacks, ratings-with-notes, flow templates ──
|
||||
@pytest.mark.asyncio
|
||||
async def test_favorites_toggle_and_tenant_isolation(tmp_path, monkeypatch):
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
client = await _client(_FakeServer())
|
||||
# empty to start
|
||||
assert (await (await client.get("/v1/favorites", headers={"X-Org": "acme"})).json()) == {"org": "acme", "favorites": {}}
|
||||
# favorite a run (key = run id) — echoes back, stamps a ts
|
||||
r = await client.post("/v1/favorites", json={"key": "run-1", "favorite": True}, headers={"X-Org": "acme"})
|
||||
assert r.status == 200 and (await r.json()) == {"ok": True, "key": "run-1", "favorite": True}
|
||||
favs = (await (await client.get("/v1/favorites", headers={"X-Org": "acme"})).json())["favorites"]
|
||||
assert "run-1" in favs and isinstance(favs["run-1"]["ts"], int)
|
||||
# favorite defaults to True when omitted
|
||||
await client.post("/v1/favorites", json={"key": "run-2"}, headers={"X-Org": "acme"})
|
||||
assert set((await (await client.get("/v1/favorites", headers={"X-Org": "acme"})).json())["favorites"]) == {"run-1", "run-2"}
|
||||
# un-favorite removes it (the heart toggles both ways)
|
||||
await client.post("/v1/favorites", json={"key": "run-1", "favorite": False}, headers={"X-Org": "acme"})
|
||||
assert list((await (await client.get("/v1/favorites", headers={"X-Org": "acme"})).json())["favorites"]) == ["run-2"]
|
||||
# validation: empty key -> 400
|
||||
assert (await client.post("/v1/favorites", json={"key": ""}, headers={"X-Org": "acme"})).status == 400
|
||||
# tenant isolation: globex sees none of acme's favorites
|
||||
assert (await (await client.get("/v1/favorites", headers={"X-Org": "globex"})).json())["favorites"] == {}
|
||||
await client.close()
|
||||
|
||||
|
||||
# Stacks are now the per-org SQLite store (Procreate folders) — the full REST
|
||||
# surface, tenancy, one-stack-per-asset, order persistence and delete semantics
|
||||
# live in test_stacks.py, which supersedes the earlier JSON whole-set store.
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_stars_notes_path_and_tenant_isolation(tmp_path, monkeypatch):
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
# rate resolves via _library_root -> the org needs a library.json
|
||||
(Path(sh.folder_paths.get_org_output_directory("acme")) / "library.json").write_text('{"assets":[]}')
|
||||
(Path(sh.folder_paths.get_org_output_directory("globex")) / "library.json").write_text('{"assets":[]}')
|
||||
client = await _client(_FakeServer())
|
||||
# 0-5 stars + notes + the run's output path, keyed by run id (the flywheel row)
|
||||
r = await client.post("/v1/library/rate", json={
|
||||
"key": "run-1", "stars": 4, "notes": "hands look off", "path": "fixes/x_00001_.png"},
|
||||
headers={"X-Org": "acme"})
|
||||
assert r.status == 200
|
||||
entry = (await r.json())["entry"]
|
||||
assert entry["stars"] == 4 and entry["notes"] == "hands look off" and entry["path"] == "fixes/x_00001_.png"
|
||||
# visible on reopen; a note-only update keeps the stars (partial update)
|
||||
await client.post("/v1/library/rate", json={"key": "run-1", "notes": "better now"}, headers={"X-Org": "acme"})
|
||||
ratings = await (await client.get("/v1/library/ratings", headers={"X-Org": "acme"})).json()
|
||||
assert ratings["run-1"]["stars"] == 4 and ratings["run-1"]["notes"] == "better now"
|
||||
# a 0-star rating is valid (0-5 inclusive)
|
||||
assert (await client.post("/v1/library/rate", json={"key": "run-1", "stars": 0}, headers={"X-Org": "acme"})).status == 200
|
||||
# validation: out of range -> 400; missing key -> 400
|
||||
assert (await client.post("/v1/library/rate", json={"key": "r", "stars": 6}, headers={"X-Org": "acme"})).status == 400
|
||||
assert (await client.post("/v1/library/rate", json={"key": "", "stars": 3}, headers={"X-Org": "acme"})).status == 400
|
||||
# tenant isolation: globex never sees acme's rating
|
||||
assert (await (await client.get("/v1/library/ratings", headers={"X-Org": "globex"})).json()) == {}
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_template_accepts_flow_payload_not_just_graph(tmp_path, monkeypatch):
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
client = await _client(_FakeServer())
|
||||
# a flow template (no node graph) — what "favorite -> turn into a template" saves
|
||||
flow = {"kind": "fix", "prompt": "make the straps thicker", "refs": ["a.png"], "parents": ["a.png"]}
|
||||
r = await client.post("/v1/templates", json={"name": "Thicken Straps", "flow": flow}, headers={"X-Org": "acme"})
|
||||
assert r.status == 200 and (await r.json())["slug"] == "thicken-straps"
|
||||
# lists (not renderable — no runnable graph) and round-trips the captured flow
|
||||
lst = (await (await client.get("/v1/templates", headers={"X-Org": "acme"})).json())["templates"]
|
||||
assert lst[0]["name"] == "Thicken Straps" and lst[0]["renderable"] is False
|
||||
got = await (await client.get("/v1/templates/thicken-straps", headers={"X-Org": "acme"})).json()
|
||||
assert got["flow"] == flow and "graph" not in got
|
||||
# neither graph nor flow -> 400 (a template must capture something)
|
||||
assert (await client.post("/v1/templates", json={"name": "empty"}, headers={"X-Org": "acme"})).status == 400
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_critique_persists_thread_and_records_lesson(tmp_path, monkeypatch):
|
||||
"""A critique message: 200, replies (canned fallback when the vision call is
|
||||
unavailable), persists a user+assistant thread on the image, and writes the
|
||||
user's words as the note _mistakes reads — so critiquing trains the next render."""
|
||||
from middleware import quality_gate as qg
|
||||
monkeypatch.setattr(qg, "_downscale", lambda p: None) # skip the network vision call
|
||||
outd, _ = _orgaware(tmp_path, monkeypatch)
|
||||
root = Path(sh.folder_paths.get_org_output_directory("acme"))
|
||||
(root / "library.json").write_text('{"assets":[]}')
|
||||
(root / "fixes").mkdir(parents=True, exist_ok=True)
|
||||
from PIL import Image
|
||||
Image.new("RGB", (8, 8)).save(root / "fixes" / "x_00001_.png", "PNG")
|
||||
client = await _client(_FakeServer())
|
||||
rel = "fixes/x_00001_.png"
|
||||
r = await client.post("/v1/library/critique", json={"path": rel, "message": "skin is too orange"}, headers={"X-Org": "acme"})
|
||||
assert r.status == 200
|
||||
body = await r.json()
|
||||
assert body["reply"] and len(body["thread"]) == 2
|
||||
assert body["thread"][0]["role"] == "user" and body["thread"][1]["role"] == "assistant"
|
||||
# the note the flywheel reads carries the critique; reopening returns the thread
|
||||
ratings = await (await client.get("/v1/library/ratings", headers={"X-Org": "acme"})).json()
|
||||
assert "skin is too orange" in ratings[rel]["notes"]
|
||||
got = await (await client.get("/v1/library/critique?path=" + rel, headers={"X-Org": "acme"})).json()
|
||||
assert len(got["thread"]) == 2
|
||||
# guards: unknown asset 404, missing fields 400, tenant isolation
|
||||
assert (await client.post("/v1/library/critique", json={"path": "nope.png", "message": "x"}, headers={"X-Org": "acme"})).status == 404
|
||||
assert (await client.post("/v1/library/critique", json={"path": rel, "message": ""}, headers={"X-Org": "acme"})).status == 400
|
||||
other = await (await client.get("/v1/library/critique?path=" + rel, headers={"X-Org": "globex"})).json()
|
||||
assert other["thread"] == []
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_name_sanitized_against_html_injection(tmp_path, monkeypatch):
|
||||
"""A hostile upload filename (HTML/JS metacharacters) is collapsed to a safe
|
||||
charset at ingest — the stored path can never carry a stored-XSS payload."""
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
root = Path(sh.folder_paths.get_org_output_directory("acme"))
|
||||
(root / "library.json").write_text('{"assets":[]}')
|
||||
client = await _client(_FakeServer())
|
||||
png = b"\x89PNG\r\n\x1a\n" + b"\x00" * 64
|
||||
r = await client.post('/v1/library/upload?name=<img src=x onerror=alert(1)>.png', data=png, headers={"X-Org": "acme"})
|
||||
assert r.status == 200
|
||||
path = (await r.json())["path"]
|
||||
assert "<" not in path and ">" not in path and '"' not in path and "'" not in path
|
||||
assert path.endswith(".png")
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_thumbnail_serves_image_content_type(tmp_path, monkeypatch):
|
||||
"""The grid thumbnail MUST carry an image Content-Type. aiohttp's FileResponse
|
||||
guesses from the extension via mimetypes, which does not know .webp, so it served
|
||||
application/octet-stream — which the browser refuses to render (blank grid,
|
||||
'assets won't load'). Pin it."""
|
||||
from PIL import Image
|
||||
_orgaware(tmp_path, monkeypatch)
|
||||
root = Path(sh.folder_paths.get_org_output_directory("acme"))
|
||||
(root / "library.json").write_text('{"assets":[]}')
|
||||
(root / "renders").mkdir(parents=True, exist_ok=True)
|
||||
Image.new("RGB", (64, 96)).save(root / "renders" / "shot.png", "PNG")
|
||||
client = await _client(_FakeServer())
|
||||
r = await client.get("/v1/library/file?thumb=1&path=renders/shot.png", headers={"X-Org": "acme"})
|
||||
assert r.status == 200
|
||||
assert r.headers.get("Content-Type") == "image/webp" # NOT application/octet-stream
|
||||
await client.close()
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
"""Stored-XSS regression for the queue / worklog rows (middleware/studio_home.html).
|
||||
|
||||
A malicious job field — j.node (← X-Target-GPU), j.prefix/title (← SaveImage prefix),
|
||||
desc (← prompt text), j.id (← client prompt_id) — must render INERT. Two layers:
|
||||
(1) executable: the file's OWN esc() neutralizes an <img onerror> payload (run under
|
||||
node; skipped only if node is unavailable);
|
||||
(2) static: the row builders route every field through esc() and drive row actions via
|
||||
esc'd data-* attributes + one delegated listener — never a raw-concatenated onclick.
|
||||
"""
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
HTML = (Path(__file__).resolve().parents[2] / "middleware" / "studio_home.html").read_text()
|
||||
PAYLOAD = "<img src=x onerror=alert(1)>"
|
||||
|
||||
|
||||
def test_file_esc_map_neutralizes_payload():
|
||||
"""Apply the FILE's OWN esc() character map (extracted from studio_home.html) to the
|
||||
payload and assert it renders INERT — every '<'/'>' neutralized, no live tag survives.
|
||||
Pure-Python (no node dep) so it runs in CI; weakening the map (e.g. dropping '<')
|
||||
fails here."""
|
||||
m = re.search(r"const esc=.*?replace\(/\[([^\]]+)\]/g,\s*c=>\((\{[^}]+\})", HTML)
|
||||
assert m, "esc() helper/map not found"
|
||||
charclass, mapsrc = m.group(1), m.group(2)
|
||||
emap = dict(re.findall(r"'(.)':'([^']+)'", mapsrc))
|
||||
assert {"<", ">", "&", '"'} <= set(charclass) and {"<", ">"} <= emap.keys()
|
||||
out = "".join(emap.get(ch, ch) for ch in PAYLOAD)
|
||||
assert "<img" not in out and "<" not in out # every '<' neutralized
|
||||
assert emap["<"] in out # escaped entity present, inert text
|
||||
|
||||
|
||||
def test_gjRow_escapes_fields_and_delegates_cancel():
|
||||
assert "${esc(j.prefix||'render')}" in HTML # SaveImage prefix escaped
|
||||
assert "esc((j.node&&j.node!=='queued')?j.node:'your GPU')" in HTML # X-Target-GPU/identity escaped
|
||||
assert 'data-cancel="${esc(j.id)}"' in HTML # id via esc'd data attr
|
||||
|
||||
|
||||
def test_qitem_escapes_title_desc_and_delegates_cancel():
|
||||
assert "${esc(title)}" in HTML
|
||||
assert "${esc(desc)||'(no description)'}" in HTML
|
||||
assert 'data-cancel="${esc(id)}"' in HTML
|
||||
|
||||
|
||||
def test_no_raw_id_onclick_sinks_remain():
|
||||
for sink in ("cancelJob('${j.id}')", "cancelJob('${id}')", "cancelJob('${it.id}')",
|
||||
"openAmend('${it.id}')", "openContext('${it.id}')", "delWork('${it.id}')"):
|
||||
assert sink not in HTML, f"raw client-id onclick sink still present: {sink}"
|
||||
|
||||
|
||||
def test_one_delegated_listener_wires_data_actions():
|
||||
assert "t.closest('[data-cancel]')" in HTML
|
||||
assert "cancelJob(el.getAttribute('data-cancel'))" in HTML
|
||||
assert "openAmend(el.getAttribute('data-amend'))" in HTML
|
||||
Vendored
+76
-44
@@ -1,52 +1,57 @@
|
||||
import { useEffect, useRef, useState } from 'react'
|
||||
import {
|
||||
chat, getSession, getLibrary, getRenderQueue, gpuOnline, thumbURL,
|
||||
type ChatMessage, type Session, type Asset, type RenderJob,
|
||||
} from './api'
|
||||
import { chat, getSession, gpuOnline, type ChatMessage, type Session } from './api'
|
||||
import { Rail } from './rail'
|
||||
import { useLayout, startDrag, clamp, DEFAULT, RAIL_MIN, RAIL_MAX } from './layout'
|
||||
|
||||
// Studio-chat: the specialized chat surface. You say what to create; the cloud
|
||||
// /v1/chat "create" capability drives studio's render tools (fix/compose) and the
|
||||
// outputs land in the library strip. Increment 1 uses a minimal own-UI; the next
|
||||
// increment swaps the transcript/composer for @hanzo/chat's <Chat> on @hanzo/gui.
|
||||
// /v1/agent "create" preset drives studio's render tools (fix/compose) and the runs
|
||||
// land in the queue/history rail on the right — Up next while they render, History
|
||||
// once done. The main area is the conversation / current generation; the rail is the
|
||||
// whole YouTube-watch-page-style queue + history + stacks + ratings surface.
|
||||
|
||||
const EXAMPLES = [
|
||||
'make the Valentina back crotchless',
|
||||
'put design 13 on model 01',
|
||||
'brighten the studio lighting',
|
||||
'compose design 4 and design 9 into one look',
|
||||
]
|
||||
|
||||
export function App() {
|
||||
const [session, setSession] = useState<Session>({})
|
||||
const [messages, setMessages] = useState<ChatMessage[]>([])
|
||||
const [conversationId, setConversationId] = useState<string>()
|
||||
const [input, setInput] = useState('')
|
||||
const [busy, setBusy] = useState(false)
|
||||
const [assets, setAssets] = useState<Asset[]>([])
|
||||
const [jobs, setJobs] = useState<RenderJob[]>([])
|
||||
const [gpu, setGpu] = useState(false)
|
||||
const [acct, setAcct] = useState(false)
|
||||
const scroller = useRef<HTMLDivElement>(null)
|
||||
const composer = useRef<HTMLTextAreaElement>(null)
|
||||
|
||||
// Layout is keyed by org+user so each person's dividers/collapse/filter persist.
|
||||
const who = session.authenticated !== undefined ? `${session.org || 'default'}:${session.user || 'anon'}` : ''
|
||||
const { layout, patch, toggle, reset, beginDrag, endDrag } = useLayout(who)
|
||||
|
||||
useEffect(() => { getSession().then(setSession) }, [])
|
||||
|
||||
// Poll the library + in-flight queue so dispatched renders appear as they land.
|
||||
useEffect(() => {
|
||||
let alive = true
|
||||
const tick = async () => {
|
||||
if (!alive) return
|
||||
const [lib, q, on] = await Promise.all([getLibrary(), getRenderQueue(), gpuOnline()])
|
||||
if (!alive) return
|
||||
setAssets(lib.assets.filter(a => a.status !== 'deleted').slice(0, 24))
|
||||
setJobs(q.jobs); setGpu(on)
|
||||
}
|
||||
tick()
|
||||
const t = setInterval(tick, 5000)
|
||||
const tick = () => gpuOnline().then(on => { if (alive) setGpu(on) })
|
||||
tick(); const t = setInterval(tick, 15000)
|
||||
return () => { alive = false; clearInterval(t) }
|
||||
}, [])
|
||||
|
||||
useEffect(() => { scroller.current?.scrollTo(0, scroller.current.scrollHeight) }, [messages, busy])
|
||||
|
||||
const ready = !!input.trim() && !busy
|
||||
|
||||
async function send() {
|
||||
const text = input.trim()
|
||||
if (!text || busy) return
|
||||
const next = [...messages, { role: 'user' as const, content: text }]
|
||||
setMessages(next); setInput(''); setBusy(true)
|
||||
try {
|
||||
const r = await chat(next)
|
||||
const r = await chat(next, conversationId)
|
||||
setConversationId(r.conversationId)
|
||||
const note = r.actions.length
|
||||
? `\n\n${r.actions.map(a => a.error ? `⚠︎ ${a.name}: ${a.error}` : `▸ ${a.name} dispatched`).join('\n')}`
|
||||
? `\n\n${r.actions.map(a => a.error ? `⚠︎ ${a.name}: ${a.error}` : `▸ ${a.name} dispatched — see Up next →`).join('\n')}`
|
||||
: ''
|
||||
setMessages([...next, { role: 'assistant', content: (r.reply || '…') + note }])
|
||||
} catch (e) {
|
||||
@@ -56,13 +61,47 @@ export function App() {
|
||||
}
|
||||
}
|
||||
|
||||
function insertRef(text: string) {
|
||||
setInput(v => (v.trim() ? `${v.trim()} ${text}` : text) + ' ')
|
||||
composer.current?.focus()
|
||||
}
|
||||
function useExample(ex: string) {
|
||||
setInput(ex); composer.current?.focus()
|
||||
}
|
||||
|
||||
function onVSplit(e: React.PointerEvent) {
|
||||
beginDrag()
|
||||
const startW = layout.railW
|
||||
startDrag(e, 'x', d => patch({ railW: clamp(startW - d, RAIL_MIN, RAIL_MAX) }), endDrag)
|
||||
}
|
||||
|
||||
const initials = (session.user || session.org || session.email || '?').replace(/[^a-zA-Z0-9]/g, '').slice(0, 2).toUpperCase() || '?'
|
||||
|
||||
return (
|
||||
<div className="app">
|
||||
<header className="top">
|
||||
<div className="brand">Hanzo <b>Studio</b></div>
|
||||
<div className="spacer" />
|
||||
<div className={`gpu ${gpu ? 'on' : 'off'}`}>{gpu ? 'GPU online' : 'no GPU'}</div>
|
||||
<div className="who">{session.authenticated ? (session.org || session.user || 'signed in') : 'sign in'}</div>
|
||||
<div className="acct-wrap">
|
||||
<button className="avatar" aria-label="account" aria-haspopup="menu"
|
||||
title={session.email || session.user || 'account'} onClick={() => setAcct(a => !a)}>
|
||||
{initials}
|
||||
</button>
|
||||
{acct && (
|
||||
<div className="rc-menu acct-menu" role="menu" onMouseLeave={() => setAcct(false)}>
|
||||
<div className="acct-id">
|
||||
<div className="acct-name">{session.user || (session.authenticated ? 'signed in' : 'guest')}</div>
|
||||
{session.email && <div className="acct-sub">{session.email}</div>}
|
||||
{session.org && <div className="acct-sub">org · {session.org}</div>}
|
||||
</div>
|
||||
<a role="menuitem" href="https://console.hanzo.ai" target="_blank" rel="noreferrer">Console</a>
|
||||
{session.authenticated
|
||||
? <a role="menuitem" href="/logout">Sign out</a>
|
||||
: <a role="menuitem" href={session.iam_url || '/login'}>Sign in with Hanzo</a>}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</header>
|
||||
|
||||
<main className="body">
|
||||
@@ -71,8 +110,11 @@ export function App() {
|
||||
{messages.length === 0 && (
|
||||
<div className="hero">
|
||||
<h1>What should we create?</h1>
|
||||
<p>Describe an edit or a new shot — “make the Valentina back crotchless”,
|
||||
“put design 13 on model 01”. I’ll run it on your GPU and it lands below.</p>
|
||||
<p>Describe an edit or a new shot. Runs appear in <b>Up next</b> on the right as they render,
|
||||
and finished results land in <b>History</b> — rate them to teach the next pass.</p>
|
||||
<div className="examples">
|
||||
{EXAMPLES.map(ex => <button key={ex} className="ex-chip" onClick={() => useExample(ex)}>{ex}</button>)}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
{messages.map((m, i) => (
|
||||
@@ -82,30 +124,20 @@ export function App() {
|
||||
</div>
|
||||
<div className="composer">
|
||||
<textarea
|
||||
value={input} placeholder="Describe what to create or change…"
|
||||
ref={composer} value={input} placeholder="Describe what to create or change…"
|
||||
onChange={e => setInput(e.target.value)}
|
||||
onKeyDown={e => { if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); send() } }}
|
||||
/>
|
||||
<button onClick={send} disabled={busy || !input.trim()}>Send</button>
|
||||
<button className={`send${ready ? ' ready' : ''}`} onClick={send} disabled={busy}>Send</button>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
<aside className="rail">
|
||||
{jobs.length > 0 && (
|
||||
<div className="queue">
|
||||
<div className="rail-h">Rendering ({jobs.length})</div>
|
||||
{jobs.map(jb => <div key={jb.id} className="job">{jb.prefix || jb.id.slice(0, 8)}</div>)}
|
||||
</div>
|
||||
)}
|
||||
<div className="rail-h">Library</div>
|
||||
<div className="grid">
|
||||
{assets.map(a => (
|
||||
<a key={a.path} className="cell" href={thumbURL(a.path)} target="_blank" rel="noreferrer">
|
||||
<img loading="lazy" src={thumbURL(a.path)} alt="" />
|
||||
</a>
|
||||
))}
|
||||
{assets.length === 0 && <div className="empty">Renders appear here.</div>}
|
||||
</div>
|
||||
<div className="vsplit" onPointerDown={onVSplit} onDoubleClick={() => patch({ railW: DEFAULT.railW })}
|
||||
title="drag to resize · double-click to reset" role="separator" aria-orientation="vertical" />
|
||||
|
||||
<aside className="rail" style={{ width: layout.railW }}>
|
||||
<Rail layout={layout} patch={patch} toggle={toggle} reset={reset}
|
||||
beginDrag={beginDrag} endDrag={endDrag} onUseRef={insertRef} />
|
||||
</aside>
|
||||
</main>
|
||||
</div>
|
||||
|
||||
Vendored
+145
-23
@@ -1,20 +1,34 @@
|
||||
// The studio-chat data layer. Two backends, no proxy layer between them:
|
||||
// - the cloud chat orchestrator: POST api.hanzo.ai/v1/chat (clients/chat) —
|
||||
// called DIRECTLY, no studio pass-through. The hanzo_token cookie is scoped to
|
||||
// .hanzo.ai so credentials:'include' carries it cross-origin, and the handler
|
||||
// replays it into the per-org-billed completion. (Gateway CORS must allow the
|
||||
// studio.hanzo.ai origin with credentials.)
|
||||
// - the studio engine API: /v1/session, /v1/library, /v1/render-queue, … —
|
||||
// same-origin on studio.hanzo.ai (same cookie).
|
||||
// - the agent orchestrator: POST api.hanzo.ai/v1/agent (github.com/hanzoai/agent,
|
||||
// mounted in the cloud binary) — called DIRECTLY, no studio pass-through. The
|
||||
// hanzo_token cookie is scoped to .hanzo.ai so credentials:'include' carries it
|
||||
// cross-origin, and the round replays it into the per-org-billed completion.
|
||||
// (Gateway CORS must allow the studio.hanzo.ai origin with credentials.)
|
||||
// - the studio engine API: /v1/session, /v1/worklog, /v1/queue/*, /v1/favorites,
|
||||
// /v1/stacks, … — same-origin on studio.hanzo.ai (same cookie).
|
||||
// No token handling in the browser: the cookie is the credential for both.
|
||||
// VITE_STUDIO_API / VITE_CHAT_API override the bases for local dev.
|
||||
// VITE_STUDIO_API / VITE_AGENT_API override the bases for local dev.
|
||||
|
||||
const BASE = (import.meta.env.VITE_STUDIO_API ?? '').replace(/\/$/, '')
|
||||
const CHAT_API = (import.meta.env.VITE_CHAT_API ?? 'https://api.hanzo.ai').replace(/\/$/, '')
|
||||
// Default: same-origin through the studio's /v1/agent bridge — the host-only
|
||||
// session cookie authorizes it and the middleware forwards the bearer, so no
|
||||
// parent-domain cookie and no gateway CORS dependency. VITE_AGENT_API points
|
||||
// straight at the gateway for direct mode.
|
||||
const AGENT_API = (import.meta.env.VITE_AGENT_API ?? '').replace(/\/$/, '')
|
||||
|
||||
async function j<T>(url: string, init?: RequestInit): Promise<T> {
|
||||
const r = await fetch(url, { credentials: 'include', ...init })
|
||||
if (!r.ok) throw new Error(`${url} -> ${r.status}`)
|
||||
if (!r.ok) {
|
||||
// Surface the server's own message — a 402 says "add credits at
|
||||
// pay.hanzo.ai", which beats a bare status code.
|
||||
let msg = `${url} -> ${r.status}`
|
||||
try {
|
||||
const b = await r.json()
|
||||
const m = b?.error?.message ?? b?.error ?? b?.message
|
||||
if (typeof m === 'string' && m) msg = m
|
||||
} catch { /* keep the status line */ }
|
||||
throw new Error(msg)
|
||||
}
|
||||
return r.json() as Promise<T>
|
||||
}
|
||||
|
||||
@@ -22,15 +36,16 @@ export type Role = 'user' | 'assistant'
|
||||
export interface ChatMessage { role: Role; content: string }
|
||||
export interface Action { name: string; args?: Record<string, unknown>; result?: unknown; error?: string }
|
||||
export interface Op { name: string; args?: Record<string, unknown> }
|
||||
export interface ChatReply { reply: string; actions: Action[]; ops: Op[] }
|
||||
export interface ChatReply { reply: string; actions: Action[]; ops: Op[]; conversationId: string }
|
||||
|
||||
// One turn of the create capability: the model edits/combines library assets;
|
||||
// `actions` are the renders it dispatched (each result carries a prompt_id).
|
||||
export function chat(messages: ChatMessage[], capability = 'create'): Promise<ChatReply> {
|
||||
return j<ChatReply>(`${CHAT_API}/v1/chat`, {
|
||||
// One turn of a preset: the model edits/combines library assets; `actions` are the
|
||||
// renders it dispatched (each result carries a prompt_id). Pass the prior reply's
|
||||
// conversationId to continue a thread — the orchestrator persists per-org history.
|
||||
export function chat(messages: ChatMessage[], conversationId?: string, preset = 'create'): Promise<ChatReply> {
|
||||
return j<ChatReply>(`${AGENT_API}/v1/agent`, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ capability, messages }),
|
||||
body: JSON.stringify({ preset, messages, conversationId }),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -39,14 +54,121 @@ export interface Session {
|
||||
authenticated?: boolean; iam_url?: string
|
||||
}
|
||||
export const getSession = () => j<Session>(`${BASE}/v1/session`).catch(() => ({ authenticated: false } as Session))
|
||||
|
||||
export interface Asset { path: string; status?: string; mtime?: number; view?: unknown }
|
||||
export const getLibrary = () => j<{ org: string; assets: Asset[] }>(`${BASE}/v1/library`).catch(() => ({ org: '', assets: [] }))
|
||||
|
||||
export interface RenderJob { id: string; prefix?: string; refs?: string[]; ts?: number; age?: number; status?: string }
|
||||
export const getRenderQueue = () => j<{ jobs: RenderJob[] }>(`${BASE}/v1/render-queue`).catch(() => ({ jobs: [] }))
|
||||
export const gpuOnline = () => j<{ online: boolean }>(`${BASE}/v1/gpu-status`).then(r => r.online).catch(() => false)
|
||||
|
||||
// Thumbnail URL for a library asset (the studio pod serves a downsized JPEG).
|
||||
// A run in a flow: one dispatched generation (fix/compose/rerun/render/template) with
|
||||
// its instruction, references, parents (the assets it derived from), landed output and
|
||||
// real status. The queue/history backbone — GET /v1/worklog, newest first.
|
||||
export type RunStatus = 'queued' | 'running' | 'done' | 'failed' | 'cancelled'
|
||||
export interface Run {
|
||||
id: string; ts: number; kind: string; prompt: string
|
||||
refs: string[]; uploads: string[]; parents: string[]
|
||||
output_prefix: string; output: string | null
|
||||
lane: string; node: string; status: RunStatus
|
||||
started_at?: number; finished_at?: number; error?: string
|
||||
params?: Record<string, unknown>
|
||||
}
|
||||
export const getWorklog = (q?: string) =>
|
||||
j<{ org: string; items: Run[] }>(`${BASE}/v1/worklog${q ? `?q=${encodeURIComponent(q)}` : ''}`)
|
||||
.catch(() => ({ org: '', items: [] as Run[] }))
|
||||
|
||||
// Live position + ETA overlay for the ACTIVE runs, per lane, from the org's own
|
||||
// measured medians. Joined onto the worklog rows by id.
|
||||
export interface QueueItem {
|
||||
id: string; kind: string; prompt: string; status: string; node: string
|
||||
lane?: string; position?: number; of?: number
|
||||
eta_seconds?: number; wait_seconds?: number; elapsed_seconds?: number; remaining_seconds?: number | null
|
||||
}
|
||||
export interface QueueStatus { items: Record<string, QueueItem>; medians: Record<string, number>; lanes: Record<string, number>; now: number }
|
||||
export const getQueueStatus = () =>
|
||||
j<QueueStatus>(`${BASE}/v1/queue/status`).catch(() => ({ items: {}, medians: {}, lanes: {}, now: Math.floor(Date.now() / 1000) }))
|
||||
|
||||
// The version chain of one output — ancestors (what it derived from) + descendants
|
||||
// (later iterations). The iteration history of a design flow.
|
||||
export interface LineageStep { path: string; prompt: string; kind: string; ts: number }
|
||||
export interface Lineage { ancestors: LineageStep[]; node: string; descendants: LineageStep[] }
|
||||
export const getLineage = (path: string) =>
|
||||
j<Lineage>(`${BASE}/v1/worklog/lineage?path=${encodeURIComponent(path)}`)
|
||||
.catch(() => ({ ancestors: [], node: path, descendants: [] }))
|
||||
|
||||
// Delete a queued run (tenant-scoped cancel; a foreign id 404s).
|
||||
export const cancelRun = (id: string) =>
|
||||
j<{ ok: boolean; canceled: string }>(`${BASE}/v1/queue/cancel`, {
|
||||
method: 'POST', headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ prompt_id: id }),
|
||||
})
|
||||
|
||||
// Amend a still-queued run — the ONE way to change it (atomic cancel + re-dispatch
|
||||
// with the same instruction plus added references). Adding references IS an amend.
|
||||
export const amendRun = (id: string, instruction: string, uploads: string[], paths: string[] = []) =>
|
||||
j<{ ok: boolean; prompt_id: string; amended_from: string; refs: number }>(`${BASE}/v1/queue/amend`, {
|
||||
method: 'POST', headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ prompt_id: id, instruction, uploads, paths }),
|
||||
})
|
||||
|
||||
// Stage a reference image in the org input dir (the engine's upload route); returns
|
||||
// the flat name amend/compose reference by.
|
||||
export const uploadRef = (file: File) => {
|
||||
const fd = new FormData()
|
||||
fd.append('image', file, file.name)
|
||||
return j<{ name: string; subfolder: string; type: string }>(`${BASE}/upload/image`, { method: 'POST', body: fd })
|
||||
}
|
||||
|
||||
// Ratings feed the AI flywheel: 0-5 stars + a note, keyed by run id, carrying the
|
||||
// output path. Persisted per-org; visible on reopen.
|
||||
export interface Rating { stars?: number; notes?: string; path?: string; updatedAt?: number }
|
||||
export const getRatings = () => j<Record<string, Rating>>(`${BASE}/v1/library/ratings`).catch(() => ({}))
|
||||
export const rate = (key: string, patch: { stars?: number; notes?: string; path?: string }) =>
|
||||
j<{ ok: boolean; key: string; entry: Rating }>(`${BASE}/v1/library/rate`, {
|
||||
method: 'POST', headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ key, ...patch }),
|
||||
})
|
||||
|
||||
// Soft-delete a finished result — the library status lifecycle ('deleted'), the same
|
||||
// mechanism the assets view uses. Keyed by the run's output path.
|
||||
export const softDelete = (path: string) =>
|
||||
j<{ ok: boolean; path: string; status: string }>(`${BASE}/v1/library/status`, {
|
||||
method: 'POST', headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ path, status: 'deleted' }),
|
||||
})
|
||||
// The org's soft-deleted asset paths — so History hides results deleted here or in the
|
||||
// assets view, and the delete persists across reloads.
|
||||
export const getDeletedPaths = () =>
|
||||
j<{ assets: { path: string; status?: string }[] }>(`${BASE}/v1/library`)
|
||||
.then(r => (r.assets || []).filter(a => a.status === 'deleted').map(a => a.path))
|
||||
.catch(() => [] as string[])
|
||||
|
||||
// Favorites: the heart toggle, a per-org set keyed by run id. Separate concern from
|
||||
// ratings — a bookmark, not quality feedback.
|
||||
export const getFavorites = () =>
|
||||
j<{ org: string; favorites: Record<string, { ts: number }> }>(`${BASE}/v1/favorites`)
|
||||
.then(r => r.favorites).catch(() => ({} as Record<string, { ts: number }>))
|
||||
export const setFavorite = (key: string, favorite: boolean) =>
|
||||
j<{ ok: boolean; key: string; favorite: boolean }>(`${BASE}/v1/favorites`, {
|
||||
method: 'POST', headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ key, favorite }),
|
||||
})
|
||||
|
||||
// Stacks: iOS-folder groupings of runs. The client owns the arrangement; the whole
|
||||
// normalized set is POSTed back and echoed.
|
||||
export interface Stack { id: string; name: string; items: string[]; ts: number }
|
||||
export const getStacks = () =>
|
||||
j<{ org: string; stacks: Stack[] }>(`${BASE}/v1/stacks`).then(r => r.stacks).catch(() => [] as Stack[])
|
||||
export const saveStacks = (stacks: Stack[]) =>
|
||||
j<{ org: string; stacks: Stack[] }>(`${BASE}/v1/stacks`, {
|
||||
method: 'POST', headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ stacks }),
|
||||
}).then(r => r.stacks)
|
||||
|
||||
// Favorite → template: capture a reusable FLOW (kind + instruction + refs + parents).
|
||||
export interface Flow { kind: string; prompt: string; refs: string[]; parents: string[] }
|
||||
export const saveTemplate = (name: string, flow: Flow) =>
|
||||
j<{ ok: boolean; name: string; slug: string }>(`${BASE}/v1/templates`, {
|
||||
method: 'POST', headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ name, flow }),
|
||||
})
|
||||
|
||||
// Thumbnails: a landed library output, a staged input reference, and the full-res file.
|
||||
export const thumbURL = (path: string) => `${BASE}/v1/library/file?thumb=1&path=${encodeURIComponent(path)}`
|
||||
export const inputThumbURL = (name: string) => `${BASE}/v1/library/file?input=1&thumb=1&path=${encodeURIComponent(name)}`
|
||||
export const fileURL = (path: string) => `${BASE}/v1/library/file?path=${encodeURIComponent(path)}`
|
||||
|
||||
Vendored
+219
@@ -0,0 +1,219 @@
|
||||
import { useState } from 'react'
|
||||
import {
|
||||
type Run, type QueueItem, type Rating,
|
||||
thumbURL, inputThumbURL, fileURL,
|
||||
} from './api'
|
||||
|
||||
// ── Small shared bits ────────────────────────────────────────────────────────────
|
||||
export function relTime(ts: number, now: number): string {
|
||||
const s = Math.max(0, now - ts)
|
||||
if (s < 45) return 'just now'
|
||||
if (s < 3600) return `${Math.round(s / 60)}m ago`
|
||||
if (s < 86400) return `${Math.round(s / 3600)}h ago`
|
||||
return `${Math.round(s / 86400)}d ago`
|
||||
}
|
||||
export function fmtDur(sec: number): string {
|
||||
if (sec < 60) return `${Math.round(sec)}s`
|
||||
if (sec < 3600) return `${Math.round(sec / 60)}m`
|
||||
return `${(sec / 3600).toFixed(1)}h`
|
||||
}
|
||||
// A friendly identifier for a run, from its output/prefix (the finding: cards must
|
||||
// carry names like the "design 13" the prompts teach). Strips the _NNNNN_ counter.
|
||||
export function idLabel(run: Run): string {
|
||||
const raw = (run.output || run.output_prefix || run.kind || '').split('/').pop() || run.kind
|
||||
return raw.replace(/_\d+_?\.\w+$/, '').replace(/_\d+_$/, '').replace(/\.\w+$/, '') || run.kind
|
||||
}
|
||||
export function runThumb(run: Run): string | null {
|
||||
if (run.output) return thumbURL(run.output)
|
||||
if (run.parents[0]) return thumbURL(run.parents[0])
|
||||
if (run.refs[0]) return inputThumbURL(run.refs[0])
|
||||
return null
|
||||
}
|
||||
const isActive = (run: Run) => run.status === 'queued' || run.status === 'running'
|
||||
|
||||
export function Stars({ value, onChange }: { value: number; onChange?: (n: number) => void }) {
|
||||
const stars = [1, 2, 3, 4, 5]
|
||||
return (
|
||||
<span className={`stars${onChange ? ' interactive' : ''}`} role={onChange ? 'radiogroup' : undefined} aria-label="rating">
|
||||
{stars.map(n => (
|
||||
<button
|
||||
key={n} type="button" className={`star${n <= value ? ' on' : ''}`}
|
||||
disabled={!onChange} aria-label={`${n} star${n > 1 ? 's' : ''}`}
|
||||
onClick={onChange ? () => onChange(n === value ? 0 : n) : undefined}
|
||||
>★</button>
|
||||
))}
|
||||
</span>
|
||||
)
|
||||
}
|
||||
|
||||
// ── One run card ─────────────────────────────────────────────────────────────────
|
||||
export interface CardProps {
|
||||
run: Run
|
||||
qitem?: QueueItem
|
||||
rating?: Rating
|
||||
favorite: boolean
|
||||
now: number
|
||||
version?: { pos: number; len: number }
|
||||
onRemove: (run: Run) => void
|
||||
onAmend: (run: Run, files: File[]) => Promise<void>
|
||||
onRate: (run: Run, patch: { stars?: number; notes?: string; path?: string }) => void
|
||||
onFavorite: (run: Run) => void
|
||||
onAddToStack: (run: Run) => void
|
||||
onOpenVersions: (run: Run) => void
|
||||
onUseRef: (run: Run) => void
|
||||
dnd: {
|
||||
onDragStart: (e: React.DragEvent) => void
|
||||
onDragEnter: (e: React.DragEvent) => void
|
||||
onDragOver: (e: React.DragEvent) => void
|
||||
onDragLeave: (e: React.DragEvent) => void
|
||||
onDrop: (e: React.DragEvent) => void
|
||||
intent: boolean
|
||||
}
|
||||
}
|
||||
|
||||
export function RunCard(p: CardProps) {
|
||||
const { run, qitem, rating, favorite, now, version } = p
|
||||
const [menu, setMenu] = useState(false)
|
||||
const [panel, setPanel] = useState<'none' | 'rate' | 'refs'>('none')
|
||||
const [stars, setStars] = useState(rating?.stars ?? 0)
|
||||
const [notes, setNotes] = useState(rating?.notes ?? '')
|
||||
const [files, setFiles] = useState<File[]>([])
|
||||
const [busy, setBusy] = useState(false)
|
||||
const [fileOver, setFileOver] = useState(false)
|
||||
const [confirmDel, setConfirmDel] = useState(false)
|
||||
|
||||
const active = isActive(run)
|
||||
const src = runThumb(run)
|
||||
const fresh = run.status === 'done' && now - (run.finished_at || run.ts) < 120
|
||||
|
||||
// Corner badge: ETA while queued, elapsed while running, stars-or-node when done.
|
||||
let badge = ''
|
||||
if (run.status === 'running') badge = qitem?.elapsed_seconds != null ? fmtDur(qitem.elapsed_seconds) : 'running'
|
||||
else if (run.status === 'queued') badge = qitem?.eta_seconds != null ? `~${fmtDur(qitem.eta_seconds)}` : 'queued'
|
||||
else if (run.status === 'failed') badge = 'failed'
|
||||
else if (run.status === 'cancelled') badge = 'cancelled'
|
||||
else badge = rating?.stars ? `★${rating.stars}` : run.node
|
||||
|
||||
const pos = qitem?.position && qitem?.of ? `${qitem.position}/${qitem.of}` : ''
|
||||
const meta = [run.node, relTime(run.ts, now), run.status + (pos ? ` · ${pos}` : '')].filter(Boolean).join(' · ')
|
||||
|
||||
const collectFiles = (list: FileList | null) =>
|
||||
setFiles(f => [...f, ...[...(list || [])].filter(x => x.type.startsWith('image/'))])
|
||||
|
||||
function onDrop(e: React.DragEvent) {
|
||||
const dropped = [...(e.dataTransfer?.files || [])].filter(x => x.type.startsWith('image/'))
|
||||
if (dropped.length && active) { // files onto a queued card → add references
|
||||
e.preventDefault(); e.stopPropagation(); setFileOver(false)
|
||||
setPanel('refs'); setFiles(f => [...f, ...dropped]); return
|
||||
}
|
||||
setFileOver(false); p.dnd.onDrop(e) // otherwise an internal run drag → stacking
|
||||
}
|
||||
function onDragOver(e: React.DragEvent) {
|
||||
if (active && e.dataTransfer?.types?.includes('Files')) { e.preventDefault(); setFileOver(true); return }
|
||||
p.dnd.onDragOver(e) // internal run drag → let the stack drop fire
|
||||
}
|
||||
|
||||
async function applyRefs() {
|
||||
if (!files.length || busy) return
|
||||
setBusy(true)
|
||||
try { await p.onAmend(run, files); setFiles([]); setPanel('none') } finally { setBusy(false) }
|
||||
}
|
||||
function saveRate() {
|
||||
p.onRate(run, { stars, notes, path: run.output || undefined })
|
||||
setPanel('none')
|
||||
}
|
||||
// The × (and the menu's Delete) share ONE handler: a queued run cancels straight
|
||||
// away; a finished result is a destructive soft-delete, so it asks first.
|
||||
function requestDelete() {
|
||||
if (active) p.onRemove(run)
|
||||
else setConfirmDel(true)
|
||||
}
|
||||
|
||||
return (
|
||||
<article
|
||||
className={`rc rc-${run.status}${p.dnd.intent ? ' stack-intent' : ''}${fileOver ? ' file-over' : ''}`}
|
||||
draggable onDragStart={p.dnd.onDragStart}
|
||||
onDragEnter={p.dnd.onDragEnter} onDragLeave={e => { setFileOver(false); p.dnd.onDragLeave(e) }}
|
||||
onDragOver={onDragOver} onDrop={onDrop}
|
||||
>
|
||||
<a className="rc-thumb" draggable={false} href={run.output ? fileURL(run.output) : undefined}
|
||||
target="_blank" rel="noreferrer" onClick={e => { if (!run.output) e.preventDefault() }}>
|
||||
{src ? <img loading="lazy" draggable={false} src={src} alt="" /> : <span className="rc-ph">{run.kind || '—'}</span>}
|
||||
<span className={`rc-badge s-${run.status}`}>{badge}</span>
|
||||
{fresh && <span className="rc-new">New</span>}
|
||||
</a>
|
||||
|
||||
<button className="rc-x" aria-label={active ? 'cancel' : 'delete'}
|
||||
title={active ? 'Cancel' : 'Delete'} onClick={requestDelete}>×</button>
|
||||
{confirmDel && (
|
||||
<div className="rc-confirm" role="alertdialog" aria-label="confirm delete">
|
||||
<span>Delete result?</span>
|
||||
<button className="danger" onClick={() => { setConfirmDel(false); p.onRemove(run) }}>Delete</button>
|
||||
<button onClick={() => setConfirmDel(false)}>Keep</button>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="rc-body">
|
||||
<div className="rc-title" title={run.prompt || run.kind}>{run.prompt || `(${run.kind})`}</div>
|
||||
<div className="rc-meta">{meta}</div>
|
||||
<div className="rc-tags">
|
||||
<span className="rc-id">{idLabel(run)}</span>
|
||||
{version && version.len > 1 && (
|
||||
<button className="rc-vers" title="versions of this design" onClick={() => p.onOpenVersions(run)}>
|
||||
{Array.from({ length: Math.min(version.len, 6) }).map((_, i) =>
|
||||
<span key={i} className={`pip${i === version.pos - 1 ? ' on' : ''}`} />)}
|
||||
<span className="rc-vn">v{version.pos}</span>
|
||||
</button>
|
||||
)}
|
||||
{favorite && <span className="rc-fav" title="favorite">♥</span>}
|
||||
{rating?.stars ? <span className="rc-rated"><Stars value={rating.stars} /></span> : null}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<button className="rc-more" aria-label="more actions" aria-haspopup="menu"
|
||||
onClick={() => setMenu(m => !m)}>⋯</button>
|
||||
|
||||
{menu && (
|
||||
<div className="rc-menu" role="menu" onMouseLeave={() => setMenu(false)}>
|
||||
<button role="menuitem" onClick={() => { setMenu(false); requestDelete() }}>Delete</button>
|
||||
{active && <button role="menuitem" onClick={() => { setMenu(false); setPanel('refs') }}>Add references</button>}
|
||||
<button role="menuitem" onClick={() => { setMenu(false); setStars(rating?.stars ?? 0); setNotes(rating?.notes ?? ''); setPanel('rate') }}>Rate & note</button>
|
||||
<button role="menuitem" onClick={() => { setMenu(false); p.onFavorite(run) }}>{favorite ? 'Unfavorite' : 'Favorite'}</button>
|
||||
<button role="menuitem" onClick={() => { setMenu(false); p.onAddToStack(run) }}>Add to stack</button>
|
||||
<button role="menuitem" onClick={() => { setMenu(false); p.onUseRef(run) }}>Use as reference</button>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{panel === 'rate' && (
|
||||
<div className="rc-panel">
|
||||
<Stars value={stars} onChange={setStars} />
|
||||
<textarea className="rc-notes" value={notes} onChange={e => setNotes(e.target.value)}
|
||||
placeholder="What should the AI do differently next time?" />
|
||||
<div className="rc-panel-row">
|
||||
<button className="pri" onClick={saveRate}>Save</button>
|
||||
<button onClick={() => setPanel('none')}>Cancel</button>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{panel === 'refs' && (
|
||||
<div className={`rc-panel rc-drop${fileOver ? ' over' : ''}`}
|
||||
onDragOver={e => { e.preventDefault(); setFileOver(true) }}
|
||||
onDragLeave={() => setFileOver(false)}
|
||||
onDrop={e => { e.preventDefault(); e.stopPropagation(); setFileOver(false); collectFiles(e.dataTransfer.files) }}>
|
||||
<div className="rc-drop-hint">Drop images here, or</div>
|
||||
<label className="rc-pick">Choose files
|
||||
<input type="file" multiple accept="image/*" onChange={e => collectFiles(e.target.files)} />
|
||||
</label>
|
||||
{files.length > 0 && <div className="rc-files">{files.map((f, i) => <span key={i} className="chip">{f.name}</span>)}</div>}
|
||||
<div className="rc-panel-row">
|
||||
<button className="pri" disabled={!files.length || busy} onClick={applyRefs}>
|
||||
{busy ? 'Requeuing…' : `Add ${files.length || ''} & requeue`}
|
||||
</button>
|
||||
<button onClick={() => { setFiles([]); setPanel('none') }}>Cancel</button>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
</article>
|
||||
)
|
||||
}
|
||||
Vendored
+82
@@ -0,0 +1,82 @@
|
||||
// Layout state for the queue/history rail: divider positions, per-section collapse,
|
||||
// and the active filter chip — ONE mechanism, persisted in localStorage keyed by
|
||||
// org+user (from /v1/session) and restored on load. Dividers are native pointer
|
||||
// drags; nothing here depends on a DnD library.
|
||||
import { useCallback, useEffect, useRef, useState } from 'react'
|
||||
import type React from 'react'
|
||||
|
||||
export type Filter = 'all' | 'queued' | 'done' | 'favorites' | 'stacks'
|
||||
|
||||
export interface Layout {
|
||||
railW: number // px — the main | rail divider (horizontal resize)
|
||||
upFrac: number // 0..1 — the Up next | History divider (vertical resize)
|
||||
collapsed: { next: boolean; up: boolean; hist: boolean }
|
||||
filter: Filter
|
||||
}
|
||||
|
||||
export const DEFAULT: Layout = {
|
||||
railW: 360, upFrac: 0.5,
|
||||
collapsed: { next: false, up: false, hist: false }, filter: 'all',
|
||||
}
|
||||
export const RAIL_MIN = 300
|
||||
export const RAIL_MAX = 640
|
||||
|
||||
const KEY = (who: string) => `studio.layout:${who}`
|
||||
export const clamp = (v: number, lo: number, hi: number) => Math.min(hi, Math.max(lo, v))
|
||||
|
||||
export function useLayout(who: string) {
|
||||
const [layout, setLayout] = useState<Layout>(DEFAULT)
|
||||
const loaded = useRef(false)
|
||||
const dragging = useRef(false)
|
||||
|
||||
// Restore once we know who the user is (org:user); merge over defaults so a stored
|
||||
// layout from an older shape never loses a newly-added field.
|
||||
useEffect(() => {
|
||||
if (!who || loaded.current) return
|
||||
loaded.current = true
|
||||
try {
|
||||
const raw = localStorage.getItem(KEY(who))
|
||||
if (raw) {
|
||||
const s = JSON.parse(raw)
|
||||
setLayout({ ...DEFAULT, ...s, collapsed: { ...DEFAULT.collapsed, ...(s.collapsed || {}) } })
|
||||
}
|
||||
} catch { /* corrupt entry — fall back to defaults */ }
|
||||
}, [who])
|
||||
|
||||
// Persist on change — but not on every drag frame; the drag commits once on release.
|
||||
useEffect(() => {
|
||||
if (!who || !loaded.current || dragging.current) return
|
||||
try { localStorage.setItem(KEY(who), JSON.stringify(layout)) } catch { /* quota — ignore */ }
|
||||
}, [who, layout])
|
||||
|
||||
const patch = useCallback((p: Partial<Layout>) => setLayout(l => ({ ...l, ...p })), [])
|
||||
const toggle = useCallback((k: keyof Layout['collapsed']) =>
|
||||
setLayout(l => ({ ...l, collapsed: { ...l.collapsed, [k]: !l.collapsed[k] } })), [])
|
||||
const reset = useCallback(() => { setLayout(DEFAULT); dragging.current = false }, [])
|
||||
const beginDrag = useCallback(() => { dragging.current = true }, [])
|
||||
const endDrag = useCallback(() => { dragging.current = false; setLayout(l => ({ ...l })) }, [])
|
||||
|
||||
return { layout, patch, toggle, reset, beginDrag, endDrag }
|
||||
}
|
||||
|
||||
// Native divider drag. Tracks pointer delta from the grab point on the given axis and
|
||||
// hands it to `onMove` (the caller clamps + applies); `onEnd` commits/persists.
|
||||
export function startDrag(
|
||||
e: React.PointerEvent, axis: 'x' | 'y',
|
||||
onMove: (delta: number) => void, onEnd: () => void,
|
||||
) {
|
||||
e.preventDefault()
|
||||
const origin = axis === 'x' ? e.clientX : e.clientY
|
||||
const move = (ev: PointerEvent) => onMove((axis === 'x' ? ev.clientX : ev.clientY) - origin)
|
||||
const up = () => {
|
||||
window.removeEventListener('pointermove', move)
|
||||
window.removeEventListener('pointerup', up)
|
||||
document.body.style.userSelect = ''
|
||||
document.body.style.cursor = ''
|
||||
onEnd()
|
||||
}
|
||||
window.addEventListener('pointermove', move)
|
||||
window.addEventListener('pointerup', up)
|
||||
document.body.style.userSelect = 'none'
|
||||
document.body.style.cursor = axis === 'x' ? 'col-resize' : 'row-resize'
|
||||
}
|
||||
Vendored
+442
@@ -0,0 +1,442 @@
|
||||
import { useCallback, useEffect, useRef, useState } from 'react'
|
||||
import {
|
||||
getWorklog, getQueueStatus, getRatings, getFavorites, getStacks, saveStacks,
|
||||
cancelRun, amendRun, uploadRef, rate, setFavorite, saveTemplate, getLineage,
|
||||
softDelete, getDeletedPaths, thumbURL, fileURL,
|
||||
type Run, type QueueStatus, type QueueItem, type Rating, type Stack, type Lineage,
|
||||
} from './api'
|
||||
import { RunCard, runThumb, idLabel, relTime, fmtDur } from './card'
|
||||
import { type Filter, type Layout, startDrag, clamp } from './layout'
|
||||
|
||||
// The queue/history rail — the whole YouTube-watch-page surface adapted to studio's
|
||||
// generative flows. App owns the shell + composer + layout persistence; this owns all
|
||||
// queue/history/favorites/stacks data, polling and mutations, and calls back to insert
|
||||
// a reference into the composer.
|
||||
|
||||
// Version chain (pos/len) for each run's output, derived from the loaded rows — the
|
||||
// same parent→child edges the server's lineage walk uses, counted for the "vN" pips.
|
||||
function chains(rows: Run[]): Map<string, { pos: number; len: number }> {
|
||||
const producer = new Map<string, Run>()
|
||||
for (const r of rows) if (r.output && !producer.has(r.output)) producer.set(r.output, r)
|
||||
const childrenOf = new Map<string, Run[]>()
|
||||
for (const r of rows) for (const par of r.parents || []) {
|
||||
const arr = childrenOf.get(par)
|
||||
if (arr) arr.push(r); else childrenOf.set(par, [r])
|
||||
}
|
||||
const out = new Map<string, { pos: number; len: number }>()
|
||||
for (const r of rows) {
|
||||
if (!r.output) continue
|
||||
let anc = 0, cur: string | undefined = r.output
|
||||
const seen = new Set([r.output])
|
||||
for (;;) {
|
||||
const par: string | undefined = producer.get(cur!)?.parents?.[0]
|
||||
if (!par || seen.has(par)) break
|
||||
seen.add(par); anc++; cur = par
|
||||
}
|
||||
let desc = 0
|
||||
const q = [r.output], seenD = new Set([r.output])
|
||||
while (q.length) {
|
||||
for (const kid of childrenOf.get(q.shift()!) || []) {
|
||||
const key = kid.output || kid.id
|
||||
if (seenD.has(key)) continue
|
||||
seenD.add(key); desc++
|
||||
if (kid.output) q.push(kid.output)
|
||||
}
|
||||
}
|
||||
out.set(r.id, { pos: anc + 1, len: anc + 1 + desc })
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
export interface RailProps {
|
||||
layout: Layout
|
||||
patch: (p: Partial<Layout>) => void
|
||||
toggle: (k: 'next' | 'up' | 'hist') => void
|
||||
reset: () => void
|
||||
beginDrag: () => void
|
||||
endDrag: () => void
|
||||
onUseRef: (text: string) => void
|
||||
}
|
||||
|
||||
export function Rail(props: RailProps) {
|
||||
const { layout, patch, toggle, reset, beginDrag, endDrag } = props
|
||||
const filter = layout.filter
|
||||
|
||||
const [rows, setRows] = useState<Run[]>([])
|
||||
const [q, setQ] = useState<QueueStatus>({ items: {}, medians: {}, lanes: {}, now: Math.floor(Date.now() / 1000) })
|
||||
const [now, setNow] = useState(Math.floor(Date.now() / 1000))
|
||||
const [ratings, setRatings] = useState<Record<string, Rating>>({})
|
||||
const [favorites, setFavorites] = useState<Record<string, { ts: number }>>({})
|
||||
const [stacks, setStacks] = useState<Stack[]>([])
|
||||
const [deletedPaths, setDeletedPaths] = useState<Set<string>>(new Set())
|
||||
const [hidden, setHidden] = useState<Set<string>>(new Set())
|
||||
const [query, setQuery] = useState('')
|
||||
|
||||
const [openStack, setOpenStack] = useState<string | null>(null)
|
||||
const [stackPick, setStackPick] = useState<Run | null>(null)
|
||||
const [versionsFor, setVersionsFor] = useState<Run | null>(null)
|
||||
const [lineage, setLineage] = useState<Lineage | null>(null)
|
||||
const [offer, setOffer] = useState<Run | null>(null)
|
||||
const [railMenu, setRailMenu] = useState(false)
|
||||
const [err, setErr] = useState('')
|
||||
|
||||
const pollLive = useCallback(async () => {
|
||||
const [w, qs] = await Promise.all([getWorklog(query), getQueueStatus()])
|
||||
setRows(w.items); setQ(qs); setNow(qs.now)
|
||||
}, [query])
|
||||
const loadRatings = useCallback(() => getRatings().then(setRatings), [])
|
||||
const loadFavorites = useCallback(() => getFavorites().then(setFavorites), [])
|
||||
const loadStacks = useCallback(() => getStacks().then(setStacks), [])
|
||||
const loadDeleted = useCallback(() => getDeletedPaths().then(ps => setDeletedPaths(new Set(ps))), [])
|
||||
|
||||
useEffect(() => { pollLive(); const t = setInterval(pollLive, 4000); return () => clearInterval(t) }, [pollLive])
|
||||
useEffect(() => { loadRatings(); loadFavorites(); loadStacks(); loadDeleted() }, [loadRatings, loadFavorites, loadStacks, loadDeleted])
|
||||
useEffect(() => { if (!err) return; const t = setTimeout(() => setErr(''), 4000); return () => clearTimeout(t) }, [err])
|
||||
|
||||
// Active = the runs the server still counts as in-flight (present in queue/status),
|
||||
// plus just-dispatched rows not yet in the next poll — so a stalled row doesn't
|
||||
// linger in Up next forever (matches the queue's own 30-min active window).
|
||||
const active = rows.filter(r => (r.status === 'queued' || r.status === 'running') && (q.items[r.id] || now - r.ts <= 1800))
|
||||
.sort((a, b) => ((q.items[a.id]?.position ?? 999) - (q.items[b.id]?.position ?? 999)) || (a.ts - b.ts))
|
||||
const history = rows.filter(r => (r.status === 'done' || r.status === 'failed' || r.status === 'cancelled')
|
||||
&& !hidden.has(r.id) && !(r.output && deletedPaths.has(r.output)))
|
||||
const favRuns = rows.filter(r => favorites[r.id])
|
||||
const nextRun = active[0]
|
||||
const chainMap = chains(rows)
|
||||
|
||||
// ── Mutations (each refreshes the affected store) ────────────────────────────────
|
||||
// A dispatch/mutation can fail because the run or asset changed since this list was
|
||||
// fetched (deduped, finished, cancelled). The trailing refresh already re-fetches;
|
||||
// swap the raw error for a plain "pick it again" — one self-heal for every mutation.
|
||||
const STALE = /unknown asset|not found|no such job|already finished/i
|
||||
const heal = (e: Error): boolean => {
|
||||
if (!STALE.test(e.message)) return false
|
||||
setErr('That run changed — refreshed, pick it again.')
|
||||
return true
|
||||
}
|
||||
// The ONE remove handler behind both the card's × and its menu Delete: a queued run
|
||||
// cancels; a finished result soft-deletes its output (the library lifecycle) and is
|
||||
// hidden from History optimistically (reconciled from the server's deleted set).
|
||||
async function onRemove(run: Run) {
|
||||
if (run.status === 'queued' || run.status === 'running') {
|
||||
try { await cancelRun(run.id) } catch (e) { if (!heal(e as Error)) setErr((e as Error).message) }
|
||||
pollLive(); return
|
||||
}
|
||||
setHidden(h => new Set(h).add(run.id))
|
||||
const out = run.output
|
||||
if (out) {
|
||||
try { await softDelete(out); setDeletedPaths(d => new Set(d).add(out)) }
|
||||
catch (e) { if (!heal(e as Error)) setErr((e as Error).message) }
|
||||
}
|
||||
pollLive()
|
||||
}
|
||||
async function onAmend(run: Run, files: File[]) {
|
||||
const names: string[] = []
|
||||
for (const f of files) { try { const u = await uploadRef(f); if (u?.name) names.push(u.name) } catch (e) { setErr((e as Error).message) } }
|
||||
if (!names.length) return
|
||||
const instr = (run.prompt || '').trim()
|
||||
try { await amendRun(run.id, instr.length >= 3 ? instr : `${run.kind || 'edit'} update`, names) }
|
||||
catch (e) { if (!heal(e as Error)) setErr((e as Error).message) }
|
||||
pollLive()
|
||||
}
|
||||
async function onRate(run: Run, p: { stars?: number; notes?: string; path?: string }) {
|
||||
try { await rate(run.id, p) } catch (e) { setErr((e as Error).message) }
|
||||
loadRatings()
|
||||
}
|
||||
async function onFavorite(run: Run) {
|
||||
const nowFav = !favorites[run.id]
|
||||
try { await setFavorite(run.id, nowFav) } catch (e) { setErr((e as Error).message) }
|
||||
await loadFavorites()
|
||||
if (nowFav) setOffer(run) // offer to turn the flow into a template
|
||||
}
|
||||
async function onOpenVersions(run: Run) {
|
||||
setVersionsFor(run); setLineage(null)
|
||||
setLineage(await getLineage(run.output || ''))
|
||||
}
|
||||
async function persistStacks(next: Stack[]) {
|
||||
const clean = next.filter(s => s.items.length >= 2) // a folder of one dissolves
|
||||
setStacks(clean)
|
||||
try { setStacks(await saveStacks(clean)) } catch (e) { setErr((e as Error).message) }
|
||||
}
|
||||
function stackRuns(from: string, target: string) {
|
||||
const next = stacks.map(s => ({ ...s, items: s.items.filter(i => i !== from) }))
|
||||
const host = next.find(s => s.items.includes(target))
|
||||
if (host) host.items = [...host.items, from]
|
||||
else next.push({ id: `st_${Date.now().toString(36)}`, name: 'Stack', items: [target, from], ts: Math.floor(Date.now() / 1000) })
|
||||
persistStacks(next)
|
||||
}
|
||||
function addToStack(run: Run, stackId: string) {
|
||||
const next = stacks.map(s => ({ ...s, items: s.items.filter(i => i !== run.id) }))
|
||||
const host = next.find(s => s.id === stackId)
|
||||
if (host) host.items = [...host.items, run.id]
|
||||
persistStacks(next); setStackPick(null)
|
||||
}
|
||||
function unstack(run: Run) {
|
||||
persistStacks(stacks.map(s => ({ ...s, items: s.items.filter(i => i !== run.id) })))
|
||||
}
|
||||
function renameStack(id: string, name: string) {
|
||||
setStacks(s => s.map(x => x.id === id ? { ...x, name } : x))
|
||||
}
|
||||
function commitRename() { persistStacks(stacks) }
|
||||
|
||||
// ── Drag-hold stacking (native HTML5 drag + a 600ms hover timer) ─────────────────
|
||||
const dragId = useRef<string | null>(null)
|
||||
const holdTimer = useRef<number | undefined>(undefined)
|
||||
const [intentId, setIntentId] = useState<string | null>(null)
|
||||
function dndFor(run: Run) {
|
||||
return {
|
||||
onDragStart: (e: React.DragEvent) => { dragId.current = run.id; e.dataTransfer.setData('text/plain', run.id); e.dataTransfer.effectAllowed = 'move' },
|
||||
onDragOver: (e: React.DragEvent) => { if (dragId.current && dragId.current !== run.id) { e.preventDefault(); e.dataTransfer.dropEffect = 'move' } },
|
||||
onDragEnter: (e: React.DragEvent) => {
|
||||
if (!dragId.current || dragId.current === run.id) return
|
||||
e.preventDefault()
|
||||
window.clearTimeout(holdTimer.current)
|
||||
holdTimer.current = window.setTimeout(() => setIntentId(run.id), 600) // hold to form a stack
|
||||
},
|
||||
onDragLeave: () => { window.clearTimeout(holdTimer.current); setIntentId(cur => cur === run.id ? null : cur) },
|
||||
onDrop: () => {
|
||||
window.clearTimeout(holdTimer.current)
|
||||
const from = dragId.current, held = intentId === run.id
|
||||
setIntentId(null); dragId.current = null
|
||||
if (from && from !== run.id && held) stackRuns(from, run.id)
|
||||
},
|
||||
intent: intentId === run.id,
|
||||
}
|
||||
}
|
||||
|
||||
const card = (run: Run) => (
|
||||
<RunCard
|
||||
key={run.id} run={run} qitem={q.items[run.id]} rating={ratings[run.id]}
|
||||
favorite={!!favorites[run.id]} now={now} version={chainMap.get(run.id)}
|
||||
onRemove={onRemove} onAmend={onAmend} onRate={onRate} onFavorite={onFavorite}
|
||||
onAddToStack={r => setStackPick(r)} onOpenVersions={onOpenVersions}
|
||||
onUseRef={r => props.onUseRef(idLabel(r))} dnd={dndFor(run)}
|
||||
/>
|
||||
)
|
||||
|
||||
// ── Section layout (content-driven flex + the vertical resize divider) ───────────
|
||||
const sectionsRef = useRef<HTMLDivElement>(null)
|
||||
const bothOpen = filter === 'all' && !layout.collapsed.up && !layout.collapsed.hist
|
||||
const upStyle: React.CSSProperties = filter === 'queued' ? { flex: '1 1 auto' }
|
||||
: bothOpen ? { flex: `0 0 ${Math.round(layout.upFrac * 100)}%` }
|
||||
: layout.collapsed.up ? { flex: '0 0 auto' } : { flex: '1 1 auto' }
|
||||
const histStyle: React.CSSProperties = layout.collapsed.hist ? { flex: '0 0 auto' } : { flex: '1 1 0' }
|
||||
function onHSplit(e: React.PointerEvent) {
|
||||
beginDrag()
|
||||
const h = sectionsRef.current?.getBoundingClientRect().height || 1
|
||||
const start = layout.upFrac
|
||||
startDrag(e, 'y', d => patch({ upFrac: clamp(start + d / h, 0.15, 0.85) }), endDrag)
|
||||
}
|
||||
|
||||
const chips: { k: Filter; label: string; n?: number }[] = [
|
||||
{ k: 'all', label: 'All' },
|
||||
{ k: 'queued', label: 'Queued', n: active.length },
|
||||
{ k: 'done', label: 'Done' },
|
||||
{ k: 'favorites', label: 'Favorites', n: favRuns.length },
|
||||
{ k: 'stacks', label: 'Stacks', n: stacks.length },
|
||||
]
|
||||
|
||||
return (
|
||||
<>
|
||||
{(filter === 'all' || filter === 'queued') && nextRun && (
|
||||
<div className={`nextpin${layout.collapsed.next ? ' col' : ''}`}>
|
||||
<button className="pin-h" onClick={() => toggle('next')} aria-expanded={!layout.collapsed.next}>
|
||||
<span className="pin-lab">Next</span>
|
||||
<span className="pin-count">{active.length} in queue</span>
|
||||
<span className="chev">{layout.collapsed.next ? '▸' : '▾'}</span>
|
||||
</button>
|
||||
{!layout.collapsed.next && <NextBody run={nextRun} q={q.items[nextRun.id]} now={now} />}
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="chips">
|
||||
<div className="chips-row">
|
||||
{chips.map(c => (
|
||||
<button key={c.k} className={`chip-f${filter === c.k ? ' on' : ''}`} onClick={() => patch({ filter: c.k })}>
|
||||
{c.label}{c.n ? <span className="chip-n">{c.n}</span> : null}
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
<div className="rail-more-wrap">
|
||||
<button className="rail-more" aria-label="rail options" onClick={() => setRailMenu(m => !m)}>⋯</button>
|
||||
{railMenu && (
|
||||
<div className="rc-menu" role="menu" onMouseLeave={() => setRailMenu(false)}>
|
||||
<button role="menuitem" onClick={() => { setRailMenu(false); reset() }}>Reset layout</button>
|
||||
<button role="menuitem" onClick={() => { setRailMenu(false); pollLive(); loadRatings(); loadFavorites(); loadStacks() }}>Refresh</button>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="rail-main">
|
||||
{filter === 'stacks' ? (
|
||||
<div className="stacks-view">
|
||||
{stacks.length === 0 && <div className="empty">No stacks yet. Drag one card onto another and hold (~½s) to make a stack.</div>}
|
||||
{stacks.map(s => {
|
||||
const members = s.items.map(id => rows.find(r => r.id === id)).filter((r): r is Run => !!r)
|
||||
const open = openStack === s.id
|
||||
return (
|
||||
<div key={s.id} className={`stack${open ? ' open' : ''}`}>
|
||||
<button className="stack-tile" onClick={() => setOpenStack(open ? null : s.id)}>
|
||||
<span className="stack-cards">
|
||||
{members.slice(0, 3).map((m, i) => {
|
||||
const src = runThumb(m)
|
||||
return <span key={i} className="sc" style={src ? { backgroundImage: `url(${src})` } : undefined} />
|
||||
})}
|
||||
</span>
|
||||
<span className="stack-meta">
|
||||
<span className="stack-name">{s.name}</span>
|
||||
<span className="stack-count">{members.length}</span>
|
||||
</span>
|
||||
</button>
|
||||
{open && (
|
||||
<div className="stack-open">
|
||||
<input className="stack-rename" value={s.name} aria-label="stack name"
|
||||
onChange={e => renameStack(s.id, e.target.value)} onBlur={commitRename} />
|
||||
{members.map(m => (
|
||||
<div key={m.id} className="stack-member">
|
||||
{card(m)}
|
||||
<button className="unstack" title="remove from stack" onClick={() => unstack(m)}>×</button>
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)
|
||||
})}
|
||||
</div>
|
||||
) : filter === 'favorites' ? (
|
||||
<section className="sec">
|
||||
<div className="sec-body">
|
||||
{favRuns.length ? favRuns.map(card) : <div className="empty">No favorites yet. Tap ♥ on any card.</div>}
|
||||
</div>
|
||||
</section>
|
||||
) : (
|
||||
<div className="sections" ref={sectionsRef}>
|
||||
{(filter === 'all' || filter === 'queued') && (
|
||||
<section className={`sec${layout.collapsed.up ? ' col' : ''}`} style={upStyle}>
|
||||
<SecHeader label="Up next" n={active.length} collapsed={layout.collapsed.up} onToggle={() => toggle('up')} />
|
||||
{!layout.collapsed.up && (
|
||||
<div className="sec-body">{active.length ? active.map(card) : <div className="empty">Nothing queued. Describe a change to start a run.</div>}</div>
|
||||
)}
|
||||
</section>
|
||||
)}
|
||||
{bothOpen && <div className="hsplit" onPointerDown={onHSplit} onDoubleClick={() => patch({ upFrac: 0.5 })} title="drag to resize · double-click to reset" />}
|
||||
{(filter === 'all' || filter === 'done') && (
|
||||
<section className={`sec${layout.collapsed.hist ? ' col' : ''}`} style={histStyle}>
|
||||
<SecHeader label="History" n={history.length} collapsed={layout.collapsed.hist} onToggle={() => toggle('hist')}
|
||||
search={<input className="sec-search" placeholder="Search history…" value={query} onChange={e => setQuery(e.target.value)} />} />
|
||||
{!layout.collapsed.hist && (
|
||||
<div className="sec-body">{history.length ? history.map(card) : <div className="empty">No past generations{query ? ' match your search' : ' yet'}.</div>}</div>
|
||||
)}
|
||||
</section>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{stackPick && (
|
||||
<div className="ovl" onClick={() => setStackPick(null)}>
|
||||
<div className="sheet" onClick={e => e.stopPropagation()}>
|
||||
<div className="sheet-h">Add “{idLabel(stackPick)}” to a stack</div>
|
||||
{stacks.length === 0
|
||||
? <p className="muted">No stacks yet — drag one card onto another and hold to make the first one.</p>
|
||||
: stacks.map(s => <button key={s.id} className="sheet-item" onClick={() => addToStack(stackPick, s.id)}>{s.name} · {s.items.length}</button>)}
|
||||
<div className="sheet-row"><button onClick={() => setStackPick(null)}>Cancel</button></div>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{versionsFor && (
|
||||
<div className="ovl" onClick={() => setVersionsFor(null)}>
|
||||
<div className="sheet vers" onClick={e => e.stopPropagation()}>
|
||||
<div className="sheet-h">Versions of {idLabel(versionsFor)}</div>
|
||||
{!lineage ? <div className="empty">Loading…</div> : (
|
||||
<div className="vers-strip">
|
||||
{[...lineage.ancestors,
|
||||
{ path: lineage.node, prompt: versionsFor.prompt, kind: versionsFor.kind, ts: versionsFor.ts },
|
||||
...lineage.descendants].map((step, i) => (
|
||||
<a key={i} className={`vstep${step.path === lineage.node ? ' cur' : ''}`}
|
||||
href={step.path ? fileURL(step.path) : undefined} target="_blank" rel="noreferrer"
|
||||
onClick={e => { if (!step.path) e.preventDefault() }} title={step.prompt || step.kind}>
|
||||
{step.path ? <img src={thumbURL(step.path)} alt="" /> : <span className="rc-ph">?</span>}
|
||||
<span className="vlabel">v{i + 1}</span>
|
||||
</a>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
<div className="sheet-row"><button onClick={() => setVersionsFor(null)}>Close</button></div>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{offer && <TemplateOffer run={offer} onClose={() => setOffer(null)} onError={setErr} />}
|
||||
{err && <div className="rail-toast" role="alert">{err}</div>}
|
||||
</>
|
||||
)
|
||||
}
|
||||
|
||||
function NextBody({ run, q, now }: { run: Run; q?: QueueItem; now: number }) {
|
||||
const src = runThumb(run)
|
||||
const badge = run.status === 'running'
|
||||
? (q?.elapsed_seconds != null ? `running · ${fmtDur(q.elapsed_seconds)}` : 'running')
|
||||
: (q?.eta_seconds != null ? `starts in ~${fmtDur(q.eta_seconds)}` : 'queued')
|
||||
return (
|
||||
<div className="pin-body">
|
||||
<div className="pin-thumb">{src ? <img src={src} alt="" /> : <span className="rc-ph">{run.kind}</span>}</div>
|
||||
<div className="pin-txt">
|
||||
<div className="pin-prompt" title={run.prompt}>{run.prompt || `(${run.kind})`}</div>
|
||||
<div className="pin-meta">{badge} · {run.node} · {relTime(run.ts, now)}</div>
|
||||
</div>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
function SecHeader({ label, n, collapsed, onToggle, search }:
|
||||
{ label: string; n: number; collapsed: boolean; onToggle: () => void; search?: React.ReactNode }) {
|
||||
return (
|
||||
<div className="sec-h">
|
||||
<button className="sec-toggle" onClick={onToggle} aria-expanded={!collapsed}>
|
||||
<span className="chev">{collapsed ? '▸' : '▾'}</span>
|
||||
<span className="sec-lab">{label}</span>
|
||||
<span className="sec-n">{n}</span>
|
||||
</button>
|
||||
{search}
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
function TemplateOffer({ run, onClose, onError }: { run: Run; onClose: () => void; onError: (m: string) => void }) {
|
||||
const [stage, setStage] = useState<'offer' | 'name'>('offer')
|
||||
const [name, setName] = useState(idLabel(run))
|
||||
const [busy, setBusy] = useState(false)
|
||||
const [done, setDone] = useState(false)
|
||||
async function save() {
|
||||
setBusy(true)
|
||||
try {
|
||||
await saveTemplate(name.trim() || idLabel(run), { kind: run.kind, prompt: run.prompt, refs: run.refs, parents: run.parents })
|
||||
setDone(true); setTimeout(onClose, 900)
|
||||
} catch (e) { onError((e as Error).message); onClose() } finally { setBusy(false) }
|
||||
}
|
||||
return (
|
||||
<div className="ovl" onClick={onClose}>
|
||||
<div className="sheet" onClick={e => e.stopPropagation()}>
|
||||
{done ? <div className="sheet-h">Saved to templates ✓</div>
|
||||
: stage === 'offer' ? (
|
||||
<>
|
||||
<div className="sheet-h">Turn this into a template?</div>
|
||||
<p className="muted">Reuse this flow — {run.kind}, its instruction and {run.refs.length} reference{run.refs.length === 1 ? '' : 's'} — as a starting point next time.</p>
|
||||
<div className="sheet-row"><button className="pri" onClick={() => setStage('name')}>Name it…</button><button onClick={onClose}>Not now</button></div>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<div className="sheet-h">Name this template</div>
|
||||
<input className="sheet-in" value={name} autoFocus onChange={e => setName(e.target.value)} placeholder="Template name" />
|
||||
<p className="muted">Captures: {run.kind} · “{run.prompt || '(no instruction)'}” · {run.refs.length} ref{run.refs.length === 1 ? '' : 's'}</p>
|
||||
<div className="sheet-row"><button className="pri" disabled={busy || !name.trim()} onClick={save}>{busy ? 'Saving…' : 'Save template'}</button><button onClick={onClose}>Cancel</button></div>
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
Vendored
+217
-29
@@ -1,48 +1,236 @@
|
||||
:root {
|
||||
--bg: #0d0d0f; --panel: #151517; --line: #26262b; --text: #f4f4f6;
|
||||
--dim: #9a9aa2; --accent: #f4f4f6; --accent-ink: #0d0d0f;
|
||||
--bg: #0d0d0f; --panel: #151517; --panel2: #1c1c20; --line: #26262b; --text: #f4f4f6;
|
||||
--dim: #a6a6ae; --accent: #f4f4f6; --accent-ink: #0d0d0f; --brand: #4f8cff;
|
||||
--good: #35c26b; --bad: #ff6b6b; --halo: #4f8cff;
|
||||
}
|
||||
@media (prefers-color-scheme: light) {
|
||||
:root { --bg: #fafafa; --panel: #fff; --line: #e6e6ea; --text: #16161a; --dim: #6b6b73; --accent: #16161a; --accent-ink: #fff; }
|
||||
:root {
|
||||
--bg: #fafafa; --panel: #fff; --panel2: #f3f3f5; --line: #e6e6ea; --text: #16161a;
|
||||
--dim: #55565e; --accent: #16161a; --accent-ink: #fff; --brand: #2f6be0;
|
||||
--good: #1f9d57; --bad: #d33; --halo: #2f6be0;
|
||||
}
|
||||
}
|
||||
* { box-sizing: border-box; }
|
||||
html, body, #root { height: 100%; margin: 0; }
|
||||
body { background: var(--bg); color: var(--text); font: 15px/1.5 -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif; }
|
||||
body {
|
||||
background: var(--bg); color: var(--text);
|
||||
font: 15px/1.5 -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
|
||||
overflow-x: hidden;
|
||||
}
|
||||
button { font: inherit; color: inherit; }
|
||||
|
||||
.app { display: flex; flex-direction: column; height: 100%; }
|
||||
.top { display: flex; align-items: center; gap: 12px; padding: 12px 18px; border-bottom: 1px solid var(--line); }
|
||||
.top { display: flex; align-items: center; gap: 12px; padding: 10px 16px; border-bottom: 1px solid var(--line); }
|
||||
.brand { font-weight: 500; letter-spacing: .2px; } .brand b { font-weight: 700; }
|
||||
.spacer { flex: 1; }
|
||||
.gpu { font-size: 12px; padding: 3px 8px; border-radius: 999px; border: 1px solid var(--line); color: var(--dim); }
|
||||
.gpu.on { color: var(--text); } .gpu.on::before { content: "● "; color: #35c26b; }
|
||||
.who { font-size: 13px; color: var(--dim); }
|
||||
.gpu { font-size: 12px; padding: 3px 9px; border-radius: 999px; border: 1px solid var(--line); color: var(--dim); }
|
||||
.gpu.on { color: var(--text); } .gpu.on::before { content: "● "; color: var(--good); }
|
||||
|
||||
.body { flex: 1; display: grid; grid-template-columns: 1fr 320px; min-height: 0; }
|
||||
.chat { display: flex; flex-direction: column; min-height: 0; border-right: 1px solid var(--line); }
|
||||
.scroll { flex: 1; overflow-y: auto; padding: 22px; }
|
||||
.hero { max-width: 560px; margin: 8vh auto 0; text-align: center; }
|
||||
.hero h1 { font-size: 30px; font-weight: 600; margin: 0 0 10px; }
|
||||
.hero p { color: var(--dim); }
|
||||
/* Account avatar + menu — a 44px target, not bare text */
|
||||
.acct-wrap { position: relative; }
|
||||
.avatar {
|
||||
width: 44px; height: 44px; border-radius: 999px; border: 1px solid var(--line);
|
||||
background: var(--panel2); color: var(--text); font-weight: 700; font-size: 13px;
|
||||
cursor: pointer; letter-spacing: .5px;
|
||||
}
|
||||
.avatar:hover { border-color: var(--dim); }
|
||||
.acct-menu { right: 0; top: 50px; min-width: 210px; }
|
||||
.acct-id { padding: 8px 12px; border-bottom: 1px solid var(--line); }
|
||||
.acct-name { font-weight: 600; }
|
||||
.acct-sub { font-size: 12px; color: var(--dim); }
|
||||
|
||||
.body { flex: 1; display: flex; min-height: 0; }
|
||||
.chat { flex: 1 1 auto; min-width: 0; display: flex; flex-direction: column; }
|
||||
.scroll { flex: 1; overflow-y: auto; padding: 20px; }
|
||||
.hero { max-width: 580px; margin: 4vh auto 0; text-align: center; }
|
||||
.hero h1 { font-size: 27px; font-weight: 600; margin: 0 0 8px; }
|
||||
.hero p { color: var(--dim); margin: 0 0 16px; }
|
||||
.examples { display: flex; flex-wrap: wrap; gap: 8px; justify-content: center; }
|
||||
.ex-chip {
|
||||
border: 1px solid var(--line); background: var(--panel); color: var(--text);
|
||||
border-radius: 999px; padding: 7px 13px; font-size: 13px; cursor: pointer;
|
||||
}
|
||||
.ex-chip:hover { border-color: var(--brand); color: var(--brand); }
|
||||
.msg { display: flex; margin: 10px 0; } .msg.user { justify-content: flex-end; }
|
||||
.bubble { max-width: 74%; padding: 10px 14px; border-radius: 14px; white-space: pre-wrap; border: 1px solid var(--line); background: var(--panel); }
|
||||
.msg.user .bubble { background: var(--accent); color: var(--accent-ink); border-color: transparent; }
|
||||
.bubble.dim { color: var(--dim); }
|
||||
|
||||
.composer { display: flex; gap: 10px; padding: 14px; border-top: 1px solid var(--line); }
|
||||
.composer textarea { flex: 1; resize: none; height: 44px; padding: 11px 14px; border-radius: 12px; border: 1px solid var(--line); background: var(--panel); color: var(--text); font: inherit; }
|
||||
.composer button { padding: 0 18px; border-radius: 12px; border: 0; background: var(--accent); color: var(--accent-ink); font-weight: 600; cursor: pointer; }
|
||||
.composer button:disabled { opacity: .5; cursor: default; }
|
||||
|
||||
.rail { overflow-y: auto; padding: 16px; }
|
||||
.rail-h { font-size: 12px; text-transform: uppercase; letter-spacing: .6px; color: var(--dim); margin: 14px 0 8px; }
|
||||
.queue .job { font-size: 12px; padding: 6px 8px; border: 1px solid var(--line); border-radius: 8px; margin-bottom: 6px; color: var(--dim); }
|
||||
.grid { display: grid; grid-template-columns: repeat(3, 1fr); gap: 6px; }
|
||||
.cell { aspect-ratio: 3 / 4; overflow: hidden; border-radius: 8px; border: 1px solid var(--line); background: var(--panel); }
|
||||
.cell img { width: 100%; height: 100%; object-fit: cover; display: block; }
|
||||
.empty { color: var(--dim); font-size: 13px; }
|
||||
.composer { display: flex; gap: 10px; padding: 12px 14px; border-top: 1px solid var(--line); background: var(--bg); }
|
||||
.composer textarea { flex: 1; resize: none; height: 46px; padding: 12px 14px; border-radius: 12px; border: 1px solid var(--line); background: var(--panel); color: var(--text); font: inherit; }
|
||||
.composer textarea:focus { outline: none; border-color: var(--brand); }
|
||||
.send {
|
||||
padding: 0 20px; border-radius: 12px; cursor: pointer; font-weight: 600;
|
||||
border: 1px solid var(--line); background: transparent; color: var(--dim); /* empty: muted outline */
|
||||
}
|
||||
.send.ready { border-color: transparent; background: var(--brand); color: #fff; } /* text present: brand fill */
|
||||
.send:disabled { cursor: default; opacity: .6; }
|
||||
|
||||
/* ── The main | rail divider ─────────────────────────────────────────────── */
|
||||
.vsplit { flex: 0 0 6px; cursor: col-resize; background: var(--line); opacity: .5; }
|
||||
.vsplit:hover { opacity: 1; background: var(--brand); }
|
||||
|
||||
/* ── The rail ────────────────────────────────────────────────────────────── */
|
||||
.rail { border-left: 1px solid var(--line); display: flex; flex-direction: column; min-height: 0; overflow: hidden; background: var(--bg); }
|
||||
.chev { color: var(--dim); font-size: 11px; }
|
||||
|
||||
.nextpin { border-bottom: 1px solid var(--line); }
|
||||
.pin-h { width: 100%; display: flex; align-items: center; gap: 8px; padding: 9px 14px; background: var(--panel2); border: 0; cursor: pointer; text-align: left; }
|
||||
.pin-lab { font-weight: 700; font-size: 12px; text-transform: uppercase; letter-spacing: .6px; }
|
||||
.pin-count { flex: 1; color: var(--dim); font-size: 12px; }
|
||||
.pin-body { display: flex; gap: 10px; padding: 10px 14px; }
|
||||
.pin-thumb { flex: 0 0 64px; height: 64px; border-radius: 8px; overflow: hidden; border: 1px solid var(--line); background: var(--panel); display: grid; place-items: center; }
|
||||
.pin-thumb img { width: 100%; height: 100%; object-fit: cover; }
|
||||
.pin-txt { min-width: 0; }
|
||||
.pin-prompt { font-weight: 500; display: -webkit-box; -webkit-line-clamp: 2; -webkit-box-orient: vertical; overflow: hidden; }
|
||||
.pin-meta { font-size: 12px; color: var(--dim); margin-top: 3px; }
|
||||
|
||||
.chips { display: flex; align-items: center; gap: 8px; padding: 8px 12px; border-bottom: 1px solid var(--line); }
|
||||
.chips-row { flex: 1; display: flex; gap: 6px; overflow-x: auto; scrollbar-width: none; }
|
||||
.chips-row::-webkit-scrollbar { display: none; }
|
||||
.chip-f { flex: 0 0 auto; border: 1px solid var(--line); background: var(--panel); color: var(--dim); border-radius: 999px; padding: 5px 11px; font-size: 13px; cursor: pointer; display: inline-flex; align-items: center; gap: 5px; }
|
||||
.chip-f.on { background: var(--accent); color: var(--accent-ink); border-color: transparent; }
|
||||
.chip-n { font-size: 11px; background: color-mix(in srgb, var(--dim) 22%, transparent); border-radius: 999px; padding: 0 6px; }
|
||||
.chip-f.on .chip-n { background: color-mix(in srgb, var(--accent-ink) 25%, transparent); }
|
||||
.rail-more-wrap { position: relative; }
|
||||
.rail-more, .rc-more { border: 0; background: transparent; cursor: pointer; color: var(--dim); font-size: 18px; line-height: 1; padding: 4px 6px; border-radius: 6px; }
|
||||
.rail-more:hover { background: var(--panel2); color: var(--text); }
|
||||
|
||||
.rail-main { flex: 1; min-height: 0; display: flex; flex-direction: column; overflow: hidden; }
|
||||
.sections { flex: 1; min-height: 0; display: flex; flex-direction: column; overflow: hidden; }
|
||||
.sec { display: flex; flex-direction: column; min-height: 0; overflow: hidden; }
|
||||
.sec.col { flex: 0 0 auto !important; }
|
||||
.sec-h { display: flex; align-items: center; gap: 8px; padding: 8px 12px; border-bottom: 1px solid var(--line); }
|
||||
.sec-toggle { flex: 0 0 auto; display: flex; align-items: center; gap: 7px; background: transparent; border: 0; cursor: pointer; padding: 2px 4px; }
|
||||
.sec-lab { font-size: 12px; text-transform: uppercase; letter-spacing: .6px; font-weight: 700; }
|
||||
.sec-n { font-size: 11px; color: var(--dim); background: var(--panel2); border-radius: 999px; padding: 0 7px; }
|
||||
.sec-search { flex: 1; min-width: 0; border: 1px solid var(--line); background: var(--panel); color: var(--text); border-radius: 8px; padding: 5px 9px; font-size: 13px; }
|
||||
.sec-search:focus { outline: none; border-color: var(--brand); }
|
||||
.sec-body { flex: 1; min-height: 0; overflow-y: auto; padding: 8px; display: flex; flex-direction: column; gap: 8px; }
|
||||
.hsplit { flex: 0 0 6px; cursor: row-resize; background: var(--line); opacity: .5; }
|
||||
.hsplit:hover { opacity: 1; background: var(--brand); }
|
||||
.empty { color: var(--dim); font-size: 13px; padding: 14px; text-align: center; }
|
||||
|
||||
/* ── Run card ────────────────────────────────────────────────────────────── */
|
||||
.rc { position: relative; display: flex; flex-wrap: wrap; gap: 10px; padding: 8px; border: 1px solid var(--line); border-radius: 12px; background: var(--panel); }
|
||||
.rc:hover { border-color: var(--dim); }
|
||||
.rc.stack-intent { border-color: var(--halo); box-shadow: 0 0 0 3px color-mix(in srgb, var(--halo) 35%, transparent); }
|
||||
.rc.file-over { border-color: var(--brand); border-style: dashed; }
|
||||
.rc-thumb { position: relative; flex: 0 0 48%; max-width: 48%; border-radius: 8px; overflow: hidden; background: var(--panel2); border: 1px solid var(--line); display: block; aspect-ratio: 4 / 5; text-decoration: none; }
|
||||
.rc-thumb img { width: 100%; height: 100%; object-fit: cover; display: block; }
|
||||
.rc-ph { display: grid; place-items: center; width: 100%; height: 100%; color: var(--dim); font-size: 12px; }
|
||||
.rc-badge { position: absolute; left: 5px; bottom: 5px; font-size: 11px; font-weight: 600; padding: 1px 6px; border-radius: 6px; background: color-mix(in srgb, #000 62%, transparent); color: #fff; }
|
||||
.rc-badge.s-running { background: var(--good); }
|
||||
.rc-badge.s-failed { background: var(--bad); }
|
||||
.rc-new { position: absolute; right: 5px; top: 5px; font-size: 10px; font-weight: 700; text-transform: uppercase; letter-spacing: .4px; padding: 1px 6px; border-radius: 6px; background: var(--brand); color: #fff; }
|
||||
.rc-body { flex: 1; min-width: 0; display: flex; flex-direction: column; gap: 3px; }
|
||||
.rc-title { font-weight: 500; display: -webkit-box; -webkit-line-clamp: 2; -webkit-box-orient: vertical; overflow: hidden; line-height: 1.35; }
|
||||
.rc-meta { font-size: 12px; color: var(--dim); }
|
||||
.rc-tags { display: flex; align-items: center; gap: 6px; flex-wrap: wrap; margin-top: auto; }
|
||||
.rc-id { font-size: 11px; color: var(--dim); background: var(--panel2); border-radius: 6px; padding: 1px 6px; max-width: 120px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
|
||||
.rc-vers { display: inline-flex; align-items: center; gap: 3px; border: 1px solid var(--line); background: transparent; border-radius: 999px; padding: 1px 7px 1px 5px; cursor: pointer; }
|
||||
.rc-vers .pip { width: 5px; height: 5px; border-radius: 999px; background: var(--dim); opacity: .5; }
|
||||
.rc-vers .pip.on { background: var(--brand); opacity: 1; }
|
||||
.rc-vn { font-size: 11px; color: var(--dim); }
|
||||
.rc-fav { color: var(--bad); }
|
||||
.rc-more { position: absolute; right: 4px; top: 4px; }
|
||||
.rc-more:hover { background: var(--panel2); color: var(--text); }
|
||||
/* Always-visible fast-path delete/cancel (× top-left of the thumb) */
|
||||
.rc-x { position: absolute; left: 6px; top: 6px; z-index: 4; width: 22px; height: 22px; border-radius: 999px; border: 0; background: color-mix(in srgb, #000 55%, transparent); color: #fff; font-size: 16px; line-height: 1; cursor: pointer; display: grid; place-items: center; opacity: .9; }
|
||||
.rc-x:hover { background: var(--bad); opacity: 1; }
|
||||
.rc-confirm { position: absolute; left: 6px; top: 6px; z-index: 7; display: flex; align-items: center; gap: 6px; background: var(--panel); border: 1px solid var(--line); border-radius: 9px; padding: 5px 7px; box-shadow: 0 6px 20px rgba(0,0,0,.28); font-size: 12px; }
|
||||
.rc-confirm button { border: 1px solid var(--line); background: var(--panel2); color: var(--text); border-radius: 6px; padding: 3px 9px; font-size: 12px; cursor: pointer; }
|
||||
.rc-confirm .danger { background: var(--bad); color: #fff; border-color: transparent; }
|
||||
|
||||
.rc-menu { position: absolute; z-index: 20; right: 6px; top: 30px; background: var(--panel); border: 1px solid var(--line); border-radius: 10px; box-shadow: 0 8px 28px rgba(0,0,0,.22); padding: 5px; display: flex; flex-direction: column; min-width: 168px; }
|
||||
.rc-menu > button, .rc-menu > a { text-align: left; background: transparent; border: 0; border-radius: 7px; padding: 8px 10px; cursor: pointer; color: var(--text); text-decoration: none; font-size: 14px; }
|
||||
.rc-menu > button:hover, .rc-menu > a:hover { background: var(--panel2); }
|
||||
|
||||
.rc-panel { flex-basis: 100%; order: 9; border-top: 1px solid var(--line); margin-top: 6px; padding-top: 8px; display: flex; flex-direction: column; gap: 8px; }
|
||||
.rc-notes { width: 100%; min-height: 56px; resize: vertical; border: 1px solid var(--line); background: var(--panel2); color: var(--text); border-radius: 8px; padding: 8px 10px; font: inherit; font-size: 13px; }
|
||||
.rc-notes:focus { outline: none; border-color: var(--brand); }
|
||||
.rc-panel-row { display: flex; gap: 8px; }
|
||||
.rc-panel-row .pri, .sheet .pri, .send.ready { }
|
||||
.rc-panel-row button, .sheet-row button, .sheet-item { border: 1px solid var(--line); background: var(--panel); color: var(--text); border-radius: 8px; padding: 7px 13px; cursor: pointer; font-size: 13px; }
|
||||
.rc-panel-row .pri, .sheet-row .pri { background: var(--brand); color: #fff; border-color: transparent; }
|
||||
.rc-panel-row button:disabled, .sheet-row button:disabled { opacity: .5; cursor: default; }
|
||||
.rc-drop { align-items: stretch; text-align: center; }
|
||||
.rc-drop.over { outline: 2px dashed var(--brand); outline-offset: 2px; border-radius: 8px; }
|
||||
.rc-drop-hint { font-size: 12px; color: var(--dim); }
|
||||
.rc-pick { display: inline-block; border: 1px solid var(--line); border-radius: 8px; padding: 6px 12px; cursor: pointer; font-size: 13px; align-self: center; }
|
||||
.rc-pick input { display: none; }
|
||||
.rc-files { display: flex; flex-wrap: wrap; gap: 5px; }
|
||||
.chip { font-size: 11px; background: var(--panel2); border: 1px solid var(--line); border-radius: 6px; padding: 2px 7px; max-width: 100%; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
|
||||
|
||||
/* stars */
|
||||
.stars { display: inline-flex; gap: 1px; }
|
||||
.stars .star { border: 0; background: transparent; padding: 0 1px; cursor: default; color: var(--dim); font-size: 15px; line-height: 1; }
|
||||
.stars.interactive .star { cursor: pointer; font-size: 20px; }
|
||||
.stars .star.on { color: #f5b301; }
|
||||
.rc-rated .stars .star { font-size: 12px; }
|
||||
|
||||
/* ── Stacks ──────────────────────────────────────────────────────────────── */
|
||||
.stacks-view { flex: 1; min-height: 0; overflow-y: auto; padding: 10px; display: grid; grid-template-columns: repeat(auto-fill, minmax(150px, 1fr)); gap: 12px; align-content: start; }
|
||||
.stack { grid-column: 1 / -1; }
|
||||
.stack:not(.open) { grid-column: auto; }
|
||||
.stack-tile { width: 100%; border: 1px solid var(--line); background: var(--panel); border-radius: 12px; padding: 10px; cursor: pointer; display: flex; flex-direction: column; gap: 8px; }
|
||||
.stack-tile:hover { border-color: var(--dim); }
|
||||
.stack-cards { position: relative; height: 92px; }
|
||||
.stack-cards .sc { position: absolute; width: 66%; height: 84px; border-radius: 8px; border: 1px solid var(--line); background: var(--panel2) center/cover no-repeat; box-shadow: 0 2px 8px rgba(0,0,0,.15); }
|
||||
.stack-cards .sc:nth-child(1) { left: 0; top: 8px; transform: rotate(-5deg); }
|
||||
.stack-cards .sc:nth-child(2) { left: 17%; top: 4px; transform: rotate(1deg); z-index: 1; }
|
||||
.stack-cards .sc:nth-child(3) { left: 34%; top: 0; transform: rotate(6deg); z-index: 2; }
|
||||
.stack-meta { display: flex; align-items: center; justify-content: space-between; }
|
||||
.stack-name { font-weight: 600; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
|
||||
.stack-count { font-size: 12px; color: var(--dim); background: var(--panel2); border-radius: 999px; padding: 1px 8px; }
|
||||
.stack.open .stack-open { animation: folder-open .22s ease; overflow: hidden; margin-top: 10px; display: flex; flex-direction: column; gap: 8px; }
|
||||
@keyframes folder-open { from { opacity: 0; transform: translateY(-6px) scale(.98); } to { opacity: 1; transform: none; } }
|
||||
.stack-rename { border: 1px solid var(--line); background: var(--panel2); color: var(--text); border-radius: 8px; padding: 6px 9px; font: inherit; font-weight: 600; }
|
||||
.stack-rename:focus { outline: none; border-color: var(--brand); }
|
||||
.stack-member { position: relative; }
|
||||
.unstack { position: absolute; right: -6px; top: -6px; z-index: 5; width: 22px; height: 22px; border-radius: 999px; border: 1px solid var(--line); background: var(--panel); color: var(--text); cursor: pointer; line-height: 1; }
|
||||
.unstack:hover { background: var(--bad); color: #fff; border-color: transparent; }
|
||||
|
||||
/* ── Overlays / sheets ───────────────────────────────────────────────────── */
|
||||
.ovl { position: fixed; inset: 0; z-index: 60; background: rgba(0,0,0,.4); display: grid; place-items: center; padding: 16px; }
|
||||
.sheet { width: min(460px, 100%); max-height: 80vh; overflow-y: auto; background: var(--panel); border: 1px solid var(--line); border-radius: 16px; padding: 16px; display: flex; flex-direction: column; gap: 10px; box-shadow: 0 20px 60px rgba(0,0,0,.35); }
|
||||
.sheet-h { font-weight: 600; font-size: 16px; }
|
||||
.muted { color: var(--dim); font-size: 13px; margin: 0; }
|
||||
.sheet-in { border: 1px solid var(--line); background: var(--panel2); color: var(--text); border-radius: 10px; padding: 10px 12px; font: inherit; }
|
||||
.sheet-in:focus { outline: none; border-color: var(--brand); }
|
||||
.sheet-item { text-align: left; }
|
||||
.sheet-item:hover { border-color: var(--dim); }
|
||||
.sheet-row { display: flex; gap: 8px; justify-content: flex-end; margin-top: 4px; }
|
||||
.vers-strip { display: flex; gap: 10px; overflow-x: auto; padding: 4px 0 8px; }
|
||||
.vstep { flex: 0 0 auto; width: 92px; text-align: center; text-decoration: none; color: var(--dim); }
|
||||
.vstep img, .vstep .rc-ph { width: 92px; height: 92px; object-fit: cover; border-radius: 10px; border: 1px solid var(--line); }
|
||||
.vstep.cur img, .vstep.cur .rc-ph { border-color: var(--brand); box-shadow: 0 0 0 2px color-mix(in srgb, var(--brand) 40%, transparent); }
|
||||
.vlabel { display: block; font-size: 12px; margin-top: 4px; }
|
||||
.vstep.cur .vlabel { color: var(--text); font-weight: 600; }
|
||||
|
||||
.rail-toast { position: fixed; left: 50%; bottom: 20px; transform: translateX(-50%); z-index: 80; background: var(--bad); color: #fff; padding: 9px 16px; border-radius: 10px; font-size: 13px; box-shadow: 0 8px 24px rgba(0,0,0,.3); max-width: 90vw; }
|
||||
|
||||
/* ── Mobile / tablet: ONE vertical scroll — chat transcript, then the rail as the
|
||||
bottom section — with the composer FIXED to the viewport bottom + safe-area inset
|
||||
(the #1 thumb-reach fix). Dividers are a desktop affordance (hidden); collapse
|
||||
chevrons + filter chips + persistence all still apply. */
|
||||
@media (max-width: 820px) {
|
||||
.body { grid-template-columns: 1fr; }
|
||||
.rail { display: none; }
|
||||
.chat { border-right: 0; }
|
||||
.body { flex-direction: column; overflow-y: auto; -webkit-overflow-scrolling: touch;
|
||||
padding-bottom: calc(84px + env(safe-area-inset-bottom)); }
|
||||
.chat { flex: 0 0 auto; min-height: 52vh; }
|
||||
.scroll { flex: 0 0 auto; overflow: visible; }
|
||||
.composer { position: fixed; left: 0; right: 0; bottom: 0; z-index: 40;
|
||||
border-top: 1px solid var(--line); box-shadow: 0 -6px 18px rgba(0,0,0,.12);
|
||||
padding-bottom: calc(12px + env(safe-area-inset-bottom)); }
|
||||
.vsplit, .hsplit { display: none; }
|
||||
.rail { width: auto !important; flex: 0 0 auto; border-left: 0; border-top: 1px solid var(--line); overflow: visible; }
|
||||
.rail-main, .sections, .sec, .sec-body, .stacks-view { overflow: visible; min-height: 0; }
|
||||
.sec, .sec.col { flex: 0 0 auto !important; }
|
||||
.chips { position: sticky; top: 0; z-index: 5; background: var(--bg); }
|
||||
.rc-thumb { flex-basis: 44%; max-width: 44%; }
|
||||
.stacks-view { grid-template-columns: repeat(2, 1fr); }
|
||||
}
|
||||
@media (max-width: 420px) {
|
||||
.stacks-view { grid-template-columns: 1fr 1fr; }
|
||||
.sheet { border-radius: 14px; }
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user