studio: native IAM auth + multi-tenant isolation + GPU federation

Auth (middleware/iam_auth_middleware.py): local JWKS JWT validation against
Hanzo IAM (signature + exp + iss + aud=hanzo-studio; org from the `owner`
claim). Standard OIDC Authorization Code flow with PKCE and a /callback
session cookie for the standalone studio.hanzo.ai app; anonymous localhost
bypass unchanged. verify_jwt() extracted as a pure, unit-tested function.

Tenancy (folder_paths.py, app/user_manager.py): make get_public_user_directory
org-rooted so multi-tenant userdata resolves under user/orgs/{org}/... (it
previously always returned None). Scope execution outputs to output/orgs/{org}/
via set_execution_org(), bound by the prompt worker from the queue item org_id.
Also block a token subject from escaping into a System User namespace.

Federation (middleware/{worker_client,prompt_router}.py, server.py): move the
worker registry + execute surface to /v1/workers/register, /v1/workers,
/v1/worker/execute; add X-Worker-Token coordinator trust. Outbound pull channel
for NAT'd local boxes (GB10) specced in docs/federation.md.

Deploy: hanzo.yml + .github/workflows/cicd.yml (hanzoai/ci reusable) build
ghcr.io/hanzoai/studio and roll the studio operator Service CR at
studio.hanzo.ai. PyJWT[crypto] pinned.

Tests: tests-unit/middleware_test/iam_auth_test.py (RS256 sign/verify, state
CSRF, path exemptions) + folder_paths_test/org_scoping_test.py; 140 passing
across the auth/tenancy/user-manager/prompt-server suites.

.github/workflows/cicd.yml

@@ -0,0 +1,16 @@
+# Canonical CI/CD — imports the shared hanzoai/ci reusable workflow, which
+# reads hanzo.yml (images + test + deploy + kms). No per-repo build logic,
+# no GitHub-hosted runners (defaults to the Hanzo cloud arc pool).
+name: CI/CD
+
+on:
+  push:
+    branches: [main]
+    tags: ['v*']
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  cicd:
+    uses: hanzoai/ci/.github/workflows/build.yml@v1
+    secrets: inherit

LLM.md

@@ -35,4 +35,38 @@ studio/
 ## Key Files
 - `README.md` -- Project documentation
 - `pyproject.toml` -- Python project config
-- `Dockerfile` -- Container build
+- `Dockerfile` -- Container build (applies `branding/` then runs `main.py`)
+- `hanzo.yml` + `.github/workflows/cicd.yml` -- canonical CI/CD (hanzoai/ci); builds `ghcr.io/hanzoai/studio`, rolls the `studio` operator Service CR at `studio.hanzo.ai`.
+
+## Multi-tenant + IAM (cloud: studio.hanzo.ai)
+Standalone app; `console.hanzo.ai` link-outs land here already authenticated via
+shared Hanzo IAM (standard OIDC code flow). No embedding.
+
+- **Auth** (`middleware/iam_auth_middleware.py`): local **JWKS** JWT validation
+  against `${STUDIO_IAM_URL}/v1/iam/.well-known/jwks` — signature + `exp` +
+  `iss` + `aud` (client_id `hanzo-studio`). Org from the `owner` claim
+  (`organization` fallback) → `request["iam_user"]["org_id"]`. Anonymous mode:
+  off by default locally (loopback bypass); enable in cloud with
+  `STUDIO_ENABLE_IAM_AUTH=true` (the one auth switch — this is the
+  "STUDIO_AUTH=off default" behavior). Browser flow: unauth HTML → IAM
+  `authorize` (code+PKCE) → `/callback` exchanges the code, sets an httpOnly
+  session cookie (no frontend changes). `verify_jwt(...)` is a pure, unit-tested
+  function.
+- **Tenancy** (`folder_paths.py`, `app/user_manager.py`): with
+  `STUDIO_MULTI_TENANT=true`, userdata is `user/orgs/{org}/user/{user}/…` and
+  outputs are `output/orgs/{org}/…`. `get_public_user_directory(user, org)` is
+  org-rooted (fixed the bug where org userdata always resolved to `None`).
+  Execution outputs are scoped via `folder_paths.set_execution_org()`, bound by
+  the prompt worker from the queue item's `org_id`.
+- **GPU federation** (`middleware/{worker_client,prompt_router,compute_config,
+  visor_client}.py`): `/v1/workers/register` (heartbeat), `/v1/workers` (list),
+  `/v1/worker/execute` (push). Worker↔coordinator trust via
+  `X-Worker-Token` (`STUDIO_WORKER_TOKEN`, KMS-sourced). NAT'd local boxes (GB10)
+  join outbound — pull channel is specced in `docs/federation.md`.
+
+## Tests
+`tests-unit/{folder_paths_test,middleware_test,app_test,prompt_server_test}` —
+run `uv run pytest tests-unit/folder_paths_test tests-unit/middleware_test
+tests-unit/app_test tests-unit/prompt_server_test -q`. The auth/tenancy work is
+covered by `middleware_test/iam_auth_test.py` and
+`folder_paths_test/org_scoping_test.py`.

app/user_manager.py

@@ -62,6 +62,9 @@ class UserManager():
         iam_user = request.get("iam_user")
         if iam_user and iam_user.get("sub") and iam_user["sub"] != "local":
             user = iam_user["sub"]
+            # A token subject must never escape into a System User namespace.
+            if user.startswith(folder_paths.SYSTEM_USER_PREFIX):
+                raise KeyError("Unknown user: " + user)
             # Auto-register IAM users
             if user not in self.users:
                 display = iam_user.get("name") or iam_user.get("email") or user
@@ -95,7 +98,7 @@ class UserManager():
             raise KeyError("Unknown filepath type:" + type)
 
         user = self.get_request_user_id(request)
-        user_root = folder_paths.get_public_user_directory(user)
+        user_root = folder_paths.get_public_user_directory(user, org_id)
         if user_root is None:
             return None
         path = user_root

docs/federation.md

@@ -0,0 +1,168 @@
+# Hanzo Studio — GPU Federation
+
+How `studio.hanzo.ai` schedules generation jobs across a heterogeneous GPU pool:
+in-cluster cloud pods **and** local boxes (e.g. a GB10) that join over an
+**outbound-only** connection. No inbound tunnel to the local box is ever
+required.
+
+This document is the contract. Sections are marked **[implemented]** (code in
+this repo today) or **[specced]** (agreed design, not yet wired). Nothing here is
+a stub in code — specced pieces live in this doc until they are built.
+
+## 1. Roles
+
+- **Coordinator** — the `studio.hanzo.ai` deployment. Full UI, IAM auth,
+  multi-tenant storage. Owns the per-org worker registry and the job schedule.
+  Runs a small CPU footprint (it routes; it does not need a GPU).
+- **Worker** — a headless Studio started with `--worker-mode`. Reports its
+  device (GPU model, VRAM), registers with the coordinator, and executes prompts
+  it is given. A worker belongs to exactly one org (`STUDIO_ORG_ID`).
+
+A worker is the *same binary* as the coordinator; `--worker-mode` only drops the
+UI/websocket surface and turns on registration + the execute endpoint. This is
+the ComfyUI-Distributed lineage (workers are full ComfyUI instances the control
+plane drives) narrowed to **job-level** dispatch and hardened for multi-tenant +
+NAT traversal.
+
+## 2. Two dispatch models
+
+Reachability, not preference, decides which model a worker uses.
+
+### 2a. Push — for in-cluster workers **[implemented]**
+
+The coordinator can open a connection to the worker (both are pods in
+`hanzo-k8s`, or the worker has a routable `WORKER_EXTERNAL_URL`).
+
+```
+client → POST /prompt (coordinator)
+  prompt_router.route_prompt(org, body)
+    → get_available_gpu_worker(org)         # compute_config registry
+    → POST {worker.url}/v1/worker/execute   # forward_to_worker
+  worker queues + executes locally, returns {prompt_id, ...}
+```
+
+Files: `middleware/prompt_router.py`, `middleware/worker_client.py`
+(`add_worker_routes` → `/v1/worker/execute`), `server.py` (`/prompt` calls
+`route_prompt`).
+
+Limitation: requires the coordinator to reach the worker. A GB10 behind NAT is
+**not** reachable, so push cannot serve it.
+
+### 2b. Pull — for NAT'd local boxes (the GB10 pattern) **[specced]**
+
+The worker dials **out** and keeps the connection; the coordinator pushes jobs
+down that already-open channel. No inbound port on the box.
+
+Primary transport is a WebSocket the worker opens to the coordinator:
+
+```
+worker → WSS {coordinator}/v1/workers/connect      (Authorization: worker token)
+  ← {type:"job", job_id, prompt, extra_data, org_id}
+  → {type:"progress", job_id, value, max}
+  → {type:"result", job_id, status, outputs:[{name, s3_url}...]}
+```
+
+HTTP long-poll fallback for environments without WSS:
+
+```
+worker → POST /v1/jobs/claim   {worker_id, org_id}     # long-poll, ≤30s
+  ← 204 (no work) | 200 {job_id, prompt, extra_data}
+worker → POST /v1/jobs/{job_id}/result {status, outputs}
+```
+
+The coordinator keeps a per-org, per-worker job queue (a natural extension of
+`compute_config`), enqueues on `/prompt` when the target worker is pull-mode, and
+hands the job to whichever channel the worker holds open. `route_prompt` gains a
+third branch (`{"action":"queued"}` → `202`) alongside the existing
+forward/provisioning/unavailable branches.
+
+Not implemented because it does not "drop out" of ComfyUI's `--listen`
+architecture — it needs a real coordinator-side queue and a WS control channel.
+Building it does not change the endpoints below.
+
+## 3. Endpoints
+
+All under `/v1/` (also served under the `/api/v1/` alias ComfyUI adds for the
+frontend proxy). Worker↔coordinator endpoints are exempt from user IAM and
+carry the coordinator shared secret instead (`X-Worker-Token`, §5).
+
+| Method | Path | Dir | Status | Purpose |
+|---|---|---|---|---|
+| POST | `/v1/workers/register` | worker→coord | **impl** | Register / heartbeat. Body: `{worker_id, url, org_id, device, gpu_model, vram_gb, status}`. |
+| GET  | `/v1/workers` | client→coord | **impl** | List the org's workers + liveness. |
+| POST | `/v1/worker/execute` | coord→worker | **impl** | Push a prompt to a reachable worker. Same body as `/prompt`. |
+| WSS  | `/v1/workers/connect` | worker→coord | **spec** | Persistent pull channel for NAT'd workers. |
+| POST | `/v1/jobs/claim` | worker→coord | **spec** | HTTP long-poll fallback to claim a job. |
+| POST | `/v1/jobs/{id}/result` | worker→coord | **spec** | Report result + artifact URLs. |
+| GET  | `/v1/compute/config` | client→coord | **impl** | Read the org's compute profile. |
+| PUT  | `/v1/compute/config` | client→coord | **impl** | Update profile (GPU tier, auto-provision). |
+| POST | `/v1/compute/provision` | client→coord | **impl** | Autoscale a cloud GPU worker via Visor. |
+
+(Compute-profile + Visor autoscale routes are still under `/compute/*` in
+`server.py`; the `/v1/compute/*` labels above are the canonical names to fold to
+when those handlers are next touched. Worker-registry + execute already moved.)
+
+## 4. Registration & scheduling **[implemented]**
+
+- Worker sends a heartbeat every 30s (`WorkerClient._heartbeat_loop`).
+- Coordinator upserts it into the org's `compute.json`
+  (`compute_config.register_worker`), pruning workers silent >300s.
+  Liveness for scheduling is 90s (`WorkerInfo.is_alive`).
+- `get_available_gpu_worker(org)` returns the first ready, alive CUDA worker.
+- `prompt_router.route_prompt` chooses local vs worker from the prompt's
+  `device_preference` (`auto`|`cpu`|`gpu`) and the org's `gpu_enabled`.
+- If no worker and the org has `auto_provision`, Visor launches a cloud GPU VM
+  that boots a worker (`middleware/visor_client.py`).
+
+## 5. Security
+
+- **User auth** — all user-facing routes require an IAM JWT (JWKS-verified);
+  org comes from the `owner` claim. See `middleware/iam_auth_middleware.py`.
+- **Worker trust** — worker↔coordinator endpoints are IAM-exempt (they carry no
+  user) and instead require `X-Worker-Token == STUDIO_WORKER_TOKEN`
+  (`worker_client.verify_worker_token`). The token is KMS-sourced in cloud;
+  empty in a single-trust-domain local dev, where the check is skipped.
+  **[specced]** forward direction: replace the shared secret with an IAM
+  service-account token (`hanzo-studio` client-credentials grant) so each worker
+  is individually attributable and revocable.
+- **Org isolation** — a job carries its `org_id`; the worker binds it as the
+  execution org so outputs land in that org's namespace (§6). A worker only ever
+  claims jobs for its own `STUDIO_ORG_ID`.
+- **No inbound** — pull-mode boxes expose nothing; the coordinator is reached
+  over TLS via `hanzoai/ingress`.
+
+## 6. Output persistence **[specced]**
+
+Push/pull both execute on the worker, which writes to *its own* output dir
+(org-scoped via `folder_paths.set_execution_org`). For the coordinator UI to
+show results, workers must ship artifacts to shared storage:
+
+- On completion, upload outputs to `s3.lux.cloud` (hanzo bucket),
+  key `studio/{org_id}/{prompt_id}/{filename}`.
+- Report the object URLs in the `result` message; the coordinator records them
+  in history so `/view` (or a signed redirect) serves them.
+
+Until this lands, push-mode results remain on the worker and are addressable
+only if the worker is reachable — acceptable for the in-cluster pool, blocking
+for the NAT pull pool. This is the top build item after the pull channel.
+
+## 7. Joining a local box (GB10)
+
+```bash
+python main.py \
+  --worker-mode \
+  --coordinator-url https://studio.hanzo.ai \
+  --worker-id gb10-1
+# env: STUDIO_ORG_ID=<org>  STUDIO_WORKER_TOKEN=<from KMS>
+```
+
+Today (push) this works only if the coordinator can reach the box (set
+`WORKER_EXTERNAL_URL`). Once §2b lands, the box needs **no** inbound access: it
+registers, opens the WSS channel, and pulls jobs for its org.
+
+## 8. Build order
+
+1. Pull WSS channel `/v1/workers/connect` + coordinator per-worker queue (§2b).
+2. Artifact upload to `s3.lux.cloud` + result recording (§6).
+3. IAM service-account worker identity replacing the shared secret (§5).
+4. Fold `/compute/*` profile routes to `/v1/compute/*` (§3).

folder_paths.py

@@ -58,9 +58,30 @@ input_directory = os.path.join(base_path, "input")
 user_directory = os.path.join(base_path, "user")
 
 # --- Multi-tenant org-scoped storage ---
+import contextvars
+
 _org_id: str | None = args.org_id if hasattr(args, "org_id") else None
 _multi_tenant: bool = getattr(args, "multi_tenant", False)
 
+# Per-execution org, set by the prompt worker thread before a graph runs so that
+# nodes writing outputs (SaveImage -> get_output_directory) land in the org's
+# namespace. A ContextVar isolates the worker thread from concurrent async web
+# handlers, which resolve org explicitly per-request instead. Default None means
+# "no active execution org" -> fall back to the process org (_org_id) or base.
+_execution_org: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "studio_execution_org", default=None
+)
+
+
+def set_execution_org(org_id: str | None) -> None:
+    """Bind the org for the current execution context (worker thread)."""
+    _execution_org.set(org_id)
+
+
+def get_execution_org() -> str | None:
+    """Return the org bound to the current execution context, if any."""
+    return _execution_org.get()
+
 
 def set_org_id(org_id: str | None) -> None:
     """Set the active organization for multi-tenant path scoping."""
@@ -87,33 +108,33 @@ def get_org_base_path(org_id: str | None = None) -> str:
 
 
 def get_org_output_directory(org_id: str | None = None) -> str:
-    """Get org-scoped output directory."""
+    """Get org-scoped output directory. Falls back to the base output dir."""
     ob = get_org_base_path(org_id)
     if ob != base_path:
         d = os.path.join(ob, "output")
         os.makedirs(d, exist_ok=True)
         return d
-    return get_output_directory()
+    return output_directory
 
 
 def get_org_input_directory(org_id: str | None = None) -> str:
-    """Get org-scoped input directory."""
+    """Get org-scoped input directory. Falls back to the base input dir."""
     ob = get_org_base_path(org_id)
     if ob != base_path:
         d = os.path.join(ob, "input")
         os.makedirs(d, exist_ok=True)
         return d
-    return get_input_directory()
+    return input_directory
 
 
 def get_org_temp_directory(org_id: str | None = None) -> str:
-    """Get org-scoped temp directory."""
+    """Get org-scoped temp directory. Falls back to the base temp dir."""
     ob = get_org_base_path(org_id)
     if ob != base_path:
         d = os.path.join(ob, "temp")
         os.makedirs(d, exist_ok=True)
         return d
-    return get_temp_directory()
+    return temp_directory
 
 
 def get_org_user_directory(org_id: str | None = None) -> str:
@@ -196,16 +217,15 @@ def set_input_directory(input_dir: str) -> None:
     input_directory = input_dir
 
 def get_output_directory() -> str:
-    global output_directory
-    return output_directory
+    # Resolve against the current execution org (worker thread) when set; in
+    # single-tenant mode get_org_output_directory returns the base dir unchanged.
+    return get_org_output_directory(_execution_org.get())
 
 def get_temp_directory() -> str:
-    global temp_directory
-    return temp_directory
+    return get_org_temp_directory(_execution_org.get())
 
 def get_input_directory() -> str:
-    global input_directory
-    return input_directory
+    return get_org_input_directory(_execution_org.get())
 
 def get_user_directory() -> str:
     return user_directory
@@ -252,7 +272,7 @@ def get_system_user_directory(name: str = "system") -> str:
     return os.path.join(get_user_directory(), f"{SYSTEM_USER_PREFIX}{name}")
 
 
-def get_public_user_directory(user_id: str) -> str | None:
+def get_public_user_directory(user_id: str, org_id: str | None = None) -> str | None:
     """
     Get the path to a Public User directory for HTTP endpoint access.
 
@@ -260,8 +280,15 @@ def get_public_user_directory(user_id: str) -> str | None:
     System User (prefixed with '__'). All HTTP endpoints should use this
     function instead of directly constructing user paths.
 
+    In multi-tenant mode the user directory is rooted under the org namespace
+    (user/orgs/{org_id}/user/{user_id}); otherwise it is the flat
+    user/{user_id}. Passing org_id keeps the returned path a child of
+    get_org_user_directory(org_id), which callers rely on for containment
+    checks.
+
     Args:
         user_id: User identifier from HTTP request.
+        org_id: Organization the user belongs to (from the IAM token), or None.
 
     Returns:
         Absolute path to the user directory, or None if user_id is invalid
@@ -277,7 +304,7 @@ def get_public_user_directory(user_id: str) -> str | None:
         return None
     if user_id.startswith(SYSTEM_USER_PREFIX):
         return None
-    return os.path.join(get_user_directory(), user_id)
+    return os.path.join(get_org_user_directory(org_id), user_id)
 
 
 #NOTE: used in http server so don't put folders that should not be accessed remotely

hanzo.yml

@@ -0,0 +1,41 @@
+# Canonical CI/CD config for hanzoai/studio — read by both the hanzoai/ci
+# reusable workflow (.github/workflows/cicd.yml) and platform.hanzo.ai.
+# One image, unit gates, and a main-branch rollout onto the existing
+# `studio` operator Service CR (universe: infra/k8s/operator/crs/studio.yaml).
+
+images:
+  - name: studio
+    context: .
+    dockerfile: Dockerfile
+    repo: ghcr.io/hanzoai/studio
+
+# CI gates. Full requirements (CPU torch) are installed so every import
+# resolves, then the dependency-light unit suites run (auth, tenancy, user
+# manager, prompt server). Inference/execution/GPU suites run out-of-band.
+test:
+  - name: unit
+    run: |
+      uv venv --python 3.11
+      uv pip install -r requirements.txt pytest pytest-aiohttp pytest-asyncio
+      uv run pytest -q \
+        tests-unit/folder_paths_test \
+        tests-unit/middleware_test \
+        tests-unit/app_test \
+        tests-unit/prompt_server_test
+
+# Deploy: roll the freshly-built image onto the studio Service CR.
+deploy:
+  cluster: do-sfo3-hanzo-k8s
+  namespace: hanzo
+  # Quoted to dodge YAML 1.1 boolean coercion of a bare `on:` key.
+  "on": [main]
+  services:
+    - name: studio
+      image: studio
+
+# KMS (kms.hanzo.ai, Universal Auth) supplies the GHCR push token and the
+# target kubeconfig at run time. Only KMS_CLIENT_ID/KMS_CLIENT_SECRET live in
+# GitHub; everything else is fetched here.
+kms:
+  path: /deploy
+  environment: prod

main.py

@@ -267,7 +267,12 @@ def prompt_worker(q, server_instance):
             for k in sensitive:
                 extra_data[k] = sensitive[k]
 
-            e.execute(item[2], prompt_id, extra_data, item[4])
+            # Tenancy: scope this execution's output/temp/input dirs to the org.
+            folder_paths.set_execution_org(extra_data.get("org_id"))
+            try:
+                e.execute(item[2], prompt_id, extra_data, item[4])
+            finally:
+                folder_paths.set_execution_org(None)
             need_gc = True
 
             remove_sensitive = lambda prompt: prompt[:5] + prompt[6:]

middleware/iam_auth_middleware.py

@@ -1,20 +1,39 @@
 """
 IAM authentication middleware for Hanzo Studio.
 
-Validates Bearer tokens against the IAM OIDC userinfo endpoint
-(/v1/iam/oauth/userinfo) per HIP-0111.
-Localhost connections bypass auth. Static assets and health checks are exempt.
+Validates Bearer/cookie JWTs locally against the Hanzo IAM JWKS
+(``{iam_url}/v1/iam/.well-known/jwks``) — no per-request network call once the
+signing keys are cached. Signature (RS256 by default), ``exp``, ``iss`` and
+``aud`` are all verified. The organization is taken from the ``owner`` claim
+(``organization`` fallback) and exposed as ``request["iam_user"]`` for the
+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.
+
+Localhost connections bypass auth (single-user dev). Static assets, health
+checks and the worker coordinator endpoints are exempt.
 """
+import base64
+import hashlib
+import hmac
+import ipaddress
+import json
 import logging
 import os
-import ipaddress
-import hashlib
+import secrets
 import time as _time
+from urllib.parse import urlencode
 
 import aiohttp
 from aiohttp import web
 
-# Paths that never require auth (static assets, health, ws negotiation)
+logger = logging.getLogger(__name__)
+
+# Paths that never require auth (static assets, health, ws negotiation, oauth).
 PUBLIC_PATH_PREFIXES = (
     "/favicon",
     "/assets/",
@@ -25,6 +44,7 @@ PUBLIC_PATH_PREFIXES = (
 
 PUBLIC_PATHS_EXACT = {
     "/",
+    "/callback",
     "/health",
     "/ready",
     "/metrics",
@@ -32,162 +52,323 @@ PUBLIC_PATHS_EXACT = {
     "/api/ready",
     "/api/features",
     "/api/metrics",
-    # Worker registration uses its own auth (coordinator trust)
-    "/compute/workers/register",
-    "/api/compute/workers/register",
-    "/worker/execute",
-    "/api/worker/execute",
+    # Worker registration / job exchange use their own coordinator trust
+    # (STUDIO_WORKER_TOKEN); see worker_client and the federation design doc.
+    "/v1/workers/register",
+    "/api/v1/workers/register",
+    "/v1/jobs/claim",
+    "/api/v1/jobs/claim",
+    "/v1/worker/execute",
+    "/api/v1/worker/execute",
 }
 
-# Token validation cache: hash(token) -> (user_info, expires_at)
+_OIDC_SCOPES = "openid profile email"
+_STATE_COOKIE = "studio_oauth"
+_SESSION_COOKIE = "hanzo_token"
+_STATE_TTL = 600  # seconds a pending login may take
+
+# Decoded-claims cache: hash(token) -> (user_info, expires_at_monotonic)
 _token_cache: dict[str, tuple[dict, float]] = {}
-_CACHE_TTL = 60  # seconds
 
 
 def _is_loopback_request(request: web.Request) -> bool:
-    """Check if the request originates from a loopback address."""
-    peername = request.transport.get_extra_info("peername")
+    peername = request.transport.get_extra_info("peername") if request.transport else None
     if peername is None:
         return False
-    host = peername[0]
     try:
-        return ipaddress.ip_address(host).is_loopback
+        return ipaddress.ip_address(peername[0]).is_loopback
     except ValueError:
         return False
 
 
 def _is_public_path(path: str) -> bool:
-    """Check if the path is publicly accessible without auth."""
     if path in PUBLIC_PATHS_EXACT:
         return True
-    for prefix in PUBLIC_PATH_PREFIXES:
-        if path.startswith(prefix):
-            return True
-    # Static file extensions served from web root
-    if any(path.endswith(ext) for ext in (".js", ".css", ".svg", ".png", ".ico", ".woff", ".woff2", ".map", ".webm")):
+    if any(path.startswith(p) for p in PUBLIC_PATH_PREFIXES):
         return True
-    return False
+    return any(path.endswith(ext) for ext in (
+        ".js", ".css", ".svg", ".png", ".ico", ".woff", ".woff2", ".map", ".webm",
+    ))
 
 
 def _cache_key(token: str) -> str:
     return hashlib.sha256(token.encode()).hexdigest()[:32]
 
 
+def _b64url(raw: bytes) -> str:
+    return base64.urlsafe_b64encode(raw).rstrip(b"=").decode()
+
+
+def _external_base(request: web.Request) -> str:
+    """Public origin of this Studio, honoring the ingress forwarded headers."""
+    proto = request.headers.get("X-Forwarded-Proto", request.scheme)
+    host = request.headers.get("X-Forwarded-Host", request.host)
+    return f"{proto}://{host}"
+
+
+class _JwksCache:
+    """Async-fetched, kid-indexed JWKS with rate-limited refresh on cache miss."""
+
+    def __init__(self, jwks_uri: str, ttl: float = 600.0, min_refresh: float = 30.0):
+        self._uri = jwks_uri
+        self._ttl = ttl
+        self._min_refresh = min_refresh
+        self._keys: dict[str, object] = {}
+        self._fetched_at = 0.0
+        self._session: aiohttp.ClientSession | None = None
+
+    async def _session_get(self) -> aiohttp.ClientSession:
+        if self._session is None or self._session.closed:
+            self._session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=10))
+        return self._session
+
+    async def _refresh(self) -> None:
+        from jwt import PyJWKSet
+        session = await self._session_get()
+        async with session.get(self._uri) as resp:
+            if resp.status != 200:
+                raise RuntimeError(f"JWKS fetch {resp.status} from {self._uri}")
+            data = await resp.json()
+        jwk_set = PyJWKSet.from_dict(data)
+        self._keys = {k.key_id: k.key for k in jwk_set.keys if k.key_id}
+        self._fetched_at = _time.monotonic()
+
+    async def get_key(self, kid: str):
+        now = _time.monotonic()
+        stale = (now - self._fetched_at) > self._ttl
+        if (kid not in self._keys or stale) and (now - self._fetched_at) > self._min_refresh:
+            await self._refresh()
+        return self._keys.get(kid)
+
+
+def _sign_state(payload: dict, secret: str) -> str:
+    body = _b64url(json.dumps(payload, separators=(",", ":")).encode())
+    mac = _b64url(hmac.new(secret.encode(), body.encode(), hashlib.sha256).digest())
+    return f"{body}.{mac}"
+
+
+def _verify_state(blob: str, secret: str) -> dict | None:
+    try:
+        body, mac = blob.split(".", 1)
+    except ValueError:
+        return None
+    expected = _b64url(hmac.new(secret.encode(), body.encode(), hashlib.sha256).digest())
+    if not hmac.compare_digest(mac, expected):
+        return None
+    try:
+        pad = "=" * (-len(body) % 4)
+        payload = json.loads(base64.urlsafe_b64decode(body + pad))
+    except Exception:
+        return None
+    if payload.get("exp", 0) < _time.time():
+        return None
+    return payload
+
+
+# IAM signs RS256 by default; ES/RS variants are selectable per application.
+_ALGS = ["RS256", "RS512", "ES256", "ES384", "ES512"]
+
+
+def _claims_to_user(claims: dict) -> dict:
+    return {
+        "sub": claims.get("sub", ""),
+        "name": claims.get("name") or claims.get("preferred_username", ""),
+        "email": claims.get("email", ""),
+        "org_id": claims.get("owner") or claims.get("organization") or "default",
+        "avatar": claims.get("picture", ""),
+    }
+
+
+def verify_jwt(token: str, key, issuer: str, audience: str | None) -> dict | None:
+    """
+    Pure JWT verification: signature (via the resolved JWKS key), exp, iss and
+    aud. Returns normalized user info or None if the token is invalid. No I/O.
+    """
+    import jwt
+    try:
+        claims = jwt.decode(
+            token,
+            key,
+            algorithms=_ALGS,
+            audience=audience,
+            issuer=issuer,
+            options={"require": ["exp", "iss"], "verify_aud": bool(audience)},
+        )
+    except jwt.InvalidTokenError:
+        return None
+    return _claims_to_user(claims)
+
+
 def create_iam_auth_middleware(iam_url: str, localhost_bypass: bool = True):
     """
-    Factory that creates an aiohttp middleware for IAM token validation.
+    Build the aiohttp middleware that validates IAM JWTs via JWKS.
 
-    Args:
-        iam_url: Base URL of the IAM server (e.g. https://hanzo.id)
-        localhost_bypass: If True, skip auth for loopback connections
+    Config is read from the environment (KMS-sourced in cloud):
+      STUDIO_IAM_CLIENT_ID     OAuth client_id / expected audience (default hanzo-studio)
+      STUDIO_IAM_CLIENT_SECRET confidential-client secret for the code exchange
+      STUDIO_IAM_ISSUER        expected iss (default = iam_url, e.g. https://hanzo.id)
     """
     iam_url = iam_url.rstrip("/")
-    account_endpoint = f"{iam_url}/v1/iam/oauth/userinfo"
+    client_id = os.environ.get("STUDIO_IAM_CLIENT_ID", "hanzo-studio")
+    issuer = os.environ.get("STUDIO_IAM_ISSUER", iam_url)
+    jwks_uri = f"{iam_url}/v1/iam/.well-known/jwks"
+    jwks = _JwksCache(jwks_uri)
 
-    # Shared session — created lazily, reused across requests
-    _shared_session: aiohttp.ClientSession | None = None
+    async def _validate(token: str) -> dict | None:
+        """Return normalized user info for a valid token, else None."""
+        ck = _cache_key(token)
+        cached = _token_cache.get(ck)
+        if cached and cached[1] > _time.monotonic():
+            return cached[0]
 
-    async def _get_session() -> aiohttp.ClientSession:
-        nonlocal _shared_session
-        if _shared_session is None or _shared_session.closed:
-            _shared_session = aiohttp.ClientSession(
-                timeout=aiohttp.ClientTimeout(total=10),
-            )
-        return _shared_session
+        import jwt
+        try:
+            header = jwt.get_unverified_header(token)
+        except jwt.InvalidTokenError:
+            return None
+        kid = header.get("kid")
+        if not kid:
+            return None
+        try:
+            key = await jwks.get_key(kid)
+        except Exception as e:
+            logger.error("JWKS fetch failed: %s", e)
+            raise
+        if key is None:
+            return None
+
+        user = verify_jwt(token, key, issuer, client_id)
+        if user is None:
+            _token_cache.pop(ck, None)
+            return None
+
+        # Cache the validated result, capped at 60s so revocation propagates.
+        _token_cache[ck] = (user, _time.monotonic() + 60.0)
+        if len(_token_cache) > 2000:
+            now = _time.monotonic()
+            for k in [k for k, v in _token_cache.items() if v[1] < now]:
+                _token_cache.pop(k, None)
+        return user
+
+    def _authorize_redirect(request: web.Request) -> web.Response:
+        base = _external_base(request)
+        verifier = _b64url(secrets.token_bytes(32))
+        challenge = _b64url(hashlib.sha256(verifier.encode()).digest())
+        nonce = secrets.token_urlsafe(16)
+        state_payload = {
+            "n": nonce,
+            "v": verifier,
+            "t": str(request.rel_url),
+            "exp": int(_time.time()) + _STATE_TTL,
+        }
+        state = _sign_state(state_payload, client_id)  # HMAC key; secret optional
+        params = {
+            "response_type": "code",
+            "client_id": client_id,
+            "redirect_uri": f"{base}/callback",
+            "scope": _OIDC_SCOPES,
+            "state": state,
+            "code_challenge": challenge,
+            "code_challenge_method": "S256",
+        }
+        resp = web.HTTPFound(f"{iam_url}/v1/iam/oauth/authorize?{urlencode(params)}")
+        resp.set_cookie(
+            _STATE_COOKIE, state, max_age=_STATE_TTL, httponly=True,
+            secure=base.startswith("https"), samesite="Lax", path="/",
+        )
+        return resp
+
+    async def handle_callback(request: web.Request) -> web.Response:
+        """Exchange the authorization code for a token and set the session cookie."""
+        code = request.query.get("code")
+        state = request.query.get("state")
+        cookie_state = request.cookies.get(_STATE_COOKIE)
+        if not code or not state or state != cookie_state:
+            return web.Response(status=400, text="Invalid OAuth state")
+        payload = _verify_state(state, client_id)
+        if payload is None:
+            return web.Response(status=400, text="Expired or forged OAuth state")
+
+        base = _external_base(request)
+        client_secret = os.environ.get("STUDIO_IAM_CLIENT_SECRET", "")
+        form = {
+            "grant_type": "authorization_code",
+            "code": code,
+            "redirect_uri": f"{base}/callback",
+            "client_id": client_id,
+            "code_verifier": payload["v"],
+        }
+        if client_secret:
+            form["client_secret"] = client_secret
+        try:
+            async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=15)) as s:
+                async with s.post(f"{iam_url}/v1/iam/oauth/token", data=form) as r:
+                    if r.status != 200:
+                        logger.error("Token exchange failed %s: %s", r.status, await r.text())
+                        return web.Response(status=502, text="Token exchange failed")
+                    tok = await r.json()
+        except Exception as e:
+            logger.error("Token exchange error: %s", e)
+            return web.Response(status=502, text="IAM unreachable")
+
+        access = tok.get("access_token")
+        if not access or await _validate(access) is None:
+            return web.Response(status=502, text="Invalid token from IAM")
+
+        target = payload.get("t", "/")
+        if not target.startswith("/"):
+            target = "/"
+        resp = web.HTTPFound(target)
+        resp.set_cookie(
+            _SESSION_COOKIE, access, max_age=int(tok.get("expires_in", 3600)),
+            httponly=True, secure=base.startswith("https"), samesite="Lax", path="/",
+        )
+        resp.del_cookie(_STATE_COOKIE, path="/")
+        return resp
 
     @web.middleware
     async def iam_auth_middleware(request: web.Request, handler):
         path = request.path
 
-        # Always allow public paths (static assets, health checks)
         if _is_public_path(path):
             return await handler(request)
 
-        # Localhost bypass for dev mode
         if localhost_bypass and _is_loopback_request(request):
             request["iam_user"] = {
-                "sub": "local",
-                "name": "Local User",
-                "org_id": os.environ.get("STUDIO_ORG_ID", "default"),
+                "sub": "local", "name": "Local User", "email": "",
+                "org_id": os.environ.get("STUDIO_ORG_ID", "default"), "avatar": "",
             }
             return await handler(request)
 
-        # Extract Bearer token
-        auth_header = request.headers.get("Authorization", "")
         token = None
+        auth_header = request.headers.get("Authorization", "")
         if auth_header.startswith("Bearer "):
             token = auth_header[7:]
-
-        # Also check cookie for browser sessions
         if not token:
-            token = request.cookies.get("hanzo_token") or request.cookies.get("access_token")
-
-        # Also check query param for WebSocket upgrade
-        if not token and request.path == "/ws":
+            token = request.cookies.get(_SESSION_COOKIE) or request.cookies.get("access_token")
+        if not token and path == "/ws":
             token = request.rel_url.query.get("token")
 
+        wants_html = "text/html" in request.headers.get("Accept", "")
+
         if not token:
-            # Return login redirect for browser requests, 401 for API
-            accept = request.headers.get("Accept", "")
-            if "text/html" in accept:
-                raise web.HTTPFound(f"{iam_url}/login?redirect_uri={request.url}")
-            return web.json_response(
-                {"error": "Authentication required", "iam_url": iam_url},
-                status=401,
-            )
+            if wants_html:
+                return _authorize_redirect(request)
+            return web.json_response({"error": "Authentication required", "iam_url": iam_url}, status=401)
 
-        # Check token cache first
-        ck = _cache_key(token)
-        cached = _token_cache.get(ck)
-        if cached and cached[1] > _time.monotonic():
-            request["iam_user"] = cached[0]
-            return await handler(request)
-
-        # Validate token against IAM
         try:
-            session = await _get_session()
-            async with session.get(
-                account_endpoint,
-                headers={"Authorization": f"Bearer {token}"},
-            ) as resp:
-                if resp.status != 200:
-                    _token_cache.pop(ck, None)
-                    return web.json_response(
-                        {"error": "Invalid or expired token"},
-                        status=401,
-                    )
-                raw = await resp.json()
-        except Exception as e:
-            logging.error(f"IAM validation error: {e}")
-            return web.json_response(
-                {"error": "Authentication service unavailable"},
-                status=503,
-            )
+            user = await _validate(token)
+        except Exception:
+            return web.json_response({"error": "Authentication service unavailable"}, status=503)
 
-        # Extract user/org context from the OIDC userinfo response (HIP-0111).
-        # OIDC userinfo returns flat standard claims: sub, email, name, picture.
-        # NOTE: org derivation relies on a non-standard "owner"/"organization"
-        # claim. The IAM server must be configured to emit it in the requested
-        # scopes; otherwise org_id falls back to "default".
-        user_info = {
-            "sub": raw.get("sub", raw.get("id", "")),
-            "name": raw.get("name", raw.get("displayName", "")),
-            "email": raw.get("email", ""),
-            "org_id": raw.get("owner", raw.get("organization", "default")),
-            "avatar": raw.get("picture", raw.get("avatar", "")),
-        }
+        if user is None:
+            if wants_html:
+                return _authorize_redirect(request)
+            return web.json_response({"error": "Invalid or expired token"}, status=401)
 
-        # Cache the validated token
-        _token_cache[ck] = (user_info, _time.monotonic() + _CACHE_TTL)
-
-        # Evict stale entries periodically (simple, not perfect)
-        if len(_token_cache) > 1000:
-            now = _time.monotonic()
-            stale = [k for k, v in _token_cache.items() if v[1] < now]
-            for k in stale:
-                del _token_cache[k]
-
-        request["iam_user"] = user_info
+        request["iam_user"] = user
         return await handler(request)
 
+    # Expose the callback handler so server.py can register the /callback route.
+    iam_auth_middleware.handle_callback = handle_callback
     return iam_auth_middleware

middleware/prompt_router.py

@@ -90,17 +90,18 @@ async def forward_to_worker(worker: WorkerInfo, json_data: dict) -> Optional[dic
     """
     Forward a prompt to a remote worker via HTTP POST.
 
-    The worker exposes /api/worker/execute which accepts the same JSON body
+    The worker exposes /v1/worker/execute which accepts the same JSON body
     as POST /prompt and returns the same response shape.
 
     Returns the worker's JSON response or None on failure.
     """
-    url = f"{worker.url.rstrip('/')}/api/worker/execute"
+    from middleware.worker_client import _worker_headers
+    url = f"{worker.url.rstrip('/')}/v1/worker/execute"
     logger.info("Forwarding prompt to worker %s at %s", worker.worker_id, url)
 
     try:
         session = await _get_session()
-        async with session.post(url, json=json_data) as resp:
+        async with session.post(url, json=json_data, headers=_worker_headers()) as resp:
             if resp.status == 200:
                 return await resp.json()
             logger.warning(

middleware/worker_client.py

@@ -23,6 +23,23 @@ logger = logging.getLogger(__name__)
 # Heartbeat interval in seconds
 HEARTBEAT_INTERVAL = 30
 
+# Shared secret authenticating worker<->coordinator traffic. These endpoints are
+# exempt from user IAM (they carry no user token), so without this a rogue caller
+# could register a worker URL to hijack an org's jobs. Sourced from KMS in cloud;
+# empty in local/dev (single trust domain) where the check is skipped.
+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)."""
+    if not WORKER_TOKEN:
+        return True
+    return request.headers.get("X-Worker-Token", "") == WORKER_TOKEN
+
+
+def _worker_headers() -> dict:
+    return {"X-Worker-Token": WORKER_TOKEN} if WORKER_TOKEN else {}
+
 
 def get_device_info() -> dict:
     """Detect local device capabilities for worker registration."""
@@ -91,10 +108,10 @@ class WorkerClient:
             **self._device_info,
         }
 
-        url = f"{self.coordinator_url}/api/compute/workers/register"
+        url = f"{self.coordinator_url}/v1/workers/register"
         try:
             session = await self._get_session()
-            async with session.post(url, json=payload) as resp:
+            async with session.post(url, json=payload, headers=_worker_headers()) as resp:
                 if resp.status == 200:
                     logger.debug("Heartbeat sent to %s", url)
                     return True
@@ -136,16 +153,16 @@ class WorkerClient:
 
 def add_worker_routes(routes: web.RouteTableDef, prompt_server):
     """
-    Add the /worker/execute endpoint that receives prompts from the coordinator.
-
-    Note: This is registered as /worker/execute (not /api/worker/execute) because
-    server.py's add_routes() auto-prefixes all RouteTableDef routes with /api.
-    The final path will be /api/worker/execute.
+    Add the /v1/worker/execute endpoint that receives prompts from the coordinator
+    (push model, for in-cluster reachable workers). server.py also serves it under
+    the /api alias. NAT'd workers use the pull model instead (see federation doc).
     """
 
-    @routes.post("/worker/execute")
+    @routes.post("/v1/worker/execute")
     async def worker_execute(request):
         """Execute a prompt forwarded by the coordinator."""
+        if not verify_worker_token(request):
+            return web.json_response({"error": "Unauthorized worker"}, status=401)
         try:
             json_data = await request.json()
         except Exception:

requirements.txt

@@ -24,6 +24,7 @@ av>=14.2.0
 comfy-kitchen>=0.2.7
 comfy-aimdo>=0.2.0
 requests
+PyJWT[crypto]>=2.7.0
 
 #non essential dependencies:
 kornia>=0.7.1

server.py

@@ -223,12 +223,14 @@ class PromptServer():
 
         middlewares = [cache_control, deprecation_warning]
 
+        self._iam_auth = None
         if args.enable_iam_auth:
             localhost_bypass = not args.no_localhost_bypass
-            middlewares.append(create_iam_auth_middleware(
+            self._iam_auth = create_iam_auth_middleware(
                 iam_url=args.iam_url,
                 localhost_bypass=localhost_bypass,
-            ))
+            )
+            middlewares.append(self._iam_auth)
 
         # Rate limiting (after auth so we have org context)
         middlewares.append(create_rate_limit_middleware(rpm=args.rate_limit_rpm))
@@ -249,6 +251,9 @@ class PromptServer():
 
         max_upload_size = round(args.max_upload_size * 1024 * 1024)
         self.app = web.Application(client_max_size=max_upload_size, middlewares=middlewares)
+        # OIDC Authorization Code callback (standard code flow, standalone app).
+        if self._iam_auth is not None:
+            self.app.router.add_get("/callback", self._iam_auth.handle_callback)
         self.sockets = dict()
         self.sockets_metadata = dict()
         self.web_root = (
@@ -1033,6 +1038,7 @@ class PromptServer():
                         if sensitive_val in extra_data:
                             sensitive[sensitive_val] = extra_data.pop(sensitive_val)
                     extra_data["create_time"] = int(time.time() * 1000)  # timestamp in milliseconds
+                    extra_data["org_id"] = self._get_org_id(request)  # tenancy: scope outputs to org
                     self.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute, sensitive))
                     response = {"prompt_id": prompt_id, "number": number, "node_errors": valid[3]}
                     return web.json_response(response)
@@ -1143,7 +1149,7 @@ class PromptServer():
 
             return web.json_response(config.to_dict())
 
-        @routes.get("/compute/workers")
+        @routes.get("/v1/workers")
         async def get_compute_workers(request):
             """List workers and their status for the org."""
             org_id = self._get_org_id(request)
@@ -1162,9 +1168,12 @@ class PromptServer():
                 workers.append(wd)
             return web.json_response({"workers": workers})
 
-        @routes.post("/compute/workers/register")
+        @routes.post("/v1/workers/register")
         async def post_worker_register(request):
-            """Worker self-registration / heartbeat endpoint."""
+            """Worker self-registration / heartbeat endpoint (coordinator trust)."""
+            from middleware.worker_client import verify_worker_token
+            if not verify_worker_token(request):
+                return web.json_response({"error": "Unauthorized worker"}, status=401)
             try:
                 data = await request.json()
             except Exception:

tests-unit/app_test/user_manager_system_user_test.py

@@ -38,9 +38,16 @@ def user_manager(mock_user_directory):
 
 @pytest.fixture
 def mock_request():
-    """Create a mock request object."""
+    """Create a mock request object.
+
+    A real aiohttp request is a Mapping whose .get("iam_user") is None when no
+    IAM middleware ran; MagicMock's default .get is truthy, which would send
+    get_request_user_id down the IAM branch and mask the studio-user header
+    path these tests exercise. Pin it to None.
+    """
     request = MagicMock()
     request.headers = {}
+    request.get = MagicMock(return_value=None)
     return request
 
 

tests-unit/folder_paths_test/org_scoping_test.py

@@ -0,0 +1,95 @@
+"""Tests for multi-tenant org scoping in folder_paths.py
+
+Covers:
+- get_public_user_directory(user, org): org-rooted user dir + System User block
+- The containment invariant user_manager relies on (regression for the bug
+  where org userdata always resolved outside the org root -> None)
+- Execution-org context scoping get_output_directory / temp / input
+- Single-tenant behavior is unchanged
+"""
+import os
+
+import pytest
+
+import folder_paths
+
+
+@pytest.fixture
+def tenant_env(tmp_path, monkeypatch):
+    """Point folder_paths at a temp base and enable multi-tenant mode."""
+    base = str(tmp_path)
+    monkeypatch.setattr(folder_paths, "base_path", base)
+    monkeypatch.setattr(folder_paths, "output_directory", os.path.join(base, "output"))
+    monkeypatch.setattr(folder_paths, "input_directory", os.path.join(base, "input"))
+    monkeypatch.setattr(folder_paths, "temp_directory", os.path.join(base, "temp"))
+    monkeypatch.setattr(folder_paths, "user_directory", os.path.join(base, "user"))
+    monkeypatch.setattr(folder_paths, "_org_id", None)
+    folder_paths.set_execution_org(None)
+    yield base
+    folder_paths.set_execution_org(None)
+
+
+def test_single_tenant_userdir_unchanged(tenant_env, monkeypatch):
+    monkeypatch.setattr(folder_paths, "_multi_tenant", False)
+    assert folder_paths.get_public_user_directory("default") == \
+        os.path.join(tenant_env, "user", "default")
+
+
+def test_single_tenant_output_unchanged(tenant_env, monkeypatch):
+    monkeypatch.setattr(folder_paths, "_multi_tenant", False)
+    folder_paths.set_execution_org("acme")  # ignored when not multi-tenant
+    assert folder_paths.get_output_directory() == os.path.join(tenant_env, "output")
+
+
+def test_system_user_blocked_with_org(tenant_env, monkeypatch):
+    monkeypatch.setattr(folder_paths, "_multi_tenant", True)
+    assert folder_paths.get_public_user_directory("__system", "acme") is None
+    assert folder_paths.get_public_user_directory("", "acme") is None
+    assert folder_paths.get_public_user_directory(None, "acme") is None
+
+
+def test_org_user_dir_is_org_rooted(tenant_env, monkeypatch):
+    monkeypatch.setattr(folder_paths, "_multi_tenant", True)
+    path = folder_paths.get_public_user_directory("alice", "acme")
+    assert path == os.path.join(tenant_env, "orgs", "acme", "user", "alice")
+
+
+def test_containment_invariant(tenant_env, monkeypatch):
+    """user_manager requires user_root to be a child of the org user root.
+
+    This is the exact relation that regressed: get_org_user_directory(org) and
+    get_public_user_directory(user, org) must share the org root as commonpath.
+    """
+    monkeypatch.setattr(folder_paths, "_multi_tenant", True)
+    root = folder_paths.get_org_user_directory("acme")
+    user_root = folder_paths.get_public_user_directory("alice", "acme")
+    assert os.path.commonpath((root, user_root)) == root
+
+
+def test_two_orgs_isolated(tenant_env, monkeypatch):
+    monkeypatch.setattr(folder_paths, "_multi_tenant", True)
+    a = folder_paths.get_public_user_directory("alice", "acme")
+    b = folder_paths.get_public_user_directory("alice", "globex")
+    assert a != b
+    # Distinct org roots: they share only the orgs/ parent, never each other.
+    assert os.path.commonpath((a, b)) == os.path.join(tenant_env, "orgs")
+
+
+def test_execution_org_scopes_output(tenant_env, monkeypatch):
+    monkeypatch.setattr(folder_paths, "_multi_tenant", True)
+    folder_paths.set_execution_org("acme")
+    assert folder_paths.get_output_directory() == \
+        os.path.join(tenant_env, "orgs", "acme", "output")
+    assert folder_paths.get_temp_directory() == \
+        os.path.join(tenant_env, "orgs", "acme", "temp")
+    assert folder_paths.get_input_directory() == \
+        os.path.join(tenant_env, "orgs", "acme", "input")
+
+
+def test_execution_org_reset(tenant_env, monkeypatch):
+    monkeypatch.setattr(folder_paths, "_multi_tenant", True)
+    folder_paths.set_execution_org("acme")
+    assert "acme" in folder_paths.get_output_directory()
+    folder_paths.set_execution_org(None)
+    # No active execution org and no process org -> base output dir
+    assert folder_paths.get_output_directory() == os.path.join(tenant_env, "output")

tests-unit/middleware_test/iam_auth_test.py

@@ -0,0 +1,144 @@
+"""Tests for the IAM auth middleware's pure verification + session helpers.
+
+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 time
+
+import jwt
+import pytest
+from cryptography.hazmat.primitives.asymmetric import rsa
+
+from middleware import iam_auth_middleware as m
+
+ISS = "https://hanzo.id"
+AUD = "hanzo-studio"
+
+
+@pytest.fixture(scope="module")
+def keys():
+    priv = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+    return priv, priv.public_key()
+
+
+def _token(priv, **overrides):
+    now = int(time.time())
+    claims = {
+        "sub": "hanzo/alice",
+        "iss": ISS,
+        "aud": AUD,
+        "exp": now + 300,
+        "iat": now,
+        "owner": "acme",
+        "name": "Alice",
+        "email": "alice@acme.com",
+        "picture": "https://cdn/x.png",
+    }
+    claims.update(overrides)
+    for k in [k for k, v in claims.items() if v is None]:
+        del claims[k]
+    return jwt.encode(claims, priv, algorithm="RS256", headers={"kid": "test"})
+
+
+def test_valid_token_extracts_org(keys):
+    priv, pub = keys
+    user = m.verify_jwt(_token(priv), pub, ISS, AUD)
+    assert user is not None
+    assert user["org_id"] == "acme"
+    assert user["sub"] == "hanzo/alice"
+    assert user["name"] == "Alice"
+    assert user["email"] == "alice@acme.com"
+    assert user["avatar"] == "https://cdn/x.png"
+
+
+def test_organization_claim_fallback(keys):
+    priv, pub = keys
+    tok = _token(priv, owner=None, organization="globex")
+    assert m.verify_jwt(tok, pub, ISS, AUD)["org_id"] == "globex"
+
+
+def test_org_defaults_when_absent(keys):
+    priv, pub = keys
+    tok = _token(priv, owner=None, organization=None)
+    assert m.verify_jwt(tok, pub, ISS, AUD)["org_id"] == "default"
+
+
+def test_preferred_username_fallback(keys):
+    priv, pub = keys
+    tok = _token(priv, name=None, preferred_username="alice")
+    assert m.verify_jwt(tok, pub, ISS, AUD)["name"] == "alice"
+
+
+def test_wrong_audience_rejected(keys):
+    priv, pub = keys
+    assert m.verify_jwt(_token(priv, aud="someone-else"), pub, ISS, AUD) is None
+
+
+def test_wrong_issuer_rejected(keys):
+    priv, pub = keys
+    assert m.verify_jwt(_token(priv, iss="https://evil.example"), pub, ISS, AUD) is None
+
+
+def test_expired_rejected(keys):
+    priv, pub = keys
+    assert m.verify_jwt(_token(priv, exp=int(time.time()) - 10), pub, ISS, AUD) is None
+
+
+def test_missing_exp_rejected(keys):
+    priv, pub = keys
+    assert m.verify_jwt(_token(priv, exp=None), pub, ISS, AUD) is None
+
+
+def test_bad_signature_rejected(keys):
+    priv, _pub = keys
+    other = rsa.generate_private_key(public_exponent=65537, key_size=2048).public_key()
+    assert m.verify_jwt(_token(priv), other, ISS, AUD) is None
+
+
+def test_aud_array_form_accepted(keys):
+    """golang-jwt serializes a single audience as a string but multi as array;
+    accept both."""
+    priv, pub = keys
+    assert m.verify_jwt(_token(priv, aud=[AUD, "other"]), pub, ISS, AUD)["org_id"] == "acme"
+
+
+# --- signed OAuth state (CSRF) ---
+
+def test_state_roundtrip():
+    payload = {"n": "nonce", "v": "verifier", "t": "/workflows", "exp": int(time.time()) + 60}
+    blob = m._sign_state(payload, "secret")
+    got = m._verify_state(blob, "secret")
+    assert got["n"] == "nonce" and got["t"] == "/workflows"
+
+
+def test_state_tamper_rejected():
+    blob = m._sign_state({"n": "n", "exp": int(time.time()) + 60}, "secret")
+    body, mac = blob.split(".", 1)
+    assert m._verify_state(f"{body}x.{mac}", "secret") is None
+
+
+def test_state_wrong_secret_rejected():
+    blob = m._sign_state({"n": "n", "exp": int(time.time()) + 60}, "secret")
+    assert m._verify_state(blob, "other") is None
+
+
+def test_state_expired_rejected():
+    blob = m._sign_state({"n": "n", "exp": int(time.time()) - 1}, "secret")
+    assert m._verify_state(blob, "secret") is None
+
+
+# --- path exemptions ---
+
+@pytest.mark.parametrize("path", [
+    "/", "/callback", "/health", "/api/features",
+    "/assets/x.js", "/main.css", "/logo.svg",
+    "/v1/workers/register", "/v1/worker/execute",
+])
+def test_public_paths(path):
+    assert m._is_public_path(path) is True
+
+
+@pytest.mark.parametrize("path", ["/prompt", "/api/prompt", "/userdata/x", "/v1/compute/config"])
+def test_protected_paths(path):
+    assert m._is_public_path(path) is False