mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-10-15 19:32:21 +02:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
eb940f24a3 |
@@ -178,6 +178,14 @@ app.add_middleware(SecurityHeadersMiddleware)
|
||||
import asyncio as _asyncio
|
||||
from starlette.middleware.base import BaseHTTPMiddleware as _BaseHTTPMiddleware
|
||||
from starlette.responses import JSONResponse as _JSONResponse
|
||||
import anyio as _anyio
|
||||
from starlette.requests import ClientDisconnect as _ClientDisconnect
|
||||
|
||||
# SSE endpoints (e.g. /api/editor/bridge/events) end/abort mid-stream when the
|
||||
# browser disconnects; Starlette's BaseHTTPMiddleware surfaces that as
|
||||
# anyio.EndOfStream (and sometimes ClientDisconnect). Those are client-side
|
||||
# events, not server errors, so we map them to a 499 instead of a 500 traceback.
|
||||
_STREAM_CLIENT_EXC = (_anyio.EndOfStream, _ClientDisconnect)
|
||||
|
||||
REQUEST_HARD_TIMEOUT = float(os.getenv("REQUEST_HARD_TIMEOUT", "45"))
|
||||
_TIMEOUT_EXEMPT_PREFIXES = (
|
||||
@@ -190,6 +198,7 @@ _TIMEOUT_EXEMPT_PREFIXES = (
|
||||
"/api/cookbook/setup", # remote pacman/apt installs
|
||||
"/api/upload", # large files
|
||||
"/api/image", # diffusion proxies (inpaint/harmonize/upscale/etc.) — own 120s httpx timeout
|
||||
"/api/editor", # editor agent/bridge: generative steps can take 30-110s (SAM + local diffusion)
|
||||
"/api/memory/audit", # retains own 120s LLM inactivity timeout
|
||||
)
|
||||
|
||||
@@ -225,15 +234,43 @@ class _InteractiveActivityMiddleware(_BaseHTTPMiddleware):
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
# Heavy image ops that can run 30-150 s; a deploy restart must not cut these off.
|
||||
_HEAVY_IMAGE_OP_PREFIXES = (
|
||||
"/api/image/inpaint",
|
||||
"/api/image/replace-background",
|
||||
"/api/image/outpaint",
|
||||
"/api/image/extend",
|
||||
"/api/image/remove-bg",
|
||||
"/api/image/mask",
|
||||
"/api/image/upscale-local",
|
||||
"/api/image/denoise",
|
||||
"/api/image/enhance-face",
|
||||
"/api/gallery/ai-upscale",
|
||||
"/api/gallery/style-transfer",
|
||||
"/api/editor/agent",
|
||||
)
|
||||
|
||||
|
||||
class _SlowRequestLogMiddleware(_BaseHTTPMiddleware):
|
||||
async def dispatch(self, request, call_next):
|
||||
from src import activity
|
||||
path = request.url.path or ""
|
||||
heavy = any(path.startswith(p) for p in _HEAVY_IMAGE_OP_PREFIXES)
|
||||
if heavy:
|
||||
activity.begin_image_op(path)
|
||||
start = time.perf_counter()
|
||||
status = 500
|
||||
try:
|
||||
response = await call_next(request)
|
||||
status = getattr(response, "status_code", 0) or 0
|
||||
return response
|
||||
except _STREAM_CLIENT_EXC:
|
||||
# Client disconnected / SSE stream ended — not a server error.
|
||||
status = 499
|
||||
return _JSONResponse({"detail": "client disconnected"}, status_code=499)
|
||||
finally:
|
||||
if heavy:
|
||||
activity.end_image_op(path)
|
||||
elapsed = time.perf_counter() - start
|
||||
try:
|
||||
threshold = float(os.getenv("ODYSSEUS_SLOW_REQUEST_LOG_SECONDS", "0.75") or "0.75")
|
||||
@@ -821,6 +858,10 @@ app.include_router(setup_gallery_routes())
|
||||
from routes.editor_draft_routes import setup_editor_draft_routes
|
||||
app.include_router(setup_editor_draft_routes())
|
||||
|
||||
# Natural-language editor planning (planner wired to the configured chat model)
|
||||
from routes.editor_agent_routes import setup_editor_agent_routes
|
||||
app.include_router(setup_editor_agent_routes())
|
||||
|
||||
# Scheduled tasks + event bus
|
||||
from src.task_scheduler import TaskScheduler
|
||||
task_scheduler = TaskScheduler(session_manager)
|
||||
@@ -992,12 +1033,21 @@ async def serve_login(request: Request):
|
||||
@app.get("/api/version")
|
||||
async def get_version():
|
||||
from core.constants import APP_BUILD_VERSION, APP_SOURCE_COMMIT, APP_VERSION
|
||||
from src import activity
|
||||
return {
|
||||
"version": APP_VERSION,
|
||||
"build": APP_BUILD_VERSION,
|
||||
"source_commit": APP_SOURCE_COMMIT,
|
||||
"busy_image_ops": activity.heavy_image_ops_active(),
|
||||
}
|
||||
|
||||
|
||||
@app.get("/api/busy")
|
||||
async def get_busy():
|
||||
"""Report in-flight heavy work so a deploy can wait instead of dropping it."""
|
||||
from src import activity
|
||||
return activity.activity()
|
||||
|
||||
@app.get("/api/health")
|
||||
async def health_check() -> Dict[str, str]:
|
||||
return {"status": "healthy", "timestamp": datetime.now(timezone.utc).isoformat()}
|
||||
@@ -1436,4 +1486,35 @@ if __name__ == "__main__":
|
||||
bind_host = os.getenv("APP_BIND", "127.0.0.1")
|
||||
bind_port = int(os.getenv("APP_PORT", "7011"))
|
||||
|
||||
uvicorn.run(app, host=bind_host, port=bind_port, log_level="info")
|
||||
# Give in-flight requests time to finish on SIGTERM instead of being cut
|
||||
# off mid-flight. Diffusion edits run 60-120 s; outpaint/extend can run up
|
||||
# to ~330 s (provider cap 900 s), so default the graceful window to 900 s.
|
||||
# systemd's TimeoutStopSec must exceed this (see the graceful-stop drop-in).
|
||||
#
|
||||
# We hook uvicorn's signal handler so a SIGTERM also flips src.shutdown:
|
||||
# SSE / streaming generators check it and end promptly, so an open browser
|
||||
# stream cannot stall a restart for the whole window. The long window stays
|
||||
# a safety net for genuine in-flight image ops.
|
||||
from src import shutdown as _shutdown
|
||||
|
||||
config = uvicorn.Config(
|
||||
app,
|
||||
host=bind_host,
|
||||
port=bind_port,
|
||||
log_level="info",
|
||||
timeout_graceful_shutdown=int(os.getenv("ODYSSEUS_GRACEFUL_SHUTDOWN_TIMEOUT", "900") or 900),
|
||||
)
|
||||
server = uvicorn.Server(config)
|
||||
|
||||
# uvicorn's capture_signals() wires SIGINT/SIGTERM to self.handle_exit, so
|
||||
# wrapping that method lets us set the shutdown flag before the normal exit
|
||||
# path runs (this uvicorn version has no install_signal_handlers kwarg).
|
||||
_uvicorn_handle_exit = server.handle_exit
|
||||
|
||||
def _handle_shutdown_signal(signum, frame):
|
||||
_shutdown.request_shutdown()
|
||||
_uvicorn_handle_exit(signum, frame)
|
||||
|
||||
server.handle_exit = _handle_shutdown_signal
|
||||
|
||||
server.run()
|
||||
|
||||
@@ -0,0 +1,201 @@
|
||||
# Editor Agent API
|
||||
|
||||
How an agent (LLM or script) drives the PhotoCraft editor. There are three layers:
|
||||
|
||||
1. **Control API** — `window.geEditorControl`, a stable in-page command surface.
|
||||
2. **NL planner** — `window.geEditorAI`, turns a prompt into control-API steps (LLM or basic mode).
|
||||
3. **Bridge** — server-mediated SSE channel so a *server-side* agent can drive an **open** editor.
|
||||
|
||||
All three are owner-scoped by the app session; nothing is exposed to unauthenticated users.
|
||||
|
||||
---
|
||||
|
||||
## 1. Control API (`window.geEditorControl`)
|
||||
|
||||
Installed when the editor opens. Version `EDITOR_CONTROL_VERSION`.
|
||||
|
||||
```js
|
||||
await window.geEditorControl.invoke(op, params) // { ok, op, result } | { ok:false, op, error }
|
||||
await window.geEditorControl.run([{op, params}]) // sequential, stops on first failure
|
||||
window.geEditorControl.list() // [{op, description}]
|
||||
```
|
||||
|
||||
`invoke` is async. Every op that mutates the document pushes **one** undo entry via
|
||||
`_saveState`, so a single op is one undo step.
|
||||
|
||||
### Op catalog
|
||||
|
||||
| Op | Params | Notes |
|
||||
| --- | --- | --- |
|
||||
| `document.info` | – | size, layers, mode |
|
||||
| `document.importImage` | `{dataUrl,name,raster}` | load a raster image as a new layer |
|
||||
| `layers.list` | – | top-most first |
|
||||
| `layers.select` | `{id\|index\|name}` | index is top-most-first |
|
||||
| `layers.addRaster` | – | new empty pixel layer |
|
||||
| `layers.addAdjustment` | `{type,params}` | adjustment layer |
|
||||
| `layers.duplicate` / `layers.delete` | – | on the active/selected layer(s) |
|
||||
| `layers.setOpacity` | `{id?,opacity 0-1}` | |
|
||||
| `layers.setBlend` | `{id?,blendMode\|mode}` | |
|
||||
| `layers.group` / `layers.ungroup` | – | |
|
||||
| `selection.selectAll` / `deselect` / `invert` | – | |
|
||||
| `selection.rectangle` / `selection.ellipse` | `{x,y,w,h}` | document space |
|
||||
| `selection.subject` | `{label}` | SAM text grounding (falls back to a centred rect) |
|
||||
| `selection.face` | – | SAM "face" (falls back to an upper-centre ellipse) |
|
||||
| `selection.sky` | – | SAM "sky" (falls back to a top-band colour range) |
|
||||
| `selection.background` | `{fuzziness}` | border colour union |
|
||||
| `selection.samText` | `{text}` | SAM text grounding only (throws without SAM) |
|
||||
| `selection.colorRange` | `{color,fuzziness}` | |
|
||||
| `selection.hasMask` | – | `{active, space}` |
|
||||
| `adjustments.apply` | `{type,params}` | adjustment layer; camelCase types accepted |
|
||||
| `adjustments.applyMasked` | `{brightness,contrast,saturation,warmth,invert}` | pixel pass, masked if a selection exists else whole layer |
|
||||
| `filters.apply` | `{id}` | opens the slider dialog and applies |
|
||||
| `filters.applyMasked` | `{id,radius,amount,invert}` | `id`: blur-gaussian, sharpen, unsharp-mask |
|
||||
| `transform.start` / `transform.apply` / `transform.cancel` | – | |
|
||||
| `text.add` | `{text,x,y,color,size,font,align}` | |
|
||||
| `history.undo` / `history.redo` | – | |
|
||||
| `history.mark` | – | returns current undo depth (for grouping) |
|
||||
| `history.collapse` | `{mark,label}` | collapses entries above `mark` into one |
|
||||
| `export.png` | – | `{dataUrl}` of the flattened document |
|
||||
| `export.gallery` | – | saves a copy back to the gallery |
|
||||
| `screenshot` | – | `{dataUrl,width,height}` of the composited canvas |
|
||||
|
||||
Param conventions the planner must respect:
|
||||
|
||||
- `adjustments.applyMasked`: `brightness`/`contrast`/`warmth` are **-100..100** (0 neutral);
|
||||
`saturation` is a **multiplier** (1 neutral; values >3 are treated as percentages).
|
||||
- `adjustments.apply` `type` aliases are normalised (`blackAndWhite` -> `black-white`).
|
||||
- `color-balance` params are `{shadows:{r,g,b}, midtones:{r,g,b}, highlights:{r,g,b}}`;
|
||||
a bare number is coerced to `{r:n,g:n,b:n}`.
|
||||
- A masked op with **no** active selection degrades to a whole-layer edit (does not throw).
|
||||
|
||||
---
|
||||
|
||||
## 2. NL planner (`window.geEditorAI`)
|
||||
|
||||
```js
|
||||
const res = await window.geEditorAI.execute('blur the background behind the person', { onStep });
|
||||
// exec: window.geEditorAI.execute('blur the background behind the person', { onStep });
|
||||
// res = { ok, prompt, reply, planner, basicMode, steps, trace, failed }
|
||||
```
|
||||
|
||||
Flow (`static/js/editor/nl-edit.js`):
|
||||
|
||||
1. Read document state + a **before** screenshot.
|
||||
2. Ask the **server planner** (`POST /api/editor/plan`) with the prompt, state and
|
||||
screenshot. The client passes its strong guidance via `state.planner_guidance`.
|
||||
If the route is unavailable it falls back to `localPlan` (regex — labelled
|
||||
**"basic mode"** in the UI).
|
||||
3. `sanitizeSteps()`: drop no-op/export steps and, if a masked op has no preceding
|
||||
`selection.*`, insert the region implied by the prompt (face/sky/background/subject).
|
||||
4. `history.mark`, then run each step, screenshotting after each one.
|
||||
5. **Self-check** (server plans only, max 3 rounds): send the after screenshot back in
|
||||
`mode:"check"` with the original intent. If not satisfied, **revert** the previous
|
||||
attempt (undo its history entries) and apply the model's **full corrected** plan.
|
||||
6. `history.collapse` — the whole request becomes a single undoable step.
|
||||
|
||||
UI: the Quick Edit box routes multi-clause / "and ..." prompts to this executor and
|
||||
streams `Step i/n: <op>` into its status line; a `(basic mode)` suffix marks the offline planner.
|
||||
|
||||
### Server route
|
||||
|
||||
`POST /api/editor/plan` (HARNESS, `routes/editor_agent_routes.py`)
|
||||
|
||||
```
|
||||
{ prompt, state, mode?: "plan"|"check", screenshot?: "data:image/..." }
|
||||
-> { reply, planner, steps: [{op, params}], done? }
|
||||
```
|
||||
|
||||
Owner-scoped, authenticated. Uses the configured default/utility model with vision.
|
||||
If the model call is unavailable it returns a deterministic `planner:"fallback"` plan.
|
||||
|
||||
---
|
||||
|
||||
## 3. Bridge (server agent drives an open editor)
|
||||
|
||||
`static/js/editor/editor-bridge.js` subscribes the open editor to the server queue.
|
||||
Enabled automatically when the editor opens (`window.__geEditorBridge`).
|
||||
|
||||
- Subscribe: `GET /api/editor/bridge/events?session=<id>` — SSE, `data: {id,op,params}` (keepalive every 15s).
|
||||
- The page runs the op via the control API and replies:
|
||||
`POST /api/editor/bridge/result {id,ok,result,error,screenshot,ms}`.
|
||||
- Enqueue a command (server side): `POST /api/editor/bridge/command {session, op|steps, wait, timeout}`.
|
||||
With `wait:true` the call blocks and returns the editor's result.
|
||||
|
||||
Round-trip example (authenticated; `<session>` = `window.__geEditorBridge.session`):
|
||||
|
||||
```bash
|
||||
curl -s -X POST http://127.0.0.1:7011/api/editor/bridge/command \
|
||||
-H 'Content-Type: application/json' -b "odysseus_session=$TOKEN" \
|
||||
-d '{"session":"default","op":"layers.addRaster","wait":true,"timeout":8000}'
|
||||
# -> {"id":"...","ok":true,"result":{...},"screenshot":"data:image/png;base64,..."}
|
||||
```
|
||||
|
||||
Session scoping is `owner:session`, so an agent can only drive editors owned by the
|
||||
same user. The bridge is a **browser-only** flow: it requires an editor page to be open.
|
||||
|
||||
---
|
||||
|
||||
## 4. SAM dependency
|
||||
|
||||
`selection.subject` / `selection.face` / `selection.sky` / `selection.samText` use the
|
||||
server SAM mask tool (`POST /api/image/mask`). When the SAM tools are **not installed**
|
||||
the editor raises *"Object mask tools are not installed"* and the control op falls back
|
||||
to a **coarse approximation** (centred rect / upper-centre ellipse / top-band colour range).
|
||||
On a plain background the approximation is acceptable; on real photos it leaks, producing
|
||||
visible rectangular/elliptical artifacts. Install the Cookbook SAM mask dependencies for
|
||||
true subject/face/sky segmentation.
|
||||
|
||||
`selection.colorRange` and `selection.background` never need SAM (pure colour matching),
|
||||
but they are also coarse on busy scenes.
|
||||
|
||||
---
|
||||
|
||||
## 5. Two ways for an agent to use this — pick one
|
||||
|
||||
### Option A — built-in tool (in-process)
|
||||
|
||||
Register an editor-control tool inside the app (same process as the model) that calls the
|
||||
bridge/control API directly. Uses the app's normal login; nothing new is exposed. This is
|
||||
the rejected stdio-MCP design's replacement: a server-side tool that talks to the bridge
|
||||
needs either an internal token (auth-bypass) or to run in-process. **Decision pending (user).**
|
||||
|
||||
### Option B — browser-only flow (works today)
|
||||
|
||||
No new server tool. An external agent drives an open editor purely through the bridge
|
||||
endpoints using the logged-in owner's session. Restricted to the editor the user has open.
|
||||
This is the flow demonstrated above and requires **no** auth-bypass.
|
||||
|
||||
A stdio `photo_editor` MCP server was prototyped and **rejected** because the built-in MCP
|
||||
subprocess has no internal token; wiring it would require expanding an auth bypass.
|
||||
The two viable options above are left for the user to choose.
|
||||
|
||||
---
|
||||
|
||||
## 6. Known limitations / honest grading
|
||||
|
||||
Live run against real 7011 with the server LLM planner (7 prompts): all 7 executed
|
||||
(`planner:"llm"`). Graded by eye:
|
||||
|
||||
- PASS: black & white portrait; warm cinematic; sharper + more vivid; text title.
|
||||
- FAIL (region-masked, SAM-bound): "dramatic sky" (global over-saturation instead of a
|
||||
dramatic, masked sky), "blur the background" (coarse mask leaks; visible rectangle),
|
||||
"brighten the face" (coarse ellipse/rect artifact).
|
||||
|
||||
Whole-image edits are reliable; subject/face/sky region edits need SAM. The self-check
|
||||
loop works (revert-and-replan) but the region cases spend all 3 iterations.
|
||||
|
||||
## 7. Theming (control kit)
|
||||
|
||||
The editor's interactive controls (ranges, number fields, selects, checkboxes) share a
|
||||
single themed kit: `static/js/editor/control-kit.js` adds `.ge-ctl-*` classes scoped to
|
||||
the editor roots (`#gallery-editor-container`, `.ge-adj-popup`, `.ge-filter-overlay`,
|
||||
`.ge-canvas-size-overlay`, `.ge-float-window`, `.ge-transform-popup`), and
|
||||
`static/css/13-image-editor.css` (section "Editor control kit") styles them from the app
|
||||
design tokens (`static/css/00-tokens.css`) to match the app sliders/selects. It keeps
|
||||
range fills in sync and attaches a **scrub-on-label** affordance to numeric fields
|
||||
(drag the field's label to change the value; arrow keys work, Shift = x10).
|
||||
|
||||
**The PhotoCraft editor is dark-only.** `:root.light` flips the global tokens, but
|
||||
`13-image-editor.css` has no `.light` handling and the editor relies on dark surfaces,
|
||||
so light mode renders a garbled mix. Unlike the rest of Odysseus, the editor follows
|
||||
Photoshop and is not offered in a light theme.
|
||||
File diff suppressed because it is too large
Load Diff
+16
-153
@@ -17,8 +17,6 @@ import base64
|
||||
import time
|
||||
import imaplib
|
||||
import smtplib
|
||||
import socket
|
||||
import ssl
|
||||
import email as email_mod
|
||||
import email.header
|
||||
import email.utils
|
||||
@@ -38,7 +36,6 @@ from typing import Optional, List
|
||||
|
||||
from src.auth_helpers import _auth_disabled, get_current_user
|
||||
from src.secret_storage import decrypt as _decrypt
|
||||
from src.url_safety import OutboundAddressBlocked, connect_outbound_tcp
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -163,94 +160,6 @@ def _smtp_security_mode(cfg: dict) -> str:
|
||||
return "ssl"
|
||||
|
||||
|
||||
# Raised when a mail host resolves into a denied range (see _mail_private_blocked).
|
||||
MailServerAddressBlocked = OutboundAddressBlocked
|
||||
|
||||
|
||||
def _mail_private_blocked(owner: str | None) -> bool:
|
||||
"""Whether private/loopback/shared mail destinations are denied for *owner*.
|
||||
|
||||
Link-local (cloud metadata), multicast, reserved and unspecified addresses
|
||||
are always denied (src/url_safety.py). Private ranges are where a mail
|
||||
tester or saved account becomes an internal-network probe, so they are
|
||||
only allowed for an admin or in single-user mode (AUTH_ENABLED=false), the
|
||||
same principals trusted with host-level tools. Operators of multi-user
|
||||
deployments with a LAN mail server opt in with EMAIL_ALLOW_PRIVATE_IPS=true;
|
||||
EMAIL_BLOCK_PRIVATE_IPS=true denies private ranges for everyone.
|
||||
"""
|
||||
if os.getenv("EMAIL_BLOCK_PRIVATE_IPS", "").strip().lower() == "true":
|
||||
return True
|
||||
if os.getenv("EMAIL_ALLOW_PRIVATE_IPS", "").strip().lower() == "true":
|
||||
return False
|
||||
from src.tool_security import owner_is_admin_or_single_user
|
||||
|
||||
return not owner_is_admin_or_single_user((owner or "").strip() or None)
|
||||
|
||||
|
||||
def _mail_socket_timeout(timeout):
|
||||
return socket._GLOBAL_DEFAULT_TIMEOUT if timeout is None else timeout
|
||||
|
||||
|
||||
class _PolicyIMAP4(imaplib.IMAP4):
|
||||
"""IMAP4 whose socket comes from connect_outbound_tcp.
|
||||
|
||||
The mail host is resolved once and connected to only at the addresses the
|
||||
address policy approved, so DNS rebinding between check and connect is
|
||||
impossible. STARTTLS still verifies against ``self.host``.
|
||||
"""
|
||||
|
||||
def __init__(self, host, port, *, block_private: bool, timeout=None):
|
||||
self._block_private = block_private
|
||||
super().__init__(host, port, timeout=timeout)
|
||||
|
||||
def _create_socket(self, timeout):
|
||||
if timeout is not None and not timeout:
|
||||
raise ValueError("Non-blocking socket (timeout=0) is not supported")
|
||||
return connect_outbound_tcp(
|
||||
self.host, self.port, timeout=_mail_socket_timeout(timeout), block_private=self._block_private,
|
||||
)
|
||||
|
||||
|
||||
class _PolicyIMAP4_SSL(imaplib.IMAP4_SSL):
|
||||
"""IMAP4_SSL over a policy-checked socket; TLS SNI/verification use the hostname."""
|
||||
|
||||
def __init__(self, host, port, *, block_private: bool, timeout=None, ssl_context=None):
|
||||
self._block_private = block_private
|
||||
super().__init__(host, port, ssl_context=ssl_context, timeout=timeout)
|
||||
|
||||
def _create_socket(self, timeout):
|
||||
sock = _PolicyIMAP4._create_socket(self, timeout)
|
||||
return self.ssl_context.wrap_socket(sock, server_hostname=self.host)
|
||||
|
||||
|
||||
class _PolicySMTP(smtplib.SMTP):
|
||||
"""SMTP whose socket comes from connect_outbound_tcp (see _PolicyIMAP4)."""
|
||||
|
||||
def __init__(self, host="", port=0, *, block_private: bool, **kwargs):
|
||||
self._block_private = block_private
|
||||
super().__init__(host, port, **kwargs)
|
||||
|
||||
def _get_socket(self, host, port, timeout):
|
||||
if timeout is not None and not timeout:
|
||||
raise ValueError("Non-blocking socket (timeout=0) is not supported")
|
||||
return connect_outbound_tcp(
|
||||
host, port, timeout=timeout, block_private=self._block_private,
|
||||
source_address=self.source_address,
|
||||
)
|
||||
|
||||
|
||||
class _PolicySMTP_SSL(smtplib.SMTP_SSL):
|
||||
"""SMTP_SSL over a policy-checked socket; TLS SNI/verification use the hostname."""
|
||||
|
||||
def __init__(self, host="", port=0, *, block_private: bool, **kwargs):
|
||||
self._block_private = block_private
|
||||
super().__init__(host, port, **kwargs)
|
||||
|
||||
def _get_socket(self, host, port, timeout):
|
||||
sock = _PolicySMTP._get_socket(self, host, port, timeout)
|
||||
return self.context.wrap_socket(sock, server_hostname=self._host)
|
||||
|
||||
|
||||
def _send_smtp_message(cfg: dict, from_addr: str, recipients: list[str], message: str | bytes, timeout: int = 30) -> None:
|
||||
"""Send through SMTP using the configured transport security mode."""
|
||||
host = cfg["smtp_host"]
|
||||
@@ -269,15 +178,14 @@ def _send_smtp_message(cfg: dict, from_addr: str, recipients: list[str], message
|
||||
smtp.login(user, password)
|
||||
|
||||
security = _smtp_security_mode(cfg)
|
||||
block_private = _mail_private_blocked(cfg.get("owner"))
|
||||
|
||||
if security == "ssl":
|
||||
with _PolicySMTP_SSL(host, port, timeout=timeout, block_private=block_private) as smtp:
|
||||
with smtplib.SMTP_SSL(host, port, timeout=timeout) as smtp:
|
||||
_auth_smtp(smtp)
|
||||
smtp.sendmail(from_addr, recipients, message)
|
||||
return
|
||||
|
||||
with _PolicySMTP(host, port, timeout=timeout, block_private=block_private) as smtp:
|
||||
with smtplib.SMTP(host, port, timeout=timeout) as smtp:
|
||||
if security == "starttls":
|
||||
smtp.starttls()
|
||||
_auth_smtp(smtp)
|
||||
@@ -316,39 +224,6 @@ def _friendly_email_auth_error(protocol: str, host: str, error: object) -> str:
|
||||
return raw[:200]
|
||||
|
||||
|
||||
# Bound at import: callers (and tests) may swap out imaplib.IMAP4 itself.
|
||||
_IMAP_ABORT = imaplib.IMAP4.abort
|
||||
|
||||
|
||||
def _mail_connection_test_error(protocol: str, host: str, error: BaseException) -> str:
|
||||
"""User-facing result for a failed IMAP/SMTP connection test.
|
||||
|
||||
Transport failures are reported by category, never with the peer's own
|
||||
bytes: a non-mail service answering on the chosen host/port would
|
||||
otherwise have its banner echoed back (an internal-service fingerprinting
|
||||
oracle). Errors from a server that does speak the protocol (login/auth
|
||||
rejections) keep their text via _friendly_email_auth_error, since that is
|
||||
what the user needs to fix their settings.
|
||||
"""
|
||||
if isinstance(error, MailServerAddressBlocked):
|
||||
return f"{protocol} server address is not allowed by this server's network policy"
|
||||
if isinstance(error, (_IMAP_ABORT, smtplib.SMTPConnectError, smtplib.SMTPServerDisconnected)):
|
||||
return f"{protocol} server did not respond like an {protocol} server"
|
||||
if isinstance(error, ssl.SSLCertVerificationError):
|
||||
return f"{protocol} TLS certificate verification failed"
|
||||
if isinstance(error, ssl.SSLError):
|
||||
return f"{protocol} TLS handshake failed; check the port and security setting"
|
||||
if isinstance(error, (socket.timeout, TimeoutError)):
|
||||
return f"{protocol} connection timed out"
|
||||
if isinstance(error, ConnectionRefusedError):
|
||||
return f"{protocol} connection refused"
|
||||
if isinstance(error, socket.gaierror):
|
||||
return f"{protocol} server name could not be resolved"
|
||||
if isinstance(error, OSError) and not isinstance(error, smtplib.SMTPException):
|
||||
return f"{protocol} connection failed"
|
||||
return _friendly_email_auth_error(protocol, host, error)
|
||||
|
||||
|
||||
def _strip_think(text: str) -> str:
|
||||
"""Email-flavored think strip — thin wrapper over the central helper.
|
||||
|
||||
@@ -810,21 +685,18 @@ def _ensure_sender_signatures_table(conn):
|
||||
_lg.getLogger(__name__).warning(f"sender_signatures owner-migration skipped: {_mig_e}")
|
||||
|
||||
|
||||
def attachment_extract_dir(folder: str, uid: str, *, owner: str, account_id: str | None) -> Path:
|
||||
"""Containment-safe extraction directory for one message's attachments.
|
||||
def attachment_extract_dir(folder: str, uid: str) -> Path:
|
||||
"""Containment-safe extraction directory for an attachment.
|
||||
|
||||
IMAP UIDs are small per-mailbox counters, so `folder_uid` alone put every
|
||||
user's and every account's "INBOX 42" in one directory, where extractions
|
||||
overwrote each other by filename (a path handed to one user's agent could
|
||||
then hold another user's attachment). See src.mail_attachment_paths: the
|
||||
directory is keyed by (owner, account, folder, uid) and contained in
|
||||
ATTACHMENTS_DIR whatever `folder`/`uid` hold."""
|
||||
from src.mail_attachment_paths import attachment_scope_dir
|
||||
|
||||
try:
|
||||
return attachment_scope_dir(ATTACHMENTS_DIR, folder, uid, owner=owner, account_id=account_id)
|
||||
except ValueError:
|
||||
`folder` and `uid` are user-controlled (query/path params). Flatten them to
|
||||
a single safe path segment so a value like folder='../../tmp' can't escape
|
||||
ATTACHMENTS_DIR, then assert containment as belt-and-suspenders."""
|
||||
key = re.sub(r"[^A-Za-z0-9._-]", "_", f"{folder}_{uid}") or "_"
|
||||
target = (ATTACHMENTS_DIR / key).resolve()
|
||||
base = ATTACHMENTS_DIR.resolve()
|
||||
if target != base and base not in target.parents:
|
||||
raise HTTPException(400, "Invalid attachment location")
|
||||
return target
|
||||
|
||||
|
||||
def _init_scheduled_db():
|
||||
@@ -1205,8 +1077,6 @@ def _get_email_config(account_id: str | None = None, owner: str = "") -> dict:
|
||||
cfg = {
|
||||
"account_id": row.id,
|
||||
"account_name": row.name,
|
||||
# Principal for the outbound mail address policy.
|
||||
"owner": row.owner or owner or "",
|
||||
"smtp_host": row.smtp_host or "",
|
||||
"smtp_port": int(row.smtp_port or 465),
|
||||
"smtp_security": _smtp_security_mode({"smtp_security": getattr(row, "smtp_security", ""), "smtp_port": row.smtp_port}),
|
||||
@@ -1240,7 +1110,6 @@ def _get_email_config(account_id: str | None = None, owner: str = "") -> dict:
|
||||
cfg = {
|
||||
"account_id": resolved_id,
|
||||
"account_name": "legacy",
|
||||
"owner": owner or "",
|
||||
"smtp_host": settings.get("smtp_host", os.environ.get("SMTP_HOST", "")),
|
||||
"smtp_port": int(settings.get("smtp_port", os.environ.get("SMTP_PORT", "465")) or 465),
|
||||
"smtp_security": _smtp_security_mode({
|
||||
@@ -1304,16 +1173,11 @@ def _open_imap_connection(
|
||||
starttls: bool,
|
||||
timeout: int = _IMAP_TIMEOUT_SECONDS,
|
||||
ssl_context=None,
|
||||
owner: str | None = None,
|
||||
):
|
||||
"""Open an IMAP connection using the configured security mode.
|
||||
|
||||
`owner` is the principal the destination is judged for
|
||||
(see _mail_private_blocked)."""
|
||||
"""Open an IMAP connection using the configured security mode."""
|
||||
port = int(port or 993)
|
||||
block_private = _mail_private_blocked(owner)
|
||||
if starttls:
|
||||
conn = _PolicyIMAP4(host, port, timeout=timeout, block_private=block_private)
|
||||
conn = imaplib.IMAP4(host, port, timeout=timeout)
|
||||
try:
|
||||
if ssl_context:
|
||||
conn.starttls(ssl_context=ssl_context)
|
||||
@@ -1329,9 +1193,9 @@ def _open_imap_connection(
|
||||
raise
|
||||
elif port == 993:
|
||||
kwargs = {"ssl_context": ssl_context} if ssl_context else {}
|
||||
conn = _PolicyIMAP4_SSL(host, port, timeout=timeout, block_private=block_private, **kwargs)
|
||||
conn = imaplib.IMAP4_SSL(host, port, timeout=timeout, **kwargs)
|
||||
else:
|
||||
conn = _PolicyIMAP4(host, port, timeout=timeout, block_private=block_private)
|
||||
conn = imaplib.IMAP4(host, port, timeout=timeout)
|
||||
try:
|
||||
conn.sock.settimeout(timeout)
|
||||
except Exception:
|
||||
@@ -1371,7 +1235,6 @@ def _imap_connect(account_id: str | None = None, owner: str = "",
|
||||
cfg["imap_port"],
|
||||
starttls=bool(cfg.get("imap_starttls")),
|
||||
timeout=timeout,
|
||||
owner=cfg.get("owner"),
|
||||
)
|
||||
try:
|
||||
if cfg.get("oauth_provider") == "google":
|
||||
|
||||
+242
-207
@@ -36,10 +36,10 @@ from pathlib import Path
|
||||
from email.mime.text import MIMEText
|
||||
from email.mime.multipart import MIMEMultipart
|
||||
|
||||
from starlette.concurrency import run_in_threadpool
|
||||
from fastapi import APIRouter, Query, UploadFile, File, BackgroundTasks, HTTPException, Depends, Request
|
||||
from fastapi.responses import FileResponse, StreamingResponse
|
||||
from src.constants import DATA_DIR
|
||||
from src import shutdown as _shutdown
|
||||
from src.path_confinement import confine
|
||||
|
||||
from src.llm_core import llm_call_async
|
||||
@@ -61,7 +61,6 @@ from .email_helpers import (
|
||||
_fetch_sender_thread_context, _pre_retrieve_context,
|
||||
_EMAIL_REPLY_SYS_PROMPT_BASE, _POOL_HOOKS,
|
||||
_friendly_email_auth_error, _email_summary_failure_log_detail,
|
||||
_mail_connection_test_error, _mail_private_blocked, _PolicySMTP, _PolicySMTP_SSL,
|
||||
_generate_email_summary, EMAIL_SUMMARY_ERROR_CODE, EMAIL_SUMMARY_ERROR_MESSAGE,
|
||||
SendEmailRequest, ExtractStyleRequest,
|
||||
ATTACHMENTS_DIR, COMPOSE_UPLOADS_DIR, SCHEDULED_DB,
|
||||
@@ -1543,44 +1542,6 @@ def _envelope_recipients(*fields: str) -> list:
|
||||
return out
|
||||
|
||||
|
||||
_MD_LINK_URL_RE = re.compile(r"\((https?://)[^)\s]*")
|
||||
|
||||
|
||||
def _md_links_to_html(s: str) -> str:
|
||||
"""Linear ``re.sub(r"\\[([^\\]]+)\\]\\((https?://[^)\\s]+)\\)", '<a href="\\2">\\1</a>', s)``.
|
||||
|
||||
Every `[` before the next `]` shares that `]`, so only the first can match;
|
||||
and when a URL runs into whitespace/end without `)`, every later `[` whose
|
||||
`]` also lies before that stop fails the same way. Skipping those spans
|
||||
keeps a `[[[` or `[a](http://x[a](http://x` flood O(n) instead of the
|
||||
regex's O(n^2).
|
||||
"""
|
||||
out: list[str] = []
|
||||
emitted = scan = 0
|
||||
while True:
|
||||
open_at = s.find("[", scan)
|
||||
if open_at < 0:
|
||||
break
|
||||
close_at = s.find("]", open_at + 1)
|
||||
if close_at < 0:
|
||||
break
|
||||
scan = close_at + 1
|
||||
if close_at == open_at + 1:
|
||||
continue
|
||||
url = _MD_LINK_URL_RE.match(s, close_at + 1)
|
||||
if url is None:
|
||||
continue
|
||||
end = url.end()
|
||||
if end < len(s) and s[end] == ")" and end > url.end(1):
|
||||
out.append(s[emitted:open_at])
|
||||
out.append(f'<a href="{s[url.start(1):end]}">{s[open_at + 1:close_at]}</a>')
|
||||
emitted = scan = end + 1
|
||||
elif end > url.end(1):
|
||||
scan = max(scan, s.rfind("]", scan, end) + 1)
|
||||
out.append(s[emitted:])
|
||||
return "".join(out)
|
||||
|
||||
|
||||
def _md_to_email_html(text: str) -> str:
|
||||
"""Render the compose markdown body to a SAFE HTML fragment for the email's
|
||||
text/html part. Everything is HTML-escaped FIRST (so a pasted <script> /
|
||||
@@ -1596,7 +1557,7 @@ def _md_to_email_html(text: str) -> str:
|
||||
s = re.sub(r"~~([^~]+)~~", r"<del>\1</del>", s)
|
||||
s = re.sub(r"`([^`]+)`", r"<code>\1</code>", s)
|
||||
# links: text + http(s) url only (escape() already neutralised quotes)
|
||||
s = _md_links_to_html(s)
|
||||
s = re.sub(r"\[([^\]]+)\]\((https?://[^)\s]+)\)", r'<a href="\2">\1</a>', s)
|
||||
return s
|
||||
|
||||
parts: list[str] = []
|
||||
@@ -1815,13 +1776,60 @@ def setup_email_routes():
|
||||
return None
|
||||
return v[1]
|
||||
|
||||
# Stale-while-revalidate snapshot store: a last-good list response per key
|
||||
# with NO TTL, so a list view can render instantly from a snapshot (or the
|
||||
# durable index) and refresh in the background instead of blocking ~2-5 s on
|
||||
# IMAP. Kept small and separate from the TTL cache so invalidation of the
|
||||
# fresh cache never removes the fallback snapshot.
|
||||
_LIST_STALE = {}
|
||||
|
||||
def _list_stale_get(key):
|
||||
return _LIST_STALE.get(key)
|
||||
|
||||
def _list_stale_put(key, value):
|
||||
_LIST_STALE[key] = value
|
||||
if len(_LIST_STALE) > 96:
|
||||
for k in list(_LIST_STALE.keys())[:-48]:
|
||||
_LIST_STALE.pop(k, None)
|
||||
|
||||
_LIST_REFRESHING = set()
|
||||
|
||||
def _list_cache_put(key, value):
|
||||
_LIST_CACHE[key] = (_time.monotonic() + _LIST_TTL, value)
|
||||
# Also snapshot the last-good response for stale-while-revalidate.
|
||||
_list_stale_put(key, value)
|
||||
# Cap size
|
||||
if len(_LIST_CACHE) > 64:
|
||||
for k in list(_LIST_CACHE.keys())[:-32]:
|
||||
_LIST_CACHE.pop(k, None)
|
||||
|
||||
def _schedule_list_refresh(sync_args, cache_key, folder, account_id, owner):
|
||||
"""Refresh a list snapshot off the event loop, at most once per key.
|
||||
|
||||
Runs the blocking IMAP fetch in a daemon thread, then repopulates both
|
||||
the TTL cache and the stale snapshot so the next list view is warm.
|
||||
"""
|
||||
if cache_key in _LIST_REFRESHING:
|
||||
return
|
||||
_LIST_REFRESHING.add(cache_key)
|
||||
|
||||
def _run():
|
||||
try:
|
||||
refreshed = _list_emails_sync(*sync_args)
|
||||
if refreshed and not refreshed.get("error"):
|
||||
_list_cache_put(cache_key, refreshed)
|
||||
try:
|
||||
_schedule_recent_email_warm(
|
||||
refreshed.get("emails") or [], folder, account_id, owner)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as exc: # pragma: no cover - background best effort
|
||||
logger.debug("email list background refresh failed: %s", exc)
|
||||
finally:
|
||||
_LIST_REFRESHING.discard(cache_key)
|
||||
|
||||
_threading.Thread(target=_run, daemon=True, name="email-list-refresh").start()
|
||||
|
||||
def _folder_cache_get(account_id, owner):
|
||||
key = (account_id or "", owner or "")
|
||||
v = _FOLDER_CACHE.get(key)
|
||||
@@ -2848,6 +2856,46 @@ def setup_email_routes():
|
||||
owner, account_id or "", folder, filter, limit, offset, elapsed_ms,
|
||||
)
|
||||
return cached
|
||||
# Stale-while-revalidate: never block a plain list view on IMAP while
|
||||
# we already hold a recent snapshot or a durable index. Serve it now
|
||||
# with a `stale` marker and refresh in the background so the next view
|
||||
# is warm. Blocking IMAP only remains for a genuine cold start with no
|
||||
# snapshot and no indexed rows (e.g. a brand-new account/folder).
|
||||
if not from_addr and not has_attachments and not (date_from or date_to):
|
||||
snapshot = _list_stale_get(ck)
|
||||
if snapshot is None:
|
||||
try:
|
||||
indexed_emails, indexed_total, indexed_at = _email_index_list(
|
||||
owner, account_id, folder, filter, limit, offset, False,
|
||||
)
|
||||
except Exception:
|
||||
indexed_emails, indexed_total, indexed_at = [], 0, None
|
||||
if indexed_total:
|
||||
snapshot = {
|
||||
"emails": indexed_emails,
|
||||
"total": indexed_total,
|
||||
"folder": folder,
|
||||
"offset": offset,
|
||||
"sync": {
|
||||
"source": "index",
|
||||
"indexed": indexed_total,
|
||||
"updated_at": indexed_at,
|
||||
},
|
||||
}
|
||||
if snapshot is not None:
|
||||
payload = dict(snapshot)
|
||||
payload["emails"] = list(payload.get("emails") or [])
|
||||
_hide_unlinked_calendar_tags(payload["emails"])
|
||||
sync_meta = dict(payload.get("sync") or {})
|
||||
sync_meta["source"] = "cache_stale" if _list_stale_get(ck) is not None else "index"
|
||||
sync_meta["stale"] = True
|
||||
payload["sync"] = sync_meta
|
||||
_schedule_list_refresh(
|
||||
(folder, limit, offset, filter, account_id, from_addr,
|
||||
bool(has_attachments), owner, True, date_from or "", date_to or ""),
|
||||
ck, folder, account_id, owner,
|
||||
)
|
||||
return payload
|
||||
result = await _asyncio.to_thread(
|
||||
_list_emails_sync, folder, limit, offset, filter, account_id, from_addr,
|
||||
bool(has_attachments), owner, manual_refresh, date_from or "", date_to or "",
|
||||
@@ -3964,7 +4012,7 @@ def setup_email_routes():
|
||||
_row, att = fixture_att
|
||||
filename = str(att.get("filename") or f"attachment-{index}.txt")
|
||||
safe_name = re.sub(r"[^\w\s\-.]", "_", filename).strip() or f"attachment-{index}.txt"
|
||||
target_dir = attachment_extract_dir(folder, uid, owner=owner, account_id=account_id)
|
||||
target_dir = attachment_extract_dir(folder, uid)
|
||||
target_dir.mkdir(parents=True, exist_ok=True)
|
||||
filepath = target_dir / safe_name
|
||||
filepath.write_bytes(str(att.get("content") or "").encode("utf-8"))
|
||||
@@ -3983,7 +4031,7 @@ def setup_email_routes():
|
||||
msg = email_mod.message_from_bytes(raw)
|
||||
|
||||
# Extract to a per-email folder
|
||||
target_dir = attachment_extract_dir(folder, uid, owner=owner, account_id=account_id)
|
||||
target_dir = attachment_extract_dir(folder, uid)
|
||||
filepath = _extract_attachment_to_disk(msg, index, target_dir)
|
||||
if not filepath:
|
||||
return {"error": f"Attachment index {index} not found"}
|
||||
@@ -4015,7 +4063,7 @@ def setup_email_routes():
|
||||
if not attachments:
|
||||
raise HTTPException(status_code=404, detail="No downloadable attachments")
|
||||
|
||||
target_dir = attachment_extract_dir(folder, uid, owner=owner, account_id=account_id)
|
||||
target_dir = attachment_extract_dir(folder, uid)
|
||||
zip_buf = io.BytesIO()
|
||||
used_names: dict[str, int] = {}
|
||||
with zipfile.ZipFile(zip_buf, "w", compression=zipfile.ZIP_DEFLATED) as zf:
|
||||
@@ -4155,7 +4203,7 @@ def setup_email_routes():
|
||||
att for att in _list_attachments_from_msg(msg)
|
||||
if not _is_likely_signature_image_attachment(att)
|
||||
]
|
||||
target_dir = attachment_extract_dir(folder, uid, owner=owner, account_id=acct)
|
||||
target_dir = attachment_extract_dir(folder, uid)
|
||||
for att in attachments:
|
||||
idx = att.get("index")
|
||||
if idx is None:
|
||||
@@ -4220,7 +4268,7 @@ def setup_email_routes():
|
||||
continue
|
||||
if not ct.lower().startswith("image/"):
|
||||
raise HTTPException(status_code=415, detail="Content-ID is not an image")
|
||||
target_dir = attachment_extract_dir(folder, uid, owner=owner, account_id=account_id)
|
||||
target_dir = attachment_extract_dir(folder, uid)
|
||||
filepath = _extract_attachment_to_disk(msg, idx, target_dir)
|
||||
if not filepath:
|
||||
raise HTTPException(status_code=404, detail="Inline image not found")
|
||||
@@ -4258,7 +4306,7 @@ def setup_email_routes():
|
||||
raw = msg_data[0][1]
|
||||
msg = email_mod.message_from_bytes(raw)
|
||||
|
||||
target_dir = attachment_extract_dir(folder, uid, owner=owner, account_id=account_id)
|
||||
target_dir = attachment_extract_dir(folder, uid)
|
||||
filepath = _extract_attachment_to_disk(msg, index, target_dir)
|
||||
if not filepath:
|
||||
return {"error": f"Attachment index {index} not found"}
|
||||
@@ -4586,9 +4634,7 @@ def setup_email_routes():
|
||||
try:
|
||||
content = filepath.read_text(encoding="utf-8", errors="replace")
|
||||
except Exception as e:
|
||||
# OSError text includes the server-side extraction path.
|
||||
logger.warning("Failed to read text attachment %s: %s", base, e)
|
||||
return {"error": "Failed to read text file", "filename": base}
|
||||
return {"error": f"Failed to read text file: {e}", "filename": base}
|
||||
doc_id = _create_markdown_doc(content, "Imported from email attachment")
|
||||
return {"doc_id": doc_id, "filename": filepath.name}
|
||||
|
||||
@@ -4605,7 +4651,7 @@ def setup_email_routes():
|
||||
_row, att = fixture_att
|
||||
filename = str(att.get("filename") or f"attachment-{index}.txt")
|
||||
safe_name = re.sub(r"[^\w\s\-.]", "_", filename).strip() or f"attachment-{index}.txt"
|
||||
target_dir = attachment_extract_dir(folder, uid, owner=owner, account_id=account_id)
|
||||
target_dir = attachment_extract_dir(folder, uid)
|
||||
target_dir.mkdir(parents=True, exist_ok=True)
|
||||
filepath = target_dir / safe_name
|
||||
filepath.write_bytes(str(att.get("content") or "").encode("utf-8"))
|
||||
@@ -4619,7 +4665,7 @@ def setup_email_routes():
|
||||
raw = msg_data[0][1]
|
||||
msg = email_mod.message_from_bytes(raw)
|
||||
|
||||
target_dir = attachment_extract_dir(folder, uid, owner=owner, account_id=account_id)
|
||||
target_dir = attachment_extract_dir(folder, uid)
|
||||
filepath = _extract_attachment_to_disk(msg, index, target_dir)
|
||||
if not filepath:
|
||||
return {"error": f"Attachment index {index} not found"}
|
||||
@@ -5273,7 +5319,7 @@ def setup_email_routes():
|
||||
return {"success": False, "error": "Email not found"}
|
||||
raw = msg_data[0][1]
|
||||
msg = email_mod.message_from_bytes(raw)
|
||||
target_dir = attachment_extract_dir(folder, uid, owner=owner, account_id=account_id)
|
||||
target_dir = attachment_extract_dir(folder, uid)
|
||||
filepath = _extract_attachment_to_disk(msg, index, target_dir)
|
||||
if not filepath:
|
||||
return {"success": False, "error": f"Attachment index {index} not found"}
|
||||
@@ -6290,7 +6336,7 @@ def setup_email_routes():
|
||||
await queue.put({'type': 'result', **result})
|
||||
task = asyncio.create_task(run())
|
||||
try:
|
||||
while True:
|
||||
while not _shutdown.is_shutting_down():
|
||||
item = await queue.get()
|
||||
yield 'data: ' + json.dumps(item) + '\n\n'
|
||||
if item.get('type') == 'result':
|
||||
@@ -6980,158 +7026,6 @@ def setup_email_routes():
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
def _test_mail_connections(body: dict, owner: str) -> dict:
|
||||
"""Blocking IMAP/SMTP connection test for /accounts/test.
|
||||
|
||||
Runs in the threadpool: each probe can block for its full socket
|
||||
timeout, which on the event loop stalled every other request. The
|
||||
connections go through the outbound mail address policy for `owner`
|
||||
(the user who chose the hosts)."""
|
||||
imap_result = {"ok": False}
|
||||
smtp_result = None
|
||||
|
||||
imap_host = (body.get("imap_host") or "").strip()
|
||||
imap_port, imap_port_err = _coerce_port(body.get("imap_port"), 993)
|
||||
imap_user = (body.get("imap_user") or "").strip()
|
||||
imap_pass = body.get("imap_password") or ""
|
||||
imap_starttls = bool(body.get("imap_starttls"))
|
||||
oauth_provider = body.get("oauth_provider") or ""
|
||||
|
||||
google_token = None
|
||||
google_token_loaded = False
|
||||
google_ssl_context = (
|
||||
ssl.create_default_context()
|
||||
if oauth_provider == "google"
|
||||
else None
|
||||
)
|
||||
|
||||
def _google_token():
|
||||
nonlocal google_token, google_token_loaded
|
||||
if not google_token_loaded:
|
||||
google_token = _get_valid_google_token(body.get("account_id"), body)
|
||||
google_token_loaded = True
|
||||
if not google_token:
|
||||
raise RuntimeError("Google OAuth token unavailable — reconnect the account")
|
||||
return google_token
|
||||
|
||||
if imap_port_err:
|
||||
imap_result = {"ok": False, "error": imap_port_err}
|
||||
elif not (imap_host and imap_user and (imap_pass or oauth_provider == "google")):
|
||||
imap_result = {"ok": False, "error": "Need IMAP host, username, and password"}
|
||||
elif oauth_provider == "google" and _normalized_mail_host(imap_host) != _GOOGLE_OAUTH_IMAP_HOST:
|
||||
imap_result = {"ok": False, "error": "Google OAuth IMAP requires imap.gmail.com"}
|
||||
elif oauth_provider == "google" and not _google_oauth_imap_transport_allowed(imap_port, imap_starttls):
|
||||
imap_result = {"ok": False, "error": "Google OAuth IMAP requires TLS on port 993 or STARTTLS on port 143"}
|
||||
else:
|
||||
# Connection mode resolution:
|
||||
# STARTTLS on → plain IMAP4 + .starttls() (upgrade)
|
||||
# STARTTLS off + port 993 → IMAP4_SSL (implicit SSL, "IMAPS")
|
||||
# STARTTLS off + any other port → plain IMAP4 (no encryption)
|
||||
# Without the last branch, local servers exposed on a non-993
|
||||
# port (Dovecot on 31143, etc.) would always fail the SSL
|
||||
# handshake because they're not actually wrapped in TLS.
|
||||
try:
|
||||
imap_kwargs = {
|
||||
"starttls": imap_starttls,
|
||||
"timeout": _IMAP_TIMEOUT_SECONDS,
|
||||
"owner": owner,
|
||||
}
|
||||
if google_ssl_context:
|
||||
imap_kwargs["ssl_context"] = google_ssl_context
|
||||
conn = _open_imap_connection(
|
||||
imap_host,
|
||||
imap_port,
|
||||
**imap_kwargs,
|
||||
)
|
||||
try:
|
||||
if oauth_provider == "google":
|
||||
token = _google_token()
|
||||
conn.authenticate("XOAUTH2", lambda x: _xoauth2_bytes(imap_user, token))
|
||||
else:
|
||||
conn.login(imap_user, imap_pass)
|
||||
imap_result = {"ok": True}
|
||||
finally:
|
||||
try: conn.logout()
|
||||
except Exception: pass
|
||||
except Exception as e:
|
||||
imap_result = {"ok": False, "error": _mail_connection_test_error("IMAP", imap_host, e)}
|
||||
|
||||
smtp_host = (body.get("smtp_host") or "").strip()
|
||||
smtp_port, smtp_port_err = _coerce_port(body.get("smtp_port"), 465)
|
||||
if smtp_host and smtp_port_err:
|
||||
smtp_result = {"ok": False, "error": smtp_port_err}
|
||||
elif oauth_provider == "google" and smtp_host and _normalized_mail_host(smtp_host) != _GOOGLE_OAUTH_SMTP_HOST:
|
||||
smtp_result = {"ok": False, "error": "Google OAuth SMTP requires smtp.gmail.com"}
|
||||
elif (
|
||||
oauth_provider == "google"
|
||||
and smtp_host
|
||||
and not _google_oauth_smtp_transport_allowed(
|
||||
smtp_port,
|
||||
_smtp_security_mode({"smtp_security": body.get("smtp_security"), "smtp_port": smtp_port}),
|
||||
)
|
||||
):
|
||||
smtp_result = {"ok": False, "error": "Google OAuth SMTP requires TLS on port 465 or STARTTLS on port 587"}
|
||||
elif smtp_host:
|
||||
smtp_security = _smtp_security_mode({"smtp_security": body.get("smtp_security"), "smtp_port": smtp_port})
|
||||
smtp_user = (body.get("smtp_user") or imap_user).strip()
|
||||
smtp_pass = body.get("smtp_password") or imap_pass
|
||||
smtp = None
|
||||
try:
|
||||
if smtp_security == "ssl":
|
||||
smtp_kwargs = (
|
||||
{"context": google_ssl_context}
|
||||
if google_ssl_context
|
||||
else {}
|
||||
)
|
||||
smtp = _PolicySMTP_SSL(
|
||||
smtp_host,
|
||||
smtp_port,
|
||||
timeout=10,
|
||||
block_private=_mail_private_blocked(owner),
|
||||
**smtp_kwargs,
|
||||
)
|
||||
else:
|
||||
smtp = _PolicySMTP(smtp_host, smtp_port, timeout=10, block_private=_mail_private_blocked(owner))
|
||||
if smtp_security == "starttls":
|
||||
try:
|
||||
if google_ssl_context:
|
||||
smtp.starttls(context=google_ssl_context)
|
||||
else:
|
||||
smtp.starttls()
|
||||
except Exception:
|
||||
# STARTTLS failed before the auth cleanup block.
|
||||
# Close the still-open plaintext socket explicitly.
|
||||
try:
|
||||
smtp.close()
|
||||
except Exception:
|
||||
pass
|
||||
smtp = None
|
||||
raise
|
||||
if oauth_provider == "google":
|
||||
token = _google_token()
|
||||
smtp.ehlo()
|
||||
smtp.auth("XOAUTH2", lambda challenge=None: _xoauth2_raw(smtp_user, token), initial_response_ok=True)
|
||||
else:
|
||||
smtp.login(smtp_user, smtp_pass)
|
||||
smtp_result = {"ok": True}
|
||||
except Exception as e:
|
||||
smtp_result = {"ok": False, "error": _mail_connection_test_error("SMTP", smtp_host, e)}
|
||||
finally:
|
||||
if smtp is not None:
|
||||
try:
|
||||
smtp.quit()
|
||||
except Exception:
|
||||
try:
|
||||
smtp.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return {
|
||||
"ok": imap_result["ok"] and (smtp_result is None or smtp_result["ok"]),
|
||||
"imap": imap_result,
|
||||
"smtp": smtp_result,
|
||||
}
|
||||
|
||||
@router.post("/accounts/test")
|
||||
async def test_account_config(req: Request, owner: str = Depends(require_user)):
|
||||
"""Try to actually connect to the provided IMAP (and optionally SMTP)
|
||||
@@ -7195,7 +7089,148 @@ def setup_email_routes():
|
||||
# inline test payloads select the OAuth branch or supply token data.
|
||||
body = {key: value for key, value in body.items() if key not in _SERVER_OWNED_OAUTH_FIELDS}
|
||||
|
||||
return await run_in_threadpool(_test_mail_connections, body, owner)
|
||||
imap_result = {"ok": False}
|
||||
smtp_result = None
|
||||
|
||||
imap_host = (body.get("imap_host") or "").strip()
|
||||
imap_port, imap_port_err = _coerce_port(body.get("imap_port"), 993)
|
||||
imap_user = (body.get("imap_user") or "").strip()
|
||||
imap_pass = body.get("imap_password") or ""
|
||||
imap_starttls = bool(body.get("imap_starttls"))
|
||||
oauth_provider = body.get("oauth_provider") or ""
|
||||
|
||||
google_token = None
|
||||
google_token_loaded = False
|
||||
google_ssl_context = (
|
||||
ssl.create_default_context()
|
||||
if oauth_provider == "google"
|
||||
else None
|
||||
)
|
||||
|
||||
def _google_token():
|
||||
nonlocal google_token, google_token_loaded
|
||||
if not google_token_loaded:
|
||||
google_token = _get_valid_google_token(body.get("account_id"), body)
|
||||
google_token_loaded = True
|
||||
if not google_token:
|
||||
raise RuntimeError("Google OAuth token unavailable — reconnect the account")
|
||||
return google_token
|
||||
|
||||
if imap_port_err:
|
||||
imap_result = {"ok": False, "error": imap_port_err}
|
||||
elif not (imap_host and imap_user and (imap_pass or oauth_provider == "google")):
|
||||
imap_result = {"ok": False, "error": "Need IMAP host, username, and password"}
|
||||
elif oauth_provider == "google" and _normalized_mail_host(imap_host) != _GOOGLE_OAUTH_IMAP_HOST:
|
||||
imap_result = {"ok": False, "error": "Google OAuth IMAP requires imap.gmail.com"}
|
||||
elif oauth_provider == "google" and not _google_oauth_imap_transport_allowed(imap_port, imap_starttls):
|
||||
imap_result = {"ok": False, "error": "Google OAuth IMAP requires TLS on port 993 or STARTTLS on port 143"}
|
||||
else:
|
||||
# Connection mode resolution:
|
||||
# STARTTLS on → plain IMAP4 + .starttls() (upgrade)
|
||||
# STARTTLS off + port 993 → IMAP4_SSL (implicit SSL, "IMAPS")
|
||||
# STARTTLS off + any other port → plain IMAP4 (no encryption)
|
||||
# Without the last branch, local servers exposed on a non-993
|
||||
# port (Dovecot on 31143, etc.) would always fail the SSL
|
||||
# handshake because they're not actually wrapped in TLS.
|
||||
try:
|
||||
imap_kwargs = {
|
||||
"starttls": imap_starttls,
|
||||
"timeout": _IMAP_TIMEOUT_SECONDS,
|
||||
}
|
||||
if google_ssl_context:
|
||||
imap_kwargs["ssl_context"] = google_ssl_context
|
||||
conn = _open_imap_connection(
|
||||
imap_host,
|
||||
imap_port,
|
||||
**imap_kwargs,
|
||||
)
|
||||
try:
|
||||
if oauth_provider == "google":
|
||||
token = _google_token()
|
||||
conn.authenticate("XOAUTH2", lambda x: _xoauth2_bytes(imap_user, token))
|
||||
else:
|
||||
conn.login(imap_user, imap_pass)
|
||||
imap_result = {"ok": True}
|
||||
finally:
|
||||
try: conn.logout()
|
||||
except Exception: pass
|
||||
except Exception as e:
|
||||
imap_result = {"ok": False, "error": _friendly_email_auth_error("IMAP", imap_host, e)}
|
||||
|
||||
smtp_host = (body.get("smtp_host") or "").strip()
|
||||
smtp_port, smtp_port_err = _coerce_port(body.get("smtp_port"), 465)
|
||||
if smtp_host and smtp_port_err:
|
||||
smtp_result = {"ok": False, "error": smtp_port_err}
|
||||
elif oauth_provider == "google" and smtp_host and _normalized_mail_host(smtp_host) != _GOOGLE_OAUTH_SMTP_HOST:
|
||||
smtp_result = {"ok": False, "error": "Google OAuth SMTP requires smtp.gmail.com"}
|
||||
elif (
|
||||
oauth_provider == "google"
|
||||
and smtp_host
|
||||
and not _google_oauth_smtp_transport_allowed(
|
||||
smtp_port,
|
||||
_smtp_security_mode({"smtp_security": body.get("smtp_security"), "smtp_port": smtp_port}),
|
||||
)
|
||||
):
|
||||
smtp_result = {"ok": False, "error": "Google OAuth SMTP requires TLS on port 465 or STARTTLS on port 587"}
|
||||
elif smtp_host:
|
||||
smtp_security = _smtp_security_mode({"smtp_security": body.get("smtp_security"), "smtp_port": smtp_port})
|
||||
smtp_user = (body.get("smtp_user") or imap_user).strip()
|
||||
smtp_pass = body.get("smtp_password") or imap_pass
|
||||
smtp = None
|
||||
try:
|
||||
if smtp_security == "ssl":
|
||||
smtp_kwargs = (
|
||||
{"context": google_ssl_context}
|
||||
if google_ssl_context
|
||||
else {}
|
||||
)
|
||||
smtp = smtplib.SMTP_SSL(
|
||||
smtp_host,
|
||||
smtp_port,
|
||||
timeout=10,
|
||||
**smtp_kwargs,
|
||||
)
|
||||
else:
|
||||
smtp = smtplib.SMTP(smtp_host, smtp_port, timeout=10)
|
||||
if smtp_security == "starttls":
|
||||
try:
|
||||
if google_ssl_context:
|
||||
smtp.starttls(context=google_ssl_context)
|
||||
else:
|
||||
smtp.starttls()
|
||||
except Exception:
|
||||
# STARTTLS failed before the auth cleanup block.
|
||||
# Close the still-open plaintext socket explicitly.
|
||||
try:
|
||||
smtp.close()
|
||||
except Exception:
|
||||
pass
|
||||
smtp = None
|
||||
raise
|
||||
if oauth_provider == "google":
|
||||
token = _google_token()
|
||||
smtp.ehlo()
|
||||
smtp.auth("XOAUTH2", lambda challenge=None: _xoauth2_raw(smtp_user, token), initial_response_ok=True)
|
||||
else:
|
||||
smtp.login(smtp_user, smtp_pass)
|
||||
smtp_result = {"ok": True}
|
||||
except Exception as e:
|
||||
smtp_result = {"ok": False, "error": _friendly_email_auth_error("SMTP", smtp_host, e)}
|
||||
finally:
|
||||
if smtp is not None:
|
||||
try:
|
||||
smtp.quit()
|
||||
except Exception:
|
||||
try:
|
||||
smtp.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return {
|
||||
"ok": imap_result["ok"] and (smtp_result is None or smtp_result["ok"]),
|
||||
"imap": imap_result,
|
||||
"smtp": smtp_result,
|
||||
}
|
||||
|
||||
@router.post("/accounts/{account_id}/set-default")
|
||||
async def set_default_account(account_id: str, owner: str = Depends(require_user)):
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+4
-10
@@ -5,9 +5,8 @@ import shlex
|
||||
import subprocess
|
||||
from copy import deepcopy
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from fastapi import APIRouter, HTTPException
|
||||
|
||||
from core.middleware import require_admin
|
||||
from core.platform_compat import run_ssh_command
|
||||
from routes._validators import validate_remote_host, validate_ssh_port
|
||||
|
||||
@@ -199,11 +198,7 @@ def _inspect_model_path(model_path: str, host: str = "", ssh_port: str = "") ->
|
||||
|
||||
|
||||
def setup_hwfit_routes():
|
||||
# Admin-gated like the Cookbook routes this backs (/api/cookbook/gpus does
|
||||
# the same SSH host probing): `host`/`ssh_port` make the server SSH out with
|
||||
# its own key to any syntactically valid host, and `model_path` probes the
|
||||
# host filesystem. Single-user mode and internal tool calls still pass.
|
||||
router = APIRouter(prefix="/api/hwfit", tags=["hwfit"], dependencies=[Depends(require_admin)])
|
||||
router = APIRouter(prefix="/api/hwfit", tags=["hwfit"])
|
||||
|
||||
@router.get("/system")
|
||||
def get_system(host: str = "", ssh_port: str = "", platform: str = "", fresh: bool = False):
|
||||
@@ -223,7 +218,7 @@ def setup_hwfit_routes():
|
||||
fresh=true bypasses the hardware-detection cache."""
|
||||
from services.hwfit.hardware import detect_system
|
||||
from services.hwfit.fit import rank_models
|
||||
from services.hwfit.models import get_models, refresh_dynamic_catalogs
|
||||
from services.hwfit.models import get_models, model_catalog_path, refresh_dynamic_catalogs
|
||||
host, ssh_port = _validate_detection_target(host, ssh_port)
|
||||
system = deepcopy(detect_system(host=host, ssh_port=ssh_port, platform=platform, fresh=fresh))
|
||||
if system.get("error"):
|
||||
@@ -238,8 +233,7 @@ def setup_hwfit_routes():
|
||||
return {
|
||||
"system": system,
|
||||
"models": [],
|
||||
"error": "Model catalog is empty. Use Rescan in Cookbook while online to populate runtime model data; offline recommendations require a previously populated cache.",
|
||||
**({"catalog_refresh": catalog_refresh} if catalog_refresh is not None else {}),
|
||||
"error": f"Model catalog missing or empty: {model_catalog_path()}",
|
||||
}
|
||||
|
||||
if ignore_detected_gpu:
|
||||
|
||||
+22
-7
@@ -472,7 +472,7 @@ def _truthy(value: str | None) -> bool:
|
||||
|
||||
_ENDPOINT_KINDS = {"auto", "local", "api", "proxy"}
|
||||
_REFRESH_MODES = {"auto", "manual", "disabled"}
|
||||
_MODEL_TOOL_MODES = {"none", "compact", "full"}
|
||||
_MODEL_TOOL_MODES = {"none", "compact", "full", "spec"}
|
||||
_MODEL_TOOL_MODE_ALIASES = {
|
||||
"regular": "full",
|
||||
"odysseus_compact": "compact",
|
||||
@@ -552,10 +552,10 @@ def _endpoint_refresh_timeout(ep: Any, category: str) -> float:
|
||||
val = 0
|
||||
if val > 0:
|
||||
return float(max(1, min(60, val)))
|
||||
# llama.cpp and other local OpenAI-compatible servers can block briefly
|
||||
# while warming/loading. A 2s local timeout makes working endpoints flicker
|
||||
# offline before /v1/models is ready.
|
||||
return 10.0 if category == "local" else 2.0
|
||||
# Keep the background refresh short so a single dead endpoint cannot stall
|
||||
# the model picker. A timed-out refresh never clears a non-empty model
|
||||
# cache, so a briefly-warming local server is simply retried next round.
|
||||
return 2.0
|
||||
|
||||
|
||||
def _manual_refresh_timeout(ep: Any, category: str, requested: Any = None) -> float:
|
||||
@@ -1510,6 +1510,9 @@ def setup_model_routes(model_discovery):
|
||||
_refresh_inflight = {"v": False, "done": None} # coarse single-flight guard
|
||||
_REFRESH_FAILURE_BASE = 300.0
|
||||
_REFRESH_FAILURE_MAX = 3600.0
|
||||
# Even a user-triggered refresh must not re-probe an endpoint that just
|
||||
# failed/timed out on every click; keep a short cooldown.
|
||||
_REFRESH_FAILURE_MIN = 60.0
|
||||
|
||||
def _refresh_key(base: str, api_key: Optional[str], provider_auth_id: Optional[str] = None) -> str:
|
||||
# Session-backed endpoints carry no static key; include their auth id
|
||||
@@ -1557,14 +1560,20 @@ def setup_model_routes(model_discovery):
|
||||
if mode in ("manual", "disabled") and not force:
|
||||
return False, info
|
||||
fails = int(state.get("fail_count") or 0)
|
||||
if fails and not force:
|
||||
if fails:
|
||||
last_failure = float(state.get("last_failure") or 0.0)
|
||||
empty_local = (
|
||||
not cached
|
||||
and category == "local"
|
||||
and str(getattr(ep, "id", "") or "").startswith("local-")
|
||||
)
|
||||
if now - last_failure < _failure_delay(fails, empty_local=empty_local):
|
||||
delay = _failure_delay(fails, empty_local=empty_local)
|
||||
if force:
|
||||
# A user-triggered refresh gets a short cooldown instead of the
|
||||
# long background backoff, but still skips a just-failed
|
||||
# endpoint so repeated clicks cannot stall the picker.
|
||||
delay = min(delay, _REFRESH_FAILURE_MIN)
|
||||
if now - last_failure < delay:
|
||||
return False, info
|
||||
if cached and not force:
|
||||
interval = _endpoint_refresh_interval(ep, category)
|
||||
@@ -3055,4 +3064,10 @@ def setup_model_routes(model_discovery):
|
||||
|
||||
router._should_refresh_endpoint = _should_refresh_endpoint
|
||||
router._search_endpoint_catalog = search_endpoint_catalog
|
||||
# Exposed for tests: the in-memory backoff state, its key function, and the
|
||||
# background refresh driver (probes endpoints in parallel with the short
|
||||
# per-endpoint timeout).
|
||||
router._refresh_state = _refresh_state
|
||||
router._refresh_key = _refresh_key
|
||||
router._refresh_caches_bg = _refresh_caches_bg
|
||||
return router
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
"""In-flight activity counters (heavy image ops) for safe restarts/observability.
|
||||
|
||||
A deploy restart kills any in-flight long image op (a diffusion edit is 60-120 s),
|
||||
so the operator needs to know whether the process is busy before issuing
|
||||
``systemctl --user restart``. These counters are process-local and cheap.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from typing import Dict
|
||||
|
||||
_lock = threading.Lock()
|
||||
_image_ops = 0
|
||||
_image_ops_by_name: Dict[str, int] = {}
|
||||
_started_at = time.time()
|
||||
|
||||
|
||||
def begin_image_op(name: str = "image") -> None:
|
||||
global _image_ops
|
||||
with _lock:
|
||||
_image_ops += 1
|
||||
_image_ops_by_name[name] = _image_ops_by_name.get(name, 0) + 1
|
||||
|
||||
|
||||
def end_image_op(name: str = "image") -> None:
|
||||
global _image_ops
|
||||
with _lock:
|
||||
_image_ops = max(0, _image_ops - 1)
|
||||
if name in _image_ops_by_name:
|
||||
_image_ops_by_name[name] -= 1
|
||||
if _image_ops_by_name[name] <= 0:
|
||||
_image_ops_by_name.pop(name, None)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def image_op(name: str = "image"):
|
||||
begin_image_op(name)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
end_image_op(name)
|
||||
|
||||
|
||||
def activity() -> Dict[str, object]:
|
||||
with _lock:
|
||||
return {
|
||||
"image_ops": _image_ops,
|
||||
"image_ops_by_name": dict(_image_ops_by_name),
|
||||
"uptime_s": round(time.time() - _started_at, 1),
|
||||
}
|
||||
|
||||
|
||||
def heavy_image_ops_active() -> int:
|
||||
with _lock:
|
||||
return _image_ops
|
||||
|
||||
|
||||
def busy() -> bool:
|
||||
return heavy_image_ops_active() > 0
|
||||
@@ -0,0 +1,646 @@
|
||||
"""In-process agent tools for the scoped ``photo_editor`` capability family.
|
||||
|
||||
These handlers run as the logged-in owner and drive the image editor over the
|
||||
owner-scoped in-process bridge (``/api/editor/bridge/*``) plus the app's
|
||||
existing image transforms (``/api/image/*``). They are only offered to the
|
||||
model when the ``photo_editor`` family is active (the Quick Edit popup or an
|
||||
explicit photo-edit chat turn), never in the general tool list.
|
||||
|
||||
Loopback calls carry the internal token and impersonate ``owner`` (middleware
|
||||
app.py), so the editor bridge resolves the same ``owner:session`` scope the
|
||||
browser subscriber uses.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
from src.tools._common import _INTERNAL_BASE, _internal_headers
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _provider_message() -> str:
|
||||
try:
|
||||
from src.ai_interaction import IMAGE_PROVIDER_REQUIRED_MESSAGE
|
||||
|
||||
return IMAGE_PROVIDER_REQUIRED_MESSAGE
|
||||
except Exception:
|
||||
return (
|
||||
"Removing/adding objects needs an image generation provider; enable one "
|
||||
"in Settings > Models."
|
||||
)
|
||||
|
||||
|
||||
def _provider_available(owner: Optional[str]) -> bool:
|
||||
try:
|
||||
from src.ai_interaction import image_editing_provider_available
|
||||
|
||||
return bool(image_editing_provider_available(owner))
|
||||
except Exception:
|
||||
return True
|
||||
|
||||
# Local diffusion/image servers can take 90-110s per image (8 steps) and
|
||||
# longer on a cold start; the previous 120s default timed out mid-generation.
|
||||
_EDITOR_TIMEOUT = 300.0
|
||||
# A mask covering less than this fraction of the image cannot drive a useful
|
||||
# regional inpaint (tiny figures, empty SAM output): fall back to grounding or
|
||||
# whole-image instruction mode instead.
|
||||
_MIN_MASK_FRACTION = 0.002
|
||||
_AI_OPS = {
|
||||
"remove-bg": "/api/image/remove-bg",
|
||||
"upscale": "/api/image/upscale-local",
|
||||
"denoise": "/api/image/denoise",
|
||||
"enhance-face": "/api/image/enhance-face",
|
||||
"inpaint": "/api/image/inpaint",
|
||||
"generative": "/api/image/inpaint",
|
||||
}
|
||||
_PLACE_OP = "layers.placeImage"
|
||||
|
||||
|
||||
def _parse(content: Any) -> Dict[str, Any]:
|
||||
raw = content if isinstance(content, str) else json.dumps(content or {})
|
||||
try:
|
||||
parsed = json.loads(raw) if raw.strip() else {}
|
||||
except (ValueError, TypeError):
|
||||
parsed = {}
|
||||
return parsed if isinstance(parsed, dict) else {}
|
||||
|
||||
|
||||
def _session(args: Dict[str, Any], ctx: Dict[str, Any]) -> str:
|
||||
return str(args.get("session") or args.get("editor_session") or ctx.get("session_id") or "default")
|
||||
|
||||
|
||||
async def _post(path: str, payload: Dict[str, Any], owner: Optional[str], timeout: float = _EDITOR_TIMEOUT) -> Dict[str, Any]:
|
||||
headers = {**_internal_headers(owner), "Content-Type": "application/json"}
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
resp = await client.post(f"{_INTERNAL_BASE}{path}", json=payload, headers=headers)
|
||||
resp.raise_for_status()
|
||||
try:
|
||||
return resp.json()
|
||||
except ValueError:
|
||||
return {"raw": resp.text}
|
||||
|
||||
|
||||
def _as_data_url(value: Any) -> Optional[str]:
|
||||
if not isinstance(value, str) or not value:
|
||||
return None
|
||||
if value.startswith("data:"):
|
||||
return value
|
||||
if value.startswith("/"):
|
||||
return None
|
||||
return f"data:image/png;base64,{value}"
|
||||
|
||||
|
||||
def _data_url_bytes(data_url: str) -> bytes:
|
||||
_, _, b64 = data_url.partition(",")
|
||||
return base64.b64decode(b64)
|
||||
|
||||
|
||||
def _persist_png(image_bytes: bytes, owner: Optional[str], session_id: Optional[str], prompt: str, model: str) -> Dict[str, Any]:
|
||||
"""Save an edited/generated PNG and return {image_id, image_url, image_size}."""
|
||||
from core.database import GalleryImage, SessionLocal
|
||||
from src.constants import GENERATED_IMAGES_DIR
|
||||
|
||||
db = SessionLocal()
|
||||
try:
|
||||
root = Path(GENERATED_IMAGES_DIR).resolve()
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
filename = f"{uuid.uuid4().hex[:12]}.png"
|
||||
(root / filename).write_bytes(image_bytes)
|
||||
new_id = str(uuid.uuid4())
|
||||
width = height = 0
|
||||
try:
|
||||
from PIL import Image
|
||||
|
||||
with Image.open(io.BytesIO(image_bytes)) as img:
|
||||
width, height = img.size
|
||||
except Exception:
|
||||
pass
|
||||
size = f"{width}x{height}" if width and height else ""
|
||||
# The editor bridge uses an owner-scoped session id (e.g. "pewds:3")
|
||||
# that is not a chat-sessions row. GalleryImage.session_id is a FK to
|
||||
# sessions.id, so drop an unknown id instead of raising IntegrityError.
|
||||
if session_id:
|
||||
from core.database import Session as _ChatSession
|
||||
|
||||
try:
|
||||
known = db.query(_ChatSession.id).filter(_ChatSession.id == session_id).first()
|
||||
except Exception:
|
||||
known = None
|
||||
if known is None:
|
||||
session_id = None
|
||||
derived = GalleryImage(
|
||||
id=new_id,
|
||||
filename=filename,
|
||||
prompt=prompt or None,
|
||||
caption=None,
|
||||
model=model,
|
||||
size=size or None,
|
||||
quality=None,
|
||||
tags=None,
|
||||
ai_tags=None,
|
||||
session_id=session_id,
|
||||
album_id=None,
|
||||
owner=owner,
|
||||
file_hash=hashlib.sha256(image_bytes).hexdigest(),
|
||||
file_size=len(image_bytes),
|
||||
width=width or None,
|
||||
height=height or None,
|
||||
)
|
||||
db.add(derived)
|
||||
db.commit()
|
||||
return {"image_id": new_id, "image_url": f"/api/generated-image/{filename}", "image_size": size or None}
|
||||
except Exception:
|
||||
db.rollback()
|
||||
raise
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
async def _place_layer(data_url: str, name: str, owner: Optional[str], session: str, *, masked: bool = False, blend: str = "normal", opacity: float = 100.0, mode: Optional[str] = None) -> bool:
|
||||
params: Dict[str, Any] = {
|
||||
"dataUrl": data_url,
|
||||
"name": name,
|
||||
"masked": bool(masked),
|
||||
"blend": blend or "normal",
|
||||
"opacity": opacity,
|
||||
}
|
||||
# Forward the edit mode so the client can build a difference mask for
|
||||
# instruction-mode results: the whole image was regenerated under an
|
||||
# instruction prompt, so untouched pixels must stay pixel-identical (the
|
||||
# client keeps only the changed region) instead of the new layer overwriting
|
||||
# the entire canvas. Omit the keys entirely when no mode is given so the
|
||||
# client keeps its current behavior (e.g. generate_image must not diff-mask).
|
||||
if mode:
|
||||
params["mode"] = mode
|
||||
params["differenceMask"] = mode == "instruction"
|
||||
if mode == "instruction":
|
||||
# Exclude a thin edge band from the client difference mask: full-frame
|
||||
# regenerations can shift pixels along the border.
|
||||
params["edgeExcludePx"] = 3
|
||||
payload = {
|
||||
"session": session,
|
||||
"op": _PLACE_OP,
|
||||
"params": params,
|
||||
"wait": True,
|
||||
"timeout": 60,
|
||||
}
|
||||
try:
|
||||
data = await _post("/api/editor/bridge/command", payload, owner, timeout=90)
|
||||
return data.get("ok") is not False
|
||||
except Exception as exc: # pragma: no cover - defensive
|
||||
logger.warning("editor layer placement failed: %s", exc)
|
||||
return False
|
||||
|
||||
|
||||
async def _fetch_selection_mask(owner: Optional[str], session: str) -> Optional[str]:
|
||||
"""Return the editor's current selection as a data-URL mask, or None.
|
||||
|
||||
Generative ops (inpaint/generative) need a mask; the planner/judge may not
|
||||
have supplied one, so fall back to whatever is selected in the editor.
|
||||
"""
|
||||
payload = {
|
||||
"session": session,
|
||||
"op": "selection.exportMask",
|
||||
"params": {},
|
||||
"wait": True,
|
||||
"timeout": 30,
|
||||
}
|
||||
try:
|
||||
data = await _post("/api/editor/bridge/command", payload, owner, timeout=45)
|
||||
except Exception as exc: # pragma: no cover - defensive
|
||||
logger.warning("editor selection mask fetch failed: %s", exc)
|
||||
return None
|
||||
result = data.get("result") if isinstance(data, dict) else None
|
||||
if isinstance(result, dict):
|
||||
return _as_data_url(result.get("mask"))
|
||||
return None
|
||||
|
||||
|
||||
def _mask_white_fraction(data_url: Optional[str]) -> float:
|
||||
"""Fraction of a grayscale mask that is selected (white, >=128)."""
|
||||
if not data_url:
|
||||
return 0.0
|
||||
try:
|
||||
from PIL import Image as _PILImage
|
||||
|
||||
with _PILImage.open(io.BytesIO(_data_url_bytes(data_url))).convert("L") as im:
|
||||
total = im.width * im.height
|
||||
if not total:
|
||||
return 0.0
|
||||
hist = im.histogram()
|
||||
return sum(hist[128:]) / total
|
||||
except Exception:
|
||||
return 0.0
|
||||
|
||||
|
||||
def _image_is_blank(data_url: str) -> bool:
|
||||
"""True when a result has an alpha channel that is essentially transparent."""
|
||||
try:
|
||||
from PIL import Image as _PILImage
|
||||
|
||||
with _PILImage.open(io.BytesIO(_data_url_bytes(data_url))) as im:
|
||||
if im.mode not in ("RGBA", "LA") and "transparency" not in im.info:
|
||||
return False
|
||||
alpha = im.convert("RGBA").getchannel("A")
|
||||
hist = alpha.histogram()
|
||||
total = sum(hist)
|
||||
return total > 0 and (sum(hist[1:]) / total) < 0.001
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _rect_mask_data_url(image_data_url: str, box: List[float], *, pad_frac: float = 0.15) -> Optional[str]:
|
||||
"""Build a rectangular mask data URL from a [x1,y1,x2,y2] box, dilated.
|
||||
|
||||
Used when the grounding detector finds an object but SAM returns an empty
|
||||
mask (tiny figures): a dilated box guarantees a usable region.
|
||||
"""
|
||||
try:
|
||||
from PIL import Image as _PILImage, ImageDraw
|
||||
|
||||
with _PILImage.open(io.BytesIO(_data_url_bytes(image_data_url))) as im:
|
||||
w, h = im.size
|
||||
x1, y1, x2, y2 = (float(v) for v in box)
|
||||
bw, bh = x2 - x1, y2 - y1
|
||||
x1, y1 = max(0.0, x1 - bw * pad_frac), max(0.0, y1 - bh * pad_frac)
|
||||
x2, y2 = min(float(w), x2 + bw * pad_frac), min(float(h), y2 + bh * pad_frac)
|
||||
mask = _PILImage.new("L", (w, h), 0)
|
||||
ImageDraw.Draw(mask).rectangle([x1, y1, x2, y2], fill=255)
|
||||
buf = io.BytesIO()
|
||||
mask.save(buf, "PNG")
|
||||
return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def _ground_box(owner: Optional[str], image_data_url: str, label: str) -> Optional[List[float]]:
|
||||
"""Use the in-app grounding detector (OWL-ViT via /api/image/mask) for a box."""
|
||||
try:
|
||||
data = await _post("/api/image/mask", {"image": image_data_url, "text": label}, owner, timeout=90)
|
||||
except Exception as exc: # pragma: no cover - defensive
|
||||
logger.warning("editor grounding box failed: %s", exc)
|
||||
return None
|
||||
box = data.get("bbox") if isinstance(data, dict) else None
|
||||
if isinstance(box, (list, tuple)) and len(box) == 4:
|
||||
try:
|
||||
return [float(v) for v in box]
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _ensure_keep_clause(text: str) -> str:
|
||||
"""Append a keep-the-subject clause when the prompt lacks one."""
|
||||
body = str(text or "").strip()
|
||||
if not body:
|
||||
return body
|
||||
if re.search(r"\b(keep|preserve|leave|unchanged|identical|same)\b", body, re.I):
|
||||
return body
|
||||
return body.rstrip(". ") + ". Keep the main subject and everything else in the image exactly the same."
|
||||
|
||||
|
||||
def _mask_label(args: Dict[str, Any], prompt: str) -> str:
|
||||
"""Best-effort object noun for grounding a regional removal."""
|
||||
label = str(args.get("label") or args.get("object") or "").strip()
|
||||
if label:
|
||||
return label
|
||||
m = re.search(
|
||||
r"\b(?:remove|delete|erase|eliminate|get\s+rid\s+of|take\s+out)\s+(?:the\s+|a\s+|an\s+)?"
|
||||
r"([a-z][a-z \-]{1,30}?)(?:\s+(?:on|in|at|from|to|near|behind|beside|by)\b|[.,;!?]|$)",
|
||||
prompt or "",
|
||||
re.I,
|
||||
)
|
||||
if m:
|
||||
noun = m.group(1).strip().lower()
|
||||
if noun and noun not in {"background", "bg"}:
|
||||
return noun
|
||||
for noun in ("person", "people", "man", "woman", "car", "dog", "sign", "logo"):
|
||||
if re.search(rf"\b{re.escape(noun)}\b", prompt or "", re.I):
|
||||
return "person" if noun in {"person", "people", "man", "woman"} else noun
|
||||
return "object"
|
||||
|
||||
|
||||
class PlanEditTool:
|
||||
"""Plan an image edit as ordered editor operations (no side effects)."""
|
||||
|
||||
async def execute(self, content: Any, ctx: Dict[str, Any]):
|
||||
args = _parse(content)
|
||||
prompt = str(args.get("prompt") or "").strip()
|
||||
if not prompt:
|
||||
return "plan_edit: invalid", {"error": "plan_edit needs a non-empty `prompt`.", "exit_code": 1}
|
||||
state = args.get("state") if isinstance(args.get("state"), dict) else {}
|
||||
try:
|
||||
data = await _post(
|
||||
"/api/editor/plan",
|
||||
{"prompt": prompt, "state": state, "mode": "plan", "session": _session(args, ctx)},
|
||||
ctx.get("owner"),
|
||||
)
|
||||
except Exception as exc:
|
||||
return "plan_edit: error", {"error": f"plan_edit failed: {exc}", "exit_code": 1}
|
||||
steps = data.get("steps") or []
|
||||
reply = str(data.get("reply") or "")
|
||||
return f"plan_edit: {len(steps)} step(s)", {
|
||||
"output": reply or f"Planned {len(steps)} editor step(s).",
|
||||
"reply": reply,
|
||||
"planner": data.get("planner"),
|
||||
"steps": steps,
|
||||
"exit_code": 0,
|
||||
}
|
||||
|
||||
|
||||
class RunEditorCommandTool:
|
||||
"""Run a single validated editor op in the browser via the owner bridge."""
|
||||
|
||||
async def execute(self, content: Any, ctx: Dict[str, Any]):
|
||||
args = _parse(content)
|
||||
op = str(args.get("op") or "").strip()
|
||||
if not op:
|
||||
return "run_editor_command: invalid", {"error": "run_editor_command needs `op`.", "exit_code": 1}
|
||||
params = args.get("params") if isinstance(args.get("params"), dict) else {}
|
||||
payload = {
|
||||
"session": _session(args, ctx),
|
||||
"op": op,
|
||||
"params": params,
|
||||
"wait": bool(args.get("wait", True)),
|
||||
"timeout": float(args.get("timeout") or 60),
|
||||
}
|
||||
try:
|
||||
data = await _post("/api/editor/bridge/command", payload, ctx.get("owner"))
|
||||
except Exception as exc:
|
||||
return f"run_editor_command: {op}", {"error": f"run_editor_command failed: {exc}", "exit_code": 1}
|
||||
ok = data.get("ok")
|
||||
result = data.get("result")
|
||||
summary = json.dumps({"ok": ok, "error": data.get("error"), "result": result}, ensure_ascii=False)
|
||||
return f"run_editor_command: {op}", {
|
||||
"output": summary[:4000],
|
||||
"result": result,
|
||||
"screenshot": data.get("screenshot"),
|
||||
"ok": ok,
|
||||
"exit_code": 0 if ok is not False else 1,
|
||||
}
|
||||
|
||||
|
||||
class GetEditorScreenshotTool:
|
||||
"""Capture the current editor canvas as a data URL."""
|
||||
|
||||
async def execute(self, content: Any, ctx: Dict[str, Any]):
|
||||
args = _parse(content)
|
||||
payload = {
|
||||
"session": _session(args, ctx),
|
||||
"op": "screenshot",
|
||||
"params": {},
|
||||
"wait": True,
|
||||
"timeout": float(args.get("timeout") or 30),
|
||||
}
|
||||
try:
|
||||
data = await _post("/api/editor/bridge/command", payload, ctx.get("owner"), timeout=45)
|
||||
except Exception as exc:
|
||||
return "get_editor_screenshot: error", {"error": f"get_editor_screenshot failed: {exc}", "exit_code": 1}
|
||||
result = data.get("result") if isinstance(data.get("result"), dict) else {}
|
||||
shot = data.get("screenshot") or result.get("dataUrl")
|
||||
if not shot:
|
||||
return "get_editor_screenshot: unavailable", {
|
||||
"error": "No editor screenshot was returned (is the editor open?).",
|
||||
"exit_code": 1,
|
||||
}
|
||||
return "get_editor_screenshot", {"output": "Captured the editor canvas.", "screenshot": shot, "exit_code": 0}
|
||||
|
||||
|
||||
class AiEditImageTool:
|
||||
"""Apply an AI image transform (remove-bg/upscale/denoise/enhance/inpaint)."""
|
||||
|
||||
async def execute(self, content: Any, ctx: Dict[str, Any]):
|
||||
args = _parse(content)
|
||||
_t = {"start": time.perf_counter()}
|
||||
op = str(args.get("op") or "").strip().lower()
|
||||
if op not in _AI_OPS:
|
||||
return "ai_edit_image: invalid", {
|
||||
"error": f"ai_edit_image `op` must be one of {sorted(_AI_OPS)}.",
|
||||
"exit_code": 1,
|
||||
}
|
||||
owner = ctx.get("owner")
|
||||
if op in ("inpaint", "generative") and not _provider_available(owner):
|
||||
return "ai_edit_image: provider required", {
|
||||
"error": _provider_message(),
|
||||
"failure_kind": "image_provider_unavailable",
|
||||
"needs_provider": True,
|
||||
"exit_code": 1,
|
||||
}
|
||||
session = _session(args, ctx)
|
||||
image = _as_data_url(args.get("image"))
|
||||
if image is None and isinstance(args.get("image"), str) and args["image"].startswith("/"):
|
||||
# A same-origin gallery/attachment URL: fetch it (owner-scoped) as bytes.
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=60) as client:
|
||||
resp = await client.get(f"{_INTERNAL_BASE}{args['image']}", headers=_internal_headers(owner))
|
||||
resp.raise_for_status()
|
||||
image = _as_data_url("data:image/png;base64," + base64.b64encode(resp.content).decode())
|
||||
except Exception:
|
||||
image = None
|
||||
used_canvas = False
|
||||
if not image:
|
||||
shot = await GetEditorScreenshotTool().execute(json.dumps({"session": session}), ctx)
|
||||
image = (shot[1] or {}).get("screenshot") if isinstance(shot, tuple) else None
|
||||
used_canvas = bool(image)
|
||||
if not image:
|
||||
return "ai_edit_image: no image", {
|
||||
"error": "ai_edit_image needs an `image` (data URL or same-origin URL) or an open editor.",
|
||||
"exit_code": 1,
|
||||
}
|
||||
_t["image"] = time.perf_counter()
|
||||
payload: Dict[str, Any] = {"image": image}
|
||||
model_choice = str(args.get("model") or "").strip()
|
||||
mask_used = False
|
||||
place_mode: Optional[str] = None
|
||||
if op == "upscale":
|
||||
payload["scale"] = int(args.get("scale") or 2)
|
||||
elif op == "denoise":
|
||||
payload["strength"] = float(args.get("strength") or 0.5)
|
||||
elif op in ("inpaint", "generative"):
|
||||
payload["prompt"] = str(args.get("prompt") or "").strip()
|
||||
if model_choice:
|
||||
# inpaint_proxy reads `_model` to route to a specific model.
|
||||
payload["_model"] = model_choice
|
||||
mode = str(args.get("mode") or "instruction").strip().lower()
|
||||
_steps = int(args.get("steps") or 0)
|
||||
if _steps > 0:
|
||||
payload["steps"] = _steps
|
||||
mask_used = False
|
||||
mask = None
|
||||
if mode == "masked":
|
||||
mask = _as_data_url(args.get("mask"))
|
||||
if not mask:
|
||||
mask = await _fetch_selection_mask(owner, session)
|
||||
# An empty/tiny SAM mask on small figures is useless: drop it.
|
||||
if mask and _mask_white_fraction(mask) < _MIN_MASK_FRACTION:
|
||||
mask = None
|
||||
if not mask:
|
||||
# Fall back to the grounding detector (OWL-ViT): a dilated
|
||||
# box mask around the matching object.
|
||||
box = await _ground_box(owner, image, _mask_label(args, payload["prompt"]))
|
||||
if box:
|
||||
candidate = _rect_mask_data_url(image, box, pad_frac=0.15)
|
||||
if candidate and _mask_white_fraction(candidate) >= _MIN_MASK_FRACTION:
|
||||
mask = candidate
|
||||
if mask:
|
||||
payload["mask"] = mask
|
||||
payload["mode"] = "masked"
|
||||
mask_used = True
|
||||
else:
|
||||
# Nothing usable to mask -> whole-image instruction fallback
|
||||
# rather than placing a transparent/empty layer.
|
||||
mode = "instruction"
|
||||
if mode != "masked":
|
||||
# Instruction mode: whole-image edit with no tight mask. Send a
|
||||
# full-white mask so the edit endpoint accepts the request and
|
||||
# regenerates the whole frame under the instruction prompt, which
|
||||
# names what must stay unchanged. This avoids the mask-compositing
|
||||
# halos that loose selections caused.
|
||||
payload["mode"] = "instruction"
|
||||
payload.pop("mask", None)
|
||||
payload["prompt"] = _ensure_keep_clause(payload["prompt"])
|
||||
try:
|
||||
from PIL import Image as _PILImage
|
||||
|
||||
_raw = _data_url_bytes(image)
|
||||
_w, _h = _PILImage.open(io.BytesIO(_raw)).size
|
||||
_mbuf = io.BytesIO()
|
||||
_PILImage.new("L", (_w, _h), 255).save(_mbuf, "PNG")
|
||||
payload["mask"] = "data:image/png;base64," + base64.b64encode(_mbuf.getvalue()).decode()
|
||||
except Exception:
|
||||
pass
|
||||
place_mode = "masked" if mask_used else "instruction"
|
||||
_t["prep"] = time.perf_counter()
|
||||
try:
|
||||
data = await _post(_AI_OPS[op], payload, owner)
|
||||
except Exception as exc:
|
||||
logger.info(
|
||||
"image-op ai_edit_image op=%s stage=provider FAILED image_ms=%.1f prep_ms=%.1f total_ms=%.1f err=%s",
|
||||
op, (_t["image"] - _t["start"]) * 1000.0, (_t["prep"] - _t["image"]) * 1000.0,
|
||||
(time.perf_counter() - _t["start"]) * 1000.0, exc,
|
||||
)
|
||||
return f"ai_edit_image: {op}", {"error": f"ai_edit_image ({op}) failed: {exc}", "exit_code": 1}
|
||||
_t["provider"] = time.perf_counter()
|
||||
result = data.get("image") or data.get("data")
|
||||
data_url = _as_data_url(result)
|
||||
if not data_url:
|
||||
return f"ai_edit_image: {op}", {"error": f"ai_edit_image ({op}) returned no image.", "exit_code": 1}
|
||||
if _image_is_blank(data_url):
|
||||
return f"ai_edit_image: {op}", {
|
||||
"error": f"ai_edit_image ({op}) produced an empty layer.",
|
||||
"failure_kind": "empty_result",
|
||||
"exit_code": 1,
|
||||
}
|
||||
|
||||
info: Dict[str, Any] = {
|
||||
"output": f"Applied AI edit '{op}'" + (" to the current canvas." if used_canvas else "."),
|
||||
"op": op,
|
||||
"model": data.get("model") or model_choice or "",
|
||||
"exit_code": 0,
|
||||
}
|
||||
try:
|
||||
info.update(_persist_png(_data_url_bytes(data_url), owner, ctx.get("session_id"), f"ai_edit_image:{op}", f"ai_edit_image:{op}"))
|
||||
except Exception as exc:
|
||||
logger.warning("ai_edit_image persist failed: %s", exc)
|
||||
_t["persist"] = time.perf_counter()
|
||||
if bool(args.get("place", True)):
|
||||
placed = await _place_layer(
|
||||
data_url,
|
||||
f"AI {op}",
|
||||
owner,
|
||||
session,
|
||||
masked=mask_used,
|
||||
blend=str(args.get("blend") or "normal"),
|
||||
opacity=float(args.get("opacity") or 100),
|
||||
mode=place_mode,
|
||||
)
|
||||
info["placed"] = placed
|
||||
if placed:
|
||||
info["output"] += " Placed the result as a new layer."
|
||||
_t["place"] = time.perf_counter()
|
||||
logger.info(
|
||||
"image-op ai_edit_image op=%s mode=%s canvas=%s image_ms=%.1f mask_ms=%.1f "
|
||||
"model_ms=%.1f persist_ms=%.1f place_ms=%.1f total_ms=%.1f",
|
||||
op, place_mode or "instruction", used_canvas,
|
||||
(_t["image"] - _t["start"]) * 1000.0,
|
||||
(_t["prep"] - _t["image"]) * 1000.0,
|
||||
(_t["provider"] - _t["prep"]) * 1000.0,
|
||||
(_t["persist"] - _t["provider"]) * 1000.0,
|
||||
(_t["place"] - _t["persist"]) * 1000.0,
|
||||
(_t["place"] - _t["start"]) * 1000.0,
|
||||
)
|
||||
return f"ai_edit_image: {op}", info
|
||||
|
||||
|
||||
class GenerateImageLayerTool:
|
||||
"""Generate an image with the configured image model and place it as a layer."""
|
||||
|
||||
async def execute(self, content: str, ctx: dict):
|
||||
args = _parse(content)
|
||||
prompt = str(args.get("prompt") or "").strip()
|
||||
if not prompt:
|
||||
return "generate_image: invalid", {"error": "generate_image needs a non-empty `prompt`.", "exit_code": 1}
|
||||
owner = ctx.get("owner")
|
||||
if not _provider_available(owner):
|
||||
return "generate_image: provider required", {
|
||||
"error": _provider_message(),
|
||||
"failure_kind": "image_provider_unavailable",
|
||||
"needs_provider": True,
|
||||
"exit_code": 1,
|
||||
}
|
||||
session = _session(args, ctx)
|
||||
payload: Dict[str, Any] = {"prompt": prompt}
|
||||
for key in ("model", "size", "quality"):
|
||||
if args.get(key):
|
||||
payload[key] = args[key]
|
||||
try:
|
||||
from src.ai_interaction import do_generate_image
|
||||
|
||||
data = await do_generate_image(
|
||||
json.dumps(payload, ensure_ascii=False),
|
||||
session_id=session,
|
||||
owner=owner,
|
||||
incognito=True,
|
||||
)
|
||||
except Exception as exc:
|
||||
return "generate_image: error", {"error": f"generate_image failed: {exc}", "exit_code": 1}
|
||||
result = data.get("image_url") or data.get("image") or data.get("data")
|
||||
data_url = _as_data_url(result)
|
||||
if not data_url:
|
||||
return "generate_image", {"error": "generate_image returned no image.", "exit_code": 1}
|
||||
info: Dict[str, Any] = {
|
||||
"output": f"Generated an image for: {prompt[:120]}",
|
||||
"image_size": data.get("image_size"),
|
||||
"model": data.get("image_model") or str(args.get("model") or ""),
|
||||
"exit_code": 0,
|
||||
}
|
||||
try:
|
||||
info.update(_persist_png(_data_url_bytes(data_url), owner, ctx.get("session_id"), prompt, "generate_image"))
|
||||
except Exception as exc:
|
||||
logger.warning("generate_image persist failed: %s", exc)
|
||||
if bool(args.get("place", True)):
|
||||
placed = await _place_layer(
|
||||
data_url,
|
||||
str(args.get("name") or "Generated image"),
|
||||
owner,
|
||||
session,
|
||||
masked=bool(args.get("mask")),
|
||||
blend=str(args.get("blend") or "normal"),
|
||||
opacity=float(args.get("opacity") or 100),
|
||||
)
|
||||
info["placed"] = placed
|
||||
if placed:
|
||||
info["output"] += " Placed it as a new layer."
|
||||
return "generate_image", info
|
||||
@@ -0,0 +1,498 @@
|
||||
"""Editor agent/planner model resolution.
|
||||
|
||||
The image-editor Quick Edit planner (`/api/editor/plan`, `/api/editor/agent`)
|
||||
runs on a LOCAL Qwen3.5-9B vision model served on seneca when one is
|
||||
reachable, instead of the hosted chat model. This module owns:
|
||||
|
||||
* the local model base URLs / model id (settings-overridable),
|
||||
* an ordered candidate list (per-request ``agent_model`` override, then the
|
||||
local Qwen3.5-9B, then the hosted default/utility chat model),
|
||||
* a vision fallback chain: local model WITH the canvas screenshot, else a
|
||||
text-only call whose image is described by a separate VL caption model,
|
||||
* a lightweight reachability probe for the configured local bases.
|
||||
|
||||
Everything is best-effort: when no local server is reachable the planner
|
||||
transparently falls back to the previously-used hosted chat model, so the
|
||||
editor agent never regresses to "no plan".
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextvars
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# The local Qwen3.5-9B vision server (seneca). :8097 (tmux AGENT_VLM) is the
|
||||
# canonical editor-agent endpoint. Override with the ``editor_agent_base_url``
|
||||
# setting (single URL wins) or the ``ODYSSEUS_EDITOR_AGENT_BASES`` env
|
||||
# (comma-separated). NOTE: :8095 is a DIFFERENT server (Qwen3-VL-30B) used ONLY
|
||||
# for the caption fallback below — never list it here as an agent model.
|
||||
DEFAULT_LOCAL_BASES = [
|
||||
"http://100.69.120.65:8097/v1",
|
||||
]
|
||||
DEFAULT_LOCAL_MODEL = "Qwen/Qwen3.5-9B"
|
||||
|
||||
# Vision caption model used when the editor-agent model cannot accept images.
|
||||
DEFAULT_CAPTION_BASE = "http://100.69.120.65:8095/v1"
|
||||
DEFAULT_CAPTION_MODEL = "Qwen/Qwen3.5-9B"
|
||||
|
||||
# Served-model-id cache per base (populated on first use). The local servers 400
|
||||
# with "Failed to load model" when asked for an id they do not serve (e.g. the
|
||||
# client's default "Qwen3.5-9B-MLX-8bit", or a stale caption model), so we
|
||||
# resolve the requested id against GET {base}/models and fall back to one that
|
||||
# is actually served.
|
||||
_MODELS_CACHE: Dict[str, set] = {}
|
||||
|
||||
_CAPTION_INSTRUCTION = (
|
||||
"Describe this image in one or two sentences for a photo editor: the main "
|
||||
"subject, the background/scene, lighting, dominant colours, and any obvious "
|
||||
"flaws (blown highlights, colour cast, clutter, blur)."
|
||||
)
|
||||
|
||||
|
||||
def _setting(key: str, default: Any = "") -> Any:
|
||||
try:
|
||||
from src.settings import load_settings
|
||||
|
||||
value = load_settings().get(key)
|
||||
return value if value not in (None, "") else default
|
||||
except Exception:
|
||||
return default
|
||||
|
||||
|
||||
def local_bases() -> List[str]:
|
||||
override = _setting("editor_agent_base_url")
|
||||
if override:
|
||||
return [str(override)]
|
||||
env = os.getenv("ODYSSEUS_EDITOR_AGENT_BASES")
|
||||
if env:
|
||||
bases = [b.strip() for b in env.split(",") if b.strip()]
|
||||
if bases:
|
||||
return bases
|
||||
return list(DEFAULT_LOCAL_BASES)
|
||||
|
||||
|
||||
def local_model() -> str:
|
||||
return str(_setting("editor_agent_model", DEFAULT_LOCAL_MODEL))
|
||||
|
||||
|
||||
def caption_base() -> str:
|
||||
return str(_setting("editor_caption_base_url", DEFAULT_CAPTION_BASE))
|
||||
|
||||
|
||||
def caption_model() -> str:
|
||||
return str(_setting("editor_caption_model", DEFAULT_CAPTION_MODEL))
|
||||
|
||||
|
||||
# Per-request / per-plan context (contextvars keep this correct under the
|
||||
# async event loop without sharing state between concurrent editor turns).
|
||||
_override: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
|
||||
"editor_agent_model_override", default=None
|
||||
)
|
||||
_last_model: contextvars.ContextVar[Optional[Tuple[str, str]]] = contextvars.ContextVar(
|
||||
"editor_agent_last_model", default=None
|
||||
)
|
||||
|
||||
|
||||
def set_override(value: Optional[str]) -> None:
|
||||
_override.set(value or None)
|
||||
|
||||
|
||||
def get_override() -> Optional[str]:
|
||||
return _override.get()
|
||||
|
||||
|
||||
def set_last_model(model: str, source: str) -> None:
|
||||
_last_model.set((model, source))
|
||||
|
||||
|
||||
def get_last_model() -> Optional[Tuple[str, str]]:
|
||||
return _last_model.get()
|
||||
|
||||
|
||||
def _chat_url(base: str) -> str:
|
||||
b = str(base or "").rstrip("/")
|
||||
if b.endswith("/chat/completions"):
|
||||
return b
|
||||
return b + "/chat/completions"
|
||||
|
||||
|
||||
def _split_override(override: str) -> Tuple[str, str]:
|
||||
"""Parse ``base::model`` (or a bare model id)."""
|
||||
raw = str(override or "").strip()
|
||||
if "::" in raw:
|
||||
base, model = raw.split("::", 1)
|
||||
return base.strip(), model.strip()
|
||||
return "", raw
|
||||
|
||||
|
||||
def _has_image(messages: List[dict]) -> bool:
|
||||
for msg in messages:
|
||||
content = msg.get("content")
|
||||
if isinstance(content, list):
|
||||
for part in content:
|
||||
if isinstance(part, dict) and part.get("type") == "image_url":
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _first_image(messages: List[dict]) -> Optional[str]:
|
||||
for msg in messages:
|
||||
content = msg.get("content")
|
||||
if isinstance(content, list):
|
||||
for part in content:
|
||||
if isinstance(part, dict) and part.get("type") == "image_url":
|
||||
url = (part.get("image_url") or {}).get("url")
|
||||
if isinstance(url, str):
|
||||
return url
|
||||
return None
|
||||
|
||||
|
||||
def _strip_images(messages: List[dict]) -> List[dict]:
|
||||
out: List[dict] = []
|
||||
for msg in messages:
|
||||
content = msg.get("content")
|
||||
if isinstance(content, list):
|
||||
text = "\n".join(
|
||||
str(part.get("text") or "")
|
||||
for part in content
|
||||
if isinstance(part, dict) and part.get("type") == "text"
|
||||
)
|
||||
out.append({"role": msg.get("role", "user"), "content": text})
|
||||
else:
|
||||
out.append({"role": msg.get("role", "user"), "content": content})
|
||||
return out
|
||||
|
||||
|
||||
def _with_caption(messages: List[dict], caption: str) -> List[dict]:
|
||||
note = f"\n\n[Canvas screenshot caption: {caption}]"
|
||||
out = [dict(m) for m in messages]
|
||||
for msg in reversed(out):
|
||||
if msg.get("role") == "user":
|
||||
msg["content"] = str(msg.get("content") or "") + note
|
||||
break
|
||||
return out
|
||||
|
||||
|
||||
async def _post_chat(
|
||||
base: str,
|
||||
model: str,
|
||||
messages: List[dict],
|
||||
headers: Optional[dict],
|
||||
timeout: float,
|
||||
max_tokens: int = 1024,
|
||||
response_format: Optional[dict] = None,
|
||||
) -> Optional[str]:
|
||||
if not base or not model:
|
||||
return None
|
||||
payload: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"temperature": 0.2,
|
||||
"max_tokens": max_tokens,
|
||||
"stream": False,
|
||||
}
|
||||
# Constrain the planner to a strict JSON ops schema (structured outputs /
|
||||
# grammar) so the model can only emit valid editor ops. Unknown servers
|
||||
# ignore the field; the server-side schema validation is the real enforcer.
|
||||
if isinstance(response_format, dict):
|
||||
payload["response_format"] = response_format
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
resp = await client.post(_chat_url(base), json=payload, headers=headers or {})
|
||||
if resp.status_code != 200:
|
||||
logger.debug(
|
||||
"editor agent chat %s -> %s (%s)", base, resp.status_code, resp.text[:200]
|
||||
)
|
||||
return None
|
||||
data = resp.json()
|
||||
except Exception as exc:
|
||||
logger.debug("editor agent chat %s failed: %s", base, exc)
|
||||
return None
|
||||
try:
|
||||
return data["choices"][0]["message"]["content"]
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def _served_models(base: str, timeout: float = 3.0) -> set:
|
||||
"""Return the set of model ids a base serves (cached; empty on failure)."""
|
||||
if not base:
|
||||
return set()
|
||||
cached = _MODELS_CACHE.get(base)
|
||||
if cached is not None:
|
||||
return cached
|
||||
ids: set = set()
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
resp = await client.get(base.rstrip("/") + "/models")
|
||||
if resp.status_code == 200:
|
||||
ids = {
|
||||
str(m.get("id"))
|
||||
for m in (resp.json().get("data") or [])
|
||||
if isinstance(m, dict) and m.get("id")
|
||||
}
|
||||
except Exception:
|
||||
ids = set()
|
||||
_MODELS_CACHE[base] = ids
|
||||
return ids
|
||||
|
||||
|
||||
async def _resolve_model(base: str, requested: str, fallback: str) -> str:
|
||||
"""Return a model id ``base`` actually serves.
|
||||
|
||||
If the requested id is served, keep it; otherwise use ``fallback`` (when
|
||||
served) or any served id. When the served list is unknown (server down),
|
||||
keep the requested id (or ``fallback``).
|
||||
"""
|
||||
requested = str(requested or "").strip()
|
||||
fallback = str(fallback or "").strip()
|
||||
ids = await _served_models(base)
|
||||
if ids:
|
||||
if requested and requested in ids:
|
||||
return requested
|
||||
if fallback in ids:
|
||||
return fallback
|
||||
return sorted(ids)[0]
|
||||
return requested or fallback
|
||||
|
||||
|
||||
async def caption_image(data_url: str, owner: Optional[str] = None, timeout: float = 90.0) -> Optional[str]:
|
||||
if not isinstance(data_url, str) or not data_url.startswith("data:image/"):
|
||||
return None
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": _CAPTION_INSTRUCTION},
|
||||
{"type": "image_url", "image_url": {"url": data_url}},
|
||||
],
|
||||
}
|
||||
]
|
||||
base = caption_base()
|
||||
model = await _resolve_model(base, caption_model(), local_model())
|
||||
return await _post_chat(base, model, messages, None, timeout, max_tokens=200)
|
||||
|
||||
|
||||
_DONE_RE = re.compile(r"^\s*DONE\b", re.I)
|
||||
_RETRY_RE = re.compile(r"RETRY:?\s*(.+)", re.I | re.S)
|
||||
|
||||
|
||||
def _parse_self_check(text: str) -> Dict[str, Any]:
|
||||
s = str(text or "").strip()
|
||||
if _DONE_RE.match(s):
|
||||
return {"done": True, "issues": [], "retry_hint": "", "text": s}
|
||||
m = _RETRY_RE.search(s)
|
||||
if m or re.search(r"\bretry\b", s, re.I):
|
||||
hint = (m.group(1).strip() if m else s) or s
|
||||
return {"done": False, "issues": [hint], "retry_hint": hint, "text": s}
|
||||
return {"done": True, "issues": [], "retry_hint": "", "text": s}
|
||||
|
||||
|
||||
async def self_check(
|
||||
before: Any,
|
||||
after: Any,
|
||||
prompt: str,
|
||||
model: Optional[str] = None,
|
||||
owner: Optional[str] = None,
|
||||
timeout: float = 30.0,
|
||||
) -> Dict[str, Any]:
|
||||
"""Same-origin VL self-check: compare BEFORE vs AFTER against the intent.
|
||||
|
||||
The browser cannot call the local VL server directly (it sends no CORS
|
||||
headers), so ``/api/editor/self-check`` proxies here. Returns
|
||||
``{done, issues, retry_hint, text}``; on any failure (bad input, the model
|
||||
is unreachable, or a timeout) it degrades to ``done=True`` with
|
||||
``reason="unavailable"`` so the client never gets stuck retrying.
|
||||
"""
|
||||
if not (
|
||||
isinstance(before, str) and before.startswith("data:image/")
|
||||
and isinstance(after, str) and after.startswith("data:image/")
|
||||
):
|
||||
return {"done": True, "issues": [], "retry_hint": "",
|
||||
"reason": "unavailable", "text": "DONE"}
|
||||
bases = local_bases() or list(DEFAULT_LOCAL_BASES)
|
||||
base = bases[0]
|
||||
mdl = await _resolve_model(base, str(model or local_model()), local_model())
|
||||
messages = [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": str(prompt or "")},
|
||||
{"type": "image_url", "image_url": {"url": before}},
|
||||
{"type": "image_url", "image_url": {"url": after}},
|
||||
],
|
||||
}]
|
||||
text = await _post_chat(base, mdl, messages, None, timeout, max_tokens=140)
|
||||
if not text:
|
||||
return {"done": True, "issues": [], "retry_hint": "",
|
||||
"reason": "unavailable", "text": "DONE"}
|
||||
out = _parse_self_check(text)
|
||||
out["model"] = mdl
|
||||
out["source"] = "local"
|
||||
return out
|
||||
|
||||
|
||||
def _candidates(owner: Optional[str], override: Optional[str]) -> List[Dict[str, Any]]:
|
||||
cands: List[Dict[str, Any]] = []
|
||||
if override:
|
||||
base, model = _split_override(override)
|
||||
if model:
|
||||
cands.append(
|
||||
{"base": base or local_bases()[0], "model": model, "headers": None, "source": "override"}
|
||||
)
|
||||
for base in local_bases():
|
||||
cands.append({"base": base, "model": local_model(), "headers": None, "source": "local"})
|
||||
try:
|
||||
from src.endpoint_resolver import resolve_endpoint
|
||||
|
||||
for prefix in ("utility", "default"):
|
||||
try:
|
||||
url, model, headers = resolve_endpoint(prefix, owner=owner)
|
||||
except Exception:
|
||||
continue
|
||||
if url and model:
|
||||
cands.append({"base": url, "model": model, "headers": headers, "source": "hosted"})
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
return cands
|
||||
|
||||
|
||||
async def plan_with_editor_agent(
|
||||
messages: List[dict],
|
||||
owner: Optional[str] = None,
|
||||
override: Optional[str] = None,
|
||||
timeout: float = 90.0,
|
||||
response_format: Optional[dict] = None,
|
||||
) -> Optional[Tuple[str, str, str]]:
|
||||
"""Return (raw_plan_text, model_used, source) or None.
|
||||
|
||||
Tries the local model WITH the screenshot, then the same model text-only
|
||||
with a VL caption of the screenshot, then the hosted chat model.
|
||||
|
||||
``response_format`` (a JSON-schema structured-output constraint) is sent
|
||||
ONLY to the local editor-agent bases, so a hosted fallback model that does
|
||||
not support the field never rejects the request.
|
||||
"""
|
||||
has_img = _has_image(messages)
|
||||
text_messages = _strip_images(messages) if has_img else messages
|
||||
image = _first_image(messages) if has_img else None
|
||||
caption: Optional[str] = None
|
||||
_local = {str(b).rstrip("/") for b in local_bases()}
|
||||
for cand in _candidates(owner, override):
|
||||
base, model = cand["base"], cand["model"]
|
||||
headers = cand.get("headers")
|
||||
is_local = cand["source"] == "local" or str(base).rstrip("/") in _local
|
||||
rf = response_format if (response_format and is_local) else None
|
||||
extra = {"response_format": rf} if rf else {}
|
||||
if has_img:
|
||||
raw = await _post_chat(base, model, messages, headers, timeout, **extra)
|
||||
if raw:
|
||||
return raw, model, cand["source"]
|
||||
if caption is None and image:
|
||||
caption = await caption_image(image, owner)
|
||||
if caption:
|
||||
raw = await _post_chat(
|
||||
base, model, _with_caption(text_messages, caption), headers, timeout, **extra,
|
||||
)
|
||||
if raw:
|
||||
return raw, model, cand["source"] + "+caption"
|
||||
else:
|
||||
raw = await _post_chat(
|
||||
base, model, text_messages, headers, timeout, **extra
|
||||
)
|
||||
if raw:
|
||||
return raw, model, cand["source"]
|
||||
return None
|
||||
|
||||
|
||||
async def list_editor_agent_models(owner: Optional[str] = None) -> List[Dict[str, Any]]:
|
||||
"""Models available to the editor agent, for the model picker.
|
||||
|
||||
Lists the local Qwen3.5-9B (one entry per reachable base) plus the hosted
|
||||
default/utility chat model that the planner falls back to.
|
||||
"""
|
||||
models: List[Dict[str, Any]] = []
|
||||
seen: set = set()
|
||||
for base in local_bases():
|
||||
model = local_model()
|
||||
key = ("local", base, model)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
models.append({
|
||||
"id": model,
|
||||
"label": f"{model} · local",
|
||||
"source": "local",
|
||||
"vision": True,
|
||||
"base_url": base,
|
||||
})
|
||||
try:
|
||||
from src.endpoint_resolver import resolve_endpoint
|
||||
|
||||
for prefix in ("utility", "default"):
|
||||
try:
|
||||
url, model, _headers = resolve_endpoint(prefix, owner=owner)
|
||||
except Exception:
|
||||
continue
|
||||
if url and model:
|
||||
key = ("hosted", model)
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
models.append({
|
||||
"id": model,
|
||||
"label": model,
|
||||
"source": "hosted",
|
||||
"vision": None,
|
||||
"base_url": url,
|
||||
})
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
return models
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Reachability probe (best effort, non-blocking).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_PROBE: Dict[str, bool] = {}
|
||||
|
||||
|
||||
def probe_models(base: str, timeout: float = 3.0) -> bool:
|
||||
try:
|
||||
with httpx.Client(timeout=timeout) as client:
|
||||
resp = client.get(str(base).rstrip("/") + "/models")
|
||||
return resp.status_code == 200
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _probe_worker() -> None:
|
||||
for base in local_bases():
|
||||
ok = probe_models(base)
|
||||
_PROBE[base] = ok
|
||||
logger.info("editor agent local base %s reachable=%s", base, ok)
|
||||
cap = caption_base()
|
||||
_PROBE[cap] = probe_models(cap)
|
||||
logger.info("editor agent caption base %s reachable=%s", cap, _PROBE[cap])
|
||||
|
||||
|
||||
def startup_probe() -> None:
|
||||
"""Kick off a background reachability probe (never blocks startup)."""
|
||||
try:
|
||||
threading.Thread(target=_probe_worker, name="editor-agent-probe", daemon=True).start()
|
||||
except Exception:
|
||||
logger.debug("editor agent probe thread failed", exc_info=True)
|
||||
|
||||
|
||||
def probe_status() -> Dict[str, bool]:
|
||||
return dict(_PROBE)
|
||||
File diff suppressed because it is too large
Load Diff
+42
-23
@@ -68,6 +68,7 @@ def _row(index: int, message: dict[str, Any]) -> dict[str, Any]:
|
||||
})
|
||||
return {
|
||||
"uid": str(message.get("uid") or _uid(index)),
|
||||
"message_id": str(message.get("message_id") or ""),
|
||||
"account": str(message.get("account") or "fixture"),
|
||||
"account_email": str(message.get("account_email") or message.get("to") or ""),
|
||||
"account_id": str(message.get("account_id") or "fixture"),
|
||||
@@ -78,6 +79,16 @@ def _row(index: int, message: dict[str, Any]) -> dict[str, Any]:
|
||||
"folder": str(message.get("folder") or "INBOX"),
|
||||
"attachments": attachments,
|
||||
"has_attachments": bool(attachments),
|
||||
"summary": str(message.get("summary") or message.get("body") or ""),
|
||||
}
|
||||
|
||||
|
||||
def _account_matches(message: dict[str, Any], account: str | None) -> bool:
|
||||
selector = str(account or "").strip().casefold()
|
||||
return not selector or selector in {
|
||||
str(message.get("account") or "fixture").strip().casefold(),
|
||||
str(message.get("account_email") or message.get("to") or "").strip().casefold(),
|
||||
str(message.get("account_id") or "fixture").strip().casefold(),
|
||||
}
|
||||
|
||||
|
||||
@@ -108,9 +119,16 @@ def execute_fixture_email(tool: str, args: dict[str, Any], owner: str | None = N
|
||||
|
||||
bare = tool.removeprefix("mcp__email__")
|
||||
if bare == "list_email_accounts":
|
||||
accounts = {}
|
||||
for index, message in enumerate(messages):
|
||||
row = _row(index, message)
|
||||
identity = (row["account_id"], row["account"], row["account_email"])
|
||||
accounts.setdefault(identity, {"id": row["account_id"], "name": row["account"],
|
||||
"email": row["account_email"], "default": not accounts})
|
||||
return {
|
||||
"accounts": [{"id": "fixture", "name": "Fixture mailbox", "default": True}],
|
||||
"output": "Fixture mailbox (default)",
|
||||
"accounts": list(accounts.values()),
|
||||
"output": "\n".join(f"{row['name']} <{row['email']}> ({row['id']})" for row in accounts.values())
|
||||
or "No fixture email accounts found.",
|
||||
"exit_code": 0,
|
||||
}
|
||||
|
||||
@@ -121,14 +139,7 @@ def execute_fixture_email(tool: str, args: dict[str, Any], owner: str | None = N
|
||||
rows = [
|
||||
_row(i, m) for i, m in enumerate(messages)
|
||||
if _matches(m, query) and _folder_matches(m.get("folder"), folder)
|
||||
and (
|
||||
not account
|
||||
or account in {
|
||||
str(m.get("account") or "").strip().casefold(),
|
||||
str(m.get("account_email") or "").strip().casefold(),
|
||||
str(m.get("account_id") or "").strip().casefold(),
|
||||
}
|
||||
)
|
||||
and _account_matches(m, account)
|
||||
]
|
||||
limit = args.get("max_results", args.get("limit", 20))
|
||||
try:
|
||||
@@ -139,11 +150,7 @@ def execute_fixture_email(tool: str, args: dict[str, Any], owner: str | None = N
|
||||
return {"output": "No emails found.", "emails": [], "exit_code": 0}
|
||||
output_lines = [f"Found {len(rows)} email(s):", ""]
|
||||
for index, r in enumerate(rows, start=1):
|
||||
source = next(
|
||||
(m for i, m in enumerate(messages) if _uid(i) == r["uid"]),
|
||||
{},
|
||||
)
|
||||
summary = str(source.get("summary") or source.get("body") or "").strip()
|
||||
summary = r["summary"].strip()
|
||||
output_lines.extend([
|
||||
f"{index}. **{r['subject']}**",
|
||||
f" From: {r['from']}",
|
||||
@@ -159,15 +166,19 @@ def execute_fixture_email(tool: str, args: dict[str, Any], owner: str | None = N
|
||||
|
||||
if bare == "read_email":
|
||||
uid = str(args.get("uid") or "")
|
||||
message_id = str(args.get("message_id") or "")
|
||||
folder = str(args.get("folder") or "INBOX")
|
||||
try:
|
||||
index = int(uid) - 1
|
||||
except (TypeError, ValueError):
|
||||
index = -1
|
||||
if index < 0 or index >= len(messages) or not _folder_matches(messages[index].get("folder"), folder):
|
||||
return {"error": f"Email UID {uid} not found.", "exit_code": 1}
|
||||
message = messages[index]
|
||||
row = _row(index, message)
|
||||
if not uid and not message_id:
|
||||
return {"error": "read_email requires a UID or Message-ID.", "exit_code": 1}
|
||||
matches = [(m, _row(i, m)) for i, m in enumerate(messages)
|
||||
if _folder_matches(m.get("folder"), folder) and _account_matches(m, args.get("account"))
|
||||
and (not uid or str(m.get("uid") or _uid(i)) == uid)
|
||||
and (not message_id or str(m.get("message_id") or "") == message_id)]
|
||||
if not matches:
|
||||
return {"error": f"Email UID/Message-ID {uid or message_id} not found.", "exit_code": 1}
|
||||
if len(matches) != 1:
|
||||
return {"error": "Email identity is ambiguous; specify its account and folder.", "exit_code": 1}
|
||||
message, row = matches[0]
|
||||
output = (
|
||||
f"UID: {row['uid']}\nSubject: {row['subject']}\nFrom: {row['from']}\n"
|
||||
f"Date: {row['date']}\n\n{message.get('body') or ''}"
|
||||
@@ -185,6 +196,7 @@ def execute_fixture_email(tool: str, args: dict[str, Any], owner: str | None = N
|
||||
all_messages = _messages()
|
||||
owner_value = str(owner or "").strip()
|
||||
visible_index = -1
|
||||
matches = []
|
||||
for original in all_messages:
|
||||
if owner_value and original.get("owner") and original.get("owner") != owner_value:
|
||||
continue
|
||||
@@ -193,6 +205,13 @@ def execute_fixture_email(tool: str, args: dict[str, Any], owner: str | None = N
|
||||
continue
|
||||
if not _folder_matches(original.get("folder"), folder):
|
||||
continue
|
||||
if not _account_matches(original, args.get("account")):
|
||||
continue
|
||||
matches.append(original)
|
||||
if len(matches) > 1:
|
||||
return {"error": "Email identity is ambiguous; specify its account and folder.", "exit_code": 1}
|
||||
if matches:
|
||||
original = matches[0]
|
||||
if bare == "archive_email":
|
||||
original["folder"] = "Archive"
|
||||
action = "Archived"
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
"""Process-wide shutdown signalling for graceful restarts.
|
||||
|
||||
Server-Sent-Events / streaming responses never complete on their own, so
|
||||
uvicorn's graceful shutdown would wait for an open browser stream up to the
|
||||
full ``timeout_graceful_shutdown`` window before exiting. A SIGTERM/SIGINT
|
||||
handler flips this flag and tells uvicorn to exit; every streaming generator
|
||||
checks the flag on each loop and ends promptly, so a restart completes in
|
||||
seconds while ordinary in-flight requests still finish.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_SHUTTING_DOWN = threading.Event()
|
||||
|
||||
|
||||
def is_shutting_down() -> bool:
|
||||
"""True once a shutdown (SIGTERM/SIGINT) has been requested."""
|
||||
return _SHUTTING_DOWN.is_set()
|
||||
|
||||
|
||||
def request_shutdown() -> None:
|
||||
"""Mark the process as shutting down (idempotent)."""
|
||||
if not _SHUTTING_DOWN.is_set():
|
||||
_SHUTTING_DOWN.set()
|
||||
logger.info("shutdown requested; streaming responses will end promptly")
|
||||
|
||||
|
||||
def reset() -> None:
|
||||
"""Clear the flag (for tests)."""
|
||||
_SHUTTING_DOWN.clear()
|
||||
+13
-3
@@ -16,7 +16,7 @@ from typing import Dict, Optional
|
||||
from src.tools._common import _parse_tool_args
|
||||
|
||||
|
||||
async def do_edit_image(content: str, owner: Optional[str] = None) -> Dict:
|
||||
async def do_edit_image(content: str, owner: Optional[str] = None, *, incognito: bool = False) -> Dict:
|
||||
"""Create an owner-scoped edited copy of a gallery image."""
|
||||
try:
|
||||
args = _parse_tool_args(content)
|
||||
@@ -45,7 +45,7 @@ async def do_edit_image(content: str, owner: Optional[str] = None) -> Dict:
|
||||
if not load_settings().get('image_gen_enabled', True):
|
||||
return {'error': 'Image generation is disabled by the administrator.', 'exit_code': 1}
|
||||
from src.ai_interaction import do_edit_image as edit_with_model
|
||||
return await edit_with_model(str(args.get('prompt') or '').strip(), info['path'], owner=owner, size='auto')
|
||||
return await edit_with_model(str(args.get('prompt') or '').strip(), info['path'], owner=owner, size='auto', incognito=incognito)
|
||||
|
||||
from core.database import GalleryImage, SessionLocal
|
||||
from src.constants import GENERATED_IMAGES_DIR
|
||||
@@ -79,7 +79,7 @@ async def do_edit_image(content: str, owner: Optional[str] = None) -> Dict:
|
||||
session_id = source.session_id
|
||||
db.close()
|
||||
from src.ai_interaction import do_edit_image as edit_with_model
|
||||
return await edit_with_model(prompt, str(source_path), session_id=session_id, owner=owner, size='auto')
|
||||
return await edit_with_model(prompt, str(source_path), session_id=session_id, owner=owner, size='auto', incognito=incognito)
|
||||
|
||||
from PIL import Image
|
||||
|
||||
@@ -111,6 +111,16 @@ async def do_edit_image(content: str, owner: Optional[str] = None) -> Dict:
|
||||
output_bytes = output.getvalue()
|
||||
width, height = image.size
|
||||
|
||||
if incognito:
|
||||
import base64
|
||||
return {
|
||||
'output': f'Image edited ({action}). Private result, not saved to gallery.',
|
||||
'exit_code': 0, 'image_id': None,
|
||||
'image_url': 'data:image/png;base64,' + base64.b64encode(output_bytes).decode('ascii'),
|
||||
'image_prompt': source.prompt or action, 'image_model': f'edit_image:{action}',
|
||||
'image_size': f'{width}x{height}', 'image_quality': source.quality or '',
|
||||
}
|
||||
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
filename = f"{uuid.uuid4().hex[:12]}.png"
|
||||
(root / filename).write_bytes(output_bytes)
|
||||
|
||||
+9
-70
@@ -63,80 +63,19 @@ INTERNAL_UPLOAD_URL_RE = re.compile(
|
||||
r"([0-9a-fA-F]{32}(?:\.[A-Za-z0-9]+)?)"
|
||||
r"(?=$|[\s\"'<>\[\](){},;!?:&#]|\.(?![A-Za-z0-9]))"
|
||||
)
|
||||
# `<!--\s*pdf(?:_form)?_source\b[^>]*\bupload_id=["'](id)["'][^>]*-->` and
|
||||
# `\[Attachment:[^\]\r\n]*\|\s*id=(id)(?:\s*\||\s*\])` are matched by the
|
||||
# forward-only scanners below. As single regexes, every opener in a run with no
|
||||
# closing `>` / `]` rescanned to the end of that run: O(n^2) on chat content,
|
||||
# which is unbounded on persisted assistant output (CodeQL py/polynomial-redos).
|
||||
_PDF_SOURCE_OPEN_RE = re.compile(r"<!--\s*pdf(?:_form)?_source\b", re.IGNORECASE)
|
||||
_PDF_SOURCE_ID_RE = re.compile(
|
||||
r"\bupload_id=[\"']([0-9a-fA-F]{32}(?:\.[A-Za-z0-9]+)?)[\"']",
|
||||
PDF_SOURCE_UPLOAD_RE = re.compile(
|
||||
r"<!--\s*pdf(?:_form)?_source\b[^>]*\bupload_id="
|
||||
r"[\"']([0-9a-fA-F]{32}(?:\.[A-Za-z0-9]+)?)[\"'][^>]*-->",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_ATTACHMENT_REFERENCE_OPEN_RE = re.compile(r"\[Attachment:", re.IGNORECASE)
|
||||
_ATTACHMENT_REFERENCE_STOP_RE = re.compile(r"[\]\r\n]")
|
||||
_ATTACHMENT_REFERENCE_TAIL_RE = re.compile(
|
||||
r"\|\s*id=([0-9a-fA-F]{32}(?:\.[A-Za-z0-9]+)?)(?:\s*\||\s*\])",
|
||||
ATTACHMENT_REFERENCE_LINE_RE = re.compile(
|
||||
r"\[Attachment:[^\]\r\n]*\|\s*id="
|
||||
r"([0-9a-fA-F]{32}(?:\.[A-Za-z0-9]+)?)"
|
||||
r"(?:\s*\||\s*\])",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def _pdf_source_upload_ids(value: str) -> list[str]:
|
||||
"""IDs from `<!-- pdf_source ... upload_id="<id>" ... -->` comments.
|
||||
|
||||
Every opener before the next `>` shares that `>`: the comment matches only
|
||||
if `--` sits right before it, and then with the last `upload_id=` (the
|
||||
greedy `[^>]*` backtracks from the right). Otherwise all of those openers
|
||||
fail together, so the scan resumes after the `>`.
|
||||
"""
|
||||
found: list[str] = []
|
||||
pos = 0
|
||||
while True:
|
||||
opener = _PDF_SOURCE_OPEN_RE.search(value, pos)
|
||||
if opener is None:
|
||||
return found
|
||||
close = value.find(">", opener.end())
|
||||
if close < 0:
|
||||
return found
|
||||
if value[close - 2:close] == "--":
|
||||
last = None
|
||||
for last in _PDF_SOURCE_ID_RE.finditer(value, opener.end(), close):
|
||||
pass
|
||||
if last is not None:
|
||||
found.append(last.group(1))
|
||||
pos = close + 1
|
||||
|
||||
|
||||
def _attachment_reference_ids(value: str) -> list[str]:
|
||||
"""IDs from `[Attachment: name | id=<id> | ...]` reference lines.
|
||||
|
||||
The label scan stops at the first `]`/CR/LF, so every opener before that
|
||||
stop shares it and can only use the last `|` (scanning right to left)
|
||||
whose tail matches. If none does, all of those openers fail together and
|
||||
the scan resumes at the stop.
|
||||
"""
|
||||
found: list[str] = []
|
||||
pos = 0
|
||||
while True:
|
||||
opener = _ATTACHMENT_REFERENCE_OPEN_RE.search(value, pos)
|
||||
if opener is None:
|
||||
return found
|
||||
stop = _ATTACHMENT_REFERENCE_STOP_RE.search(value, opener.end())
|
||||
end = stop.start() if stop else len(value)
|
||||
tail = None
|
||||
pipe = value.rfind("|", opener.end(), end)
|
||||
while pipe >= 0:
|
||||
tail = _ATTACHMENT_REFERENCE_TAIL_RE.match(value, pipe)
|
||||
if tail is not None:
|
||||
break
|
||||
pipe = value.rfind("|", opener.end(), pipe)
|
||||
if tail is None:
|
||||
pos = end
|
||||
else:
|
||||
found.append(tail.group(1))
|
||||
pos = tail.end()
|
||||
|
||||
|
||||
def is_valid_upload_id(upload_id: str) -> bool:
|
||||
"""Return True when *upload_id* matches the canonical uploads.json id format."""
|
||||
return UPLOAD_ID_RE.fullmatch(upload_id or "") is not None
|
||||
@@ -171,8 +110,8 @@ def extract_internal_upload_ids(value: Any) -> set[str]:
|
||||
return set()
|
||||
return (
|
||||
set(INTERNAL_UPLOAD_URL_RE.findall(value))
|
||||
| set(_pdf_source_upload_ids(value))
|
||||
| set(_attachment_reference_ids(value))
|
||||
| set(PDF_SOURCE_UPLOAD_RE.findall(value))
|
||||
| set(ATTACHMENT_REFERENCE_LINE_RE.findall(value))
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
"""Tests for src/activity.py — in-flight heavy-image-op counters (safe restart)."""
|
||||
from src import activity
|
||||
|
||||
|
||||
def test_image_op_counter_increments_and_decrements():
|
||||
assert activity.heavy_image_ops_active() == 0
|
||||
activity.begin_image_op("a")
|
||||
activity.begin_image_op("a")
|
||||
assert activity.heavy_image_ops_active() == 2
|
||||
activity.end_image_op("a")
|
||||
assert activity.heavy_image_ops_active() == 1
|
||||
activity.end_image_op("a")
|
||||
assert activity.heavy_image_ops_active() == 0
|
||||
# never goes negative
|
||||
activity.end_image_op("a")
|
||||
assert activity.heavy_image_ops_active() == 0
|
||||
|
||||
|
||||
def test_image_op_context_manager_and_busy():
|
||||
assert activity.busy() is False
|
||||
with activity.image_op("/api/image/inpaint"):
|
||||
assert activity.busy() is True
|
||||
snap = activity.activity()
|
||||
assert snap["image_ops"] == 1
|
||||
assert snap["image_ops_by_name"].get("/api/image/inpaint") == 1
|
||||
assert activity.busy() is False
|
||||
assert activity.activity()["image_ops"] == 0
|
||||
|
||||
|
||||
def test_image_op_released_on_exception():
|
||||
try:
|
||||
with activity.image_op("boom"):
|
||||
raise RuntimeError("x")
|
||||
except RuntimeError:
|
||||
pass
|
||||
assert activity.heavy_image_ops_active() == 0
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,10 +1,9 @@
|
||||
"""Repository asset ownership guards for issues #1335 and #6175.
|
||||
|
||||
Public Markdown and landing-page media belong in website/, while shared
|
||||
optional README/packaging imagery belongs in assets/branding/. Retained images
|
||||
in either managed root must be referenced by tracked text, and every retained
|
||||
website video must be referenced by the site's entry point. Publication can
|
||||
operate with no imagery or videos.
|
||||
README/packaging imagery belongs in assets/branding/. Images in either managed
|
||||
root must be referenced by tracked text, and every tracked website video must
|
||||
be referenced by the site's entry point.
|
||||
"""
|
||||
import re
|
||||
import subprocess
|
||||
@@ -42,8 +41,7 @@ def _tracked(*paths_under):
|
||||
return None
|
||||
if out.returncode != 0:
|
||||
return None
|
||||
return [REPO / line for line in out.stdout.splitlines()
|
||||
if line.strip() and (REPO / line).is_file()]
|
||||
return [REPO / line for line in out.stdout.splitlines() if line.strip()]
|
||||
|
||||
|
||||
def test_no_orphan_documentation_or_branding_images():
|
||||
@@ -51,6 +49,9 @@ def test_no_orphan_documentation_or_branding_images():
|
||||
if managed_files is None:
|
||||
pytest.skip("not a git checkout")
|
||||
managed_images = [p for p in managed_files if p.suffix.lower() in IMAGE_EXTS]
|
||||
assert any("assets/branding" in p.as_posix() for p in managed_images), (
|
||||
"expected assets/branding/ to contain the shared project imagery"
|
||||
)
|
||||
|
||||
# All tracked text we might reference an image from.
|
||||
all_tracked = _tracked(".") or []
|
||||
@@ -93,6 +94,7 @@ def test_pages_site_owns_its_entrypoint_and_media():
|
||||
assert text.startswith("---\nlayout: default\n---\n"), guide
|
||||
|
||||
website_videos = [p for p in website_files if p.suffix.lower() in VIDEO_EXTS]
|
||||
assert website_videos, "expected website/ to contain the landing-page videos"
|
||||
|
||||
entrypoint = (REPO / "website/index.html").read_text(encoding="utf-8")
|
||||
unreferenced = [
|
||||
|
||||
@@ -2,6 +2,7 @@ import asyncio
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from src.tools.image import do_edit_image
|
||||
@@ -122,7 +123,8 @@ def test_edit_image_rejects_actions_without_complete_input_contract():
|
||||
assert "Use prompt, upscale or rembg" in result["error"]
|
||||
|
||||
|
||||
def test_prompt_edit_forwards_owned_pixels_and_preserves_original(monkeypatch, tmp_path):
|
||||
@pytest.mark.parametrize('incognito', [False, True])
|
||||
def test_prompt_edit_forwards_owned_pixels_and_preserves_original(monkeypatch, tmp_path, incognito):
|
||||
source_path = tmp_path / 'source.png'
|
||||
Image.new('RGB', (3, 2), 'red').save(source_path)
|
||||
original = source_path.read_bytes()
|
||||
@@ -135,9 +137,9 @@ def test_prompt_edit_forwards_owned_pixels_and_preserves_original(monkeypatch, t
|
||||
return {'image_id': 'edited', 'image_url': '/api/generated-image/edited.png'}
|
||||
from pathlib import Path
|
||||
monkeypatch.setattr('src.ai_interaction.do_edit_image', edit)
|
||||
result = asyncio.run(do_edit_image(json.dumps({'image_id': 'source-id', 'action': 'prompt', 'prompt': 'Add another cow'}), owner='alice'))
|
||||
result = asyncio.run(do_edit_image(json.dumps({'image_id': 'source-id', 'action': 'prompt', 'prompt': 'Add another cow'}), owner='alice', incognito=incognito))
|
||||
assert result['image_id'] == 'edited'
|
||||
assert calls == [('Add another cow', original, {'session_id': 'session-1', 'owner': 'alice', 'size': 'auto'})]
|
||||
assert calls == [('Add another cow', original, {'session_id': 'session-1', 'owner': 'alice', 'size': 'auto', 'incognito': incognito})]
|
||||
assert source_path.read_bytes() == original
|
||||
denied = asyncio.run(do_edit_image('{"image_id":"source-id","action":"prompt","prompt":"edit"}', owner=None))
|
||||
assert denied['error'] == 'Image not found'
|
||||
@@ -165,7 +167,7 @@ def test_uploaded_image_edit_uses_owner_scoped_reference(monkeypatch, tmp_path):
|
||||
monkeypatch.setattr('src.ai_interaction.do_edit_image', edit)
|
||||
args = '{"image_id":"odysseus://attachment/upload.jpg","action":"prompt","prompt":"Make more realistic"}'
|
||||
assert asyncio.run(do_edit_image(args, owner='alice')) == {'image_id': 'edited'}
|
||||
assert calls == [('Make more realistic', str(path), {'owner': 'alice', 'size': 'auto'})]
|
||||
assert calls == [('Make more realistic', str(path), {'owner': 'alice', 'size': 'auto', 'incognito': False})]
|
||||
assert resolutions == [('upload.jpg', {'owner': 'alice', 'allow_admin': False})]
|
||||
assert 'error' in asyncio.run(do_edit_image(args, owner='bob'))
|
||||
assert len(calls) == 1
|
||||
@@ -173,3 +175,23 @@ def test_uploaded_image_edit_uses_owner_scoped_reference(monkeypatch, tmp_path):
|
||||
disabled = asyncio.run(do_edit_image(args, owner='alice'))
|
||||
assert 'disabled' in disabled['error']
|
||||
assert len(calls) == 1
|
||||
|
||||
|
||||
def test_private_upscale_does_not_save_image_or_gallery_record(monkeypatch, tmp_path):
|
||||
source_path = tmp_path / 'source.png'
|
||||
Image.new('RGB', (3, 2), 'red').save(source_path)
|
||||
original = source_path.read_bytes()
|
||||
db = _Db(_source(source_path.name))
|
||||
monkeypatch.setattr('core.database.SessionLocal', lambda: db)
|
||||
monkeypatch.setattr('src.constants.GENERATED_IMAGES_DIR', str(tmp_path))
|
||||
result = asyncio.run(do_edit_image(
|
||||
'{"image_id":"source-id","action":"upscale","scale":2}',
|
||||
owner='alice', incognito=True,
|
||||
))
|
||||
assert result['exit_code'] == 0
|
||||
assert result['image_url'].startswith('data:image/png;base64,')
|
||||
assert not result['image_id']
|
||||
assert db.added == []
|
||||
assert not db.committed
|
||||
assert list(tmp_path.iterdir()) == [source_path]
|
||||
assert source_path.read_bytes() == original
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
"""Server health audit fix: the local editor VLM returned 400s when asked for a
|
||||
model id it does not serve (the client default 'Qwen3.5-9B-MLX-8bit', or the
|
||||
stale caption model 'Qwen/Qwen3-VL-30B'). ``_resolve_model`` maps the requested
|
||||
id to one the base actually serves."""
|
||||
import asyncio
|
||||
|
||||
from src import editor_agent_model as eam
|
||||
|
||||
|
||||
def _cache(base, ids):
|
||||
eam._MODELS_CACHE[base] = set(ids)
|
||||
|
||||
|
||||
def test_resolve_model_keeps_a_served_request():
|
||||
base = "http://x:1/v1"
|
||||
_cache(base, ["Qwen/Qwen3.5-9B", "other"])
|
||||
assert asyncio.run(eam._resolve_model(base, "Qwen/Qwen3.5-9B", "Qwen/Qwen3.5-9B")) == "Qwen/Qwen3.5-9B"
|
||||
|
||||
|
||||
def test_resolve_model_falls_back_when_requested_is_not_served():
|
||||
base = "http://x:2/v1"
|
||||
_cache(base, ["Qwen/Qwen3.5-9B"])
|
||||
# e.g. the client's MLX build name, or a stale caption model id.
|
||||
assert asyncio.run(eam._resolve_model(base, "Qwen3.5-9B-MLX-8bit", "Qwen/Qwen3.5-9B")) == "Qwen/Qwen3.5-9B"
|
||||
assert asyncio.run(eam._resolve_model(base, "Qwen/Qwen3-VL-30B", "Qwen/Qwen3.5-9B")) == "Qwen/Qwen3.5-9B"
|
||||
|
||||
|
||||
def test_resolve_model_picks_a_served_id_when_fallback_absent():
|
||||
base = "http://x:3/v1"
|
||||
_cache(base, ["B", "A"])
|
||||
assert asyncio.run(eam._resolve_model(base, "nope", "also-nope")) == "A"
|
||||
|
||||
|
||||
def test_resolve_model_keeps_request_when_served_list_unknown():
|
||||
base = "http://x:4/v1"
|
||||
_cache(base, [])
|
||||
# Server unreachable / no models listed -> don't rewrite the request.
|
||||
assert asyncio.run(eam._resolve_model(base, "some/model", "fallback")) == "some/model"
|
||||
|
||||
|
||||
def test_self_check_resolves_the_model_before_calling(monkeypatch):
|
||||
base = eam.local_bases()[0]
|
||||
_cache(base, ["Qwen/Qwen3.5-9B"])
|
||||
seen = {}
|
||||
|
||||
async def fake_post(b, model, messages, headers, timeout, **kwargs):
|
||||
seen["model"] = model
|
||||
return "DONE"
|
||||
|
||||
monkeypatch.setattr(eam, "_post_chat", fake_post)
|
||||
out = asyncio.run(eam.self_check(
|
||||
"data:image/png;base64,AAAA", "data:image/png;base64,BBBB",
|
||||
"make it brighter", model="Qwen3.5-9B-MLX-8bit"))
|
||||
assert seen["model"] == "Qwen/Qwen3.5-9B"
|
||||
assert out["done"] is True
|
||||
|
||||
|
||||
def test_caption_resolves_the_caption_model(monkeypatch):
|
||||
seen = {}
|
||||
|
||||
async def fake_post(b, model, messages, headers, timeout, **kwargs):
|
||||
seen["base"] = b
|
||||
seen["model"] = model
|
||||
return "a photo"
|
||||
|
||||
monkeypatch.setattr(eam, "_post_chat", fake_post)
|
||||
import src.editor_agent_model as m
|
||||
_cache(m.caption_base(), ["Qwen/Qwen3.5-9B"])
|
||||
txt = asyncio.run(m.caption_image("data:image/png;base64,AAAA"))
|
||||
assert txt == "a photo"
|
||||
assert seen["model"] == "Qwen/Qwen3.5-9B"
|
||||
@@ -0,0 +1,171 @@
|
||||
"""Planner routing regression suite (HARNESS-owned).
|
||||
|
||||
Turns PHOTOSHOP's bench prompt families (vague, restyle, bg-remove/scene,
|
||||
tone/auto, generative, structural) into assertions on the ROUTE CLASS the
|
||||
server planner picks, so a future planner change cannot silently re-break
|
||||
routing. The classes are:
|
||||
|
||||
ask_user -> the planner is unsure; it must ask with 2-3 choices and
|
||||
apply NO operations (never a silent tone default)
|
||||
deterministic -> a deterministic editor op / recipe (adjustments, filters,
|
||||
crop, text, rotate, fill, recolour ...)
|
||||
ai_edit_image -> a generative image-model edit (content add/remove/replace,
|
||||
restyle, look)
|
||||
generate_image -> text-to-image for a brand-new image
|
||||
|
||||
These use the pure routing helpers (``_plan_special`` / ``_generative_kind`` /
|
||||
``_fallback_plan``), so no model or network call is made.
|
||||
"""
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from routes import editor_agent_routes as eaar
|
||||
|
||||
_STATE = {"width": 900, "height": 600}
|
||||
_PROVIDER = True
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_planner():
|
||||
eaar.set_editor_planner(None)
|
||||
yield
|
||||
eaar.set_editor_planner(None)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def client():
|
||||
app = FastAPI()
|
||||
app.include_router(eaar.setup_editor_agent_routes())
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def _classify(prompt, provider=_PROVIDER):
|
||||
"""Return (route_class, ops) for a prompt using the deterministic helpers."""
|
||||
special = eaar._plan_special(prompt, _STATE, provider)
|
||||
if special:
|
||||
steps = list(special[0])
|
||||
else:
|
||||
fallback = eaar._fallback_plan(prompt)
|
||||
if fallback.get("ask_user"):
|
||||
return "ask_user", []
|
||||
steps = list(fallback.get("steps") or [])
|
||||
ops = [str(s.get("op") or "") for s in steps]
|
||||
if "generate_image" in ops:
|
||||
return "generate_image", ops
|
||||
if "ai_edit_image" in ops:
|
||||
return "ai_edit_image", ops
|
||||
return "deterministic", ops
|
||||
|
||||
|
||||
# (prompt, expected route class)
|
||||
BENCH_ROUTING = [
|
||||
# vague -> ask (never a silent tone tweak)
|
||||
("do the thing", "ask_user"),
|
||||
("give it some oomph", "ask_user"),
|
||||
("make it better", "ask_user"),
|
||||
# whole-image restyle -> model edit
|
||||
("make it look like an oil painting", "ai_edit_image"),
|
||||
("turn it into an anime", "ai_edit_image"),
|
||||
("make it a comic book illustration", "ai_edit_image"),
|
||||
("make it look like cyberpunk", "ai_edit_image"),
|
||||
# background / scene replacement -> model edit
|
||||
("put her on a beach", "ai_edit_image"),
|
||||
("make the background nicer", "ai_edit_image"),
|
||||
("make it snowy", "ai_edit_image"),
|
||||
# generative content add / remove -> model edit
|
||||
("add a full moon in the sky", "ai_edit_image"),
|
||||
("remove the distractions", "ai_edit_image"),
|
||||
("remove the person on the left", "ai_edit_image"),
|
||||
# brand-new image -> text-to-image
|
||||
("create a new image of a full moon", "generate_image"),
|
||||
("generate an image of a sunset", "generate_image"),
|
||||
("make a logo", "generate_image"),
|
||||
# tone / auto / look / structural -> deterministic
|
||||
("make it black and white", "deterministic"),
|
||||
("make it pop", "deterministic"),
|
||||
("make it look professional", "deterministic"),
|
||||
("make only the red car blue", "deterministic"),
|
||||
("make a 1080x1350 Instagram post from this", "deterministic"),
|
||||
("add a title TOKYO with a drop shadow", "deterministic"),
|
||||
("extend the canvas to 16:9", "deterministic"),
|
||||
("save as PSD with layers", "deterministic"),
|
||||
("fill the white area with red", "deterministic"),
|
||||
("rotate the whole image 90 degrees", "deterministic"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prompt,expected", BENCH_ROUTING)
|
||||
def test_bench_prompts_route_to_expected_class(prompt, expected):
|
||||
route_class, ops = _classify(prompt)
|
||||
assert route_class == expected, f"{prompt!r} -> {route_class} {ops} (want {expected})"
|
||||
|
||||
|
||||
def test_ask_plans_offer_two_to_three_choices():
|
||||
for prompt in ("do the thing", "give it some oomph", "make it better"):
|
||||
plan = eaar._fallback_plan(prompt)
|
||||
ask = plan.get("ask_user")
|
||||
assert isinstance(ask, dict), prompt
|
||||
opts = ask.get("options") or []
|
||||
assert 2 <= len(opts) <= 3, (prompt, opts)
|
||||
assert plan.get("steps") == []
|
||||
|
||||
|
||||
def test_object_removal_uses_selection_plus_model_edit():
|
||||
_cls, ops = _classify("remove the person on the left")
|
||||
assert "ai_edit_image" in ops
|
||||
assert any(op.startswith("selection.") for op in ops), ops
|
||||
|
||||
|
||||
def test_recolour_never_uses_a_flat_fill_or_a_global_tint():
|
||||
# Every recolour selects the object by description and either shifts the
|
||||
# masked hue or asks the image model for the masked region — never a flat
|
||||
# selection.fill and never a whole-image adjustments.apply.
|
||||
for prompt in ("make only the red car blue", "make the red car white", "make the flowers bright pink"):
|
||||
_cls, ops = _classify(prompt)
|
||||
assert "selection.samText" in ops, (prompt, ops)
|
||||
assert "selection.fill" not in ops, (prompt, ops)
|
||||
assert "adjustments.apply" not in ops, (prompt, ops)
|
||||
|
||||
|
||||
def test_chromatic_recolour_is_a_deterministic_masked_hue_shift():
|
||||
cls, ops = _classify("make only the red car blue")
|
||||
assert cls == "deterministic", ops
|
||||
assert ops == ["selection.samText", "adjustments.applyMasked"], ops
|
||||
|
||||
|
||||
def test_new_image_request_from_scratch_vs_edit_of_current_photo():
|
||||
# A photo edit ("make it a <style>") must NOT be treated as a new image.
|
||||
assert eaar._generative_kind("make it a comic book illustration") == "restyle"
|
||||
assert eaar._is_new_image_request("make it a comic book illustration") is False
|
||||
# A from-scratch generation still is.
|
||||
assert eaar._is_new_image_request("create an illustration of a cat") is True
|
||||
assert eaar._generative_kind("create an illustration of a cat") == "add"
|
||||
|
||||
|
||||
def test_plan_route_returns_deterministic_recolour(client):
|
||||
eaar.set_editor_planner(lambda messages, owner: json.dumps({"reply": "ok", "steps": []}))
|
||||
r = client.post("/api/editor/plan", json={"prompt": "make only the red car blue", "state": _STATE})
|
||||
assert r.status_code == 200
|
||||
data = r.json()
|
||||
assert data["planner"] == "deterministic"
|
||||
assert [s["op"] for s in data["steps"]] == ["selection.samText", "adjustments.applyMasked"]
|
||||
|
||||
|
||||
def test_plan_route_asks_for_a_vague_prompt(client):
|
||||
eaar.set_editor_planner(lambda messages, owner: json.dumps({"reply": "ok", "steps": []}))
|
||||
r = client.post("/api/editor/plan", json={"prompt": "do the thing", "state": _STATE})
|
||||
assert r.status_code == 200
|
||||
data = r.json()
|
||||
assert data.get("ask_user"), data
|
||||
assert data.get("steps") == []
|
||||
|
||||
|
||||
def test_plan_route_routes_restyle_to_model_edit(client):
|
||||
eaar.set_editor_planner(lambda messages, owner: json.dumps({"reply": "ok", "steps": []}))
|
||||
r = client.post("/api/editor/plan", json={"prompt": "turn it into an anime", "state": _STATE})
|
||||
assert r.status_code == 200
|
||||
ops = [s["op"] for s in r.json()["steps"]]
|
||||
assert "ai_edit_image" in ops, ops
|
||||
@@ -11,6 +11,7 @@ from src.clean_agent_preview import (
|
||||
from src.turn_contract import (
|
||||
editor_request_instructions, requested_capabilities, selected_tools_for_request,
|
||||
preserve_bound_editor_selected_tools, resolve_turn_contract,
|
||||
targets_bound_editor_request,
|
||||
)
|
||||
from src.tool_schemas import FUNCTION_TOOL_SCHEMAS
|
||||
from src.tool_policy import ToolPolicy
|
||||
@@ -72,6 +73,16 @@ def test_direct_edits_share_the_capability_router(prompt):
|
||||
assert not active_editor_suggestion_request(doc, prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('prompt', [
|
||||
'ga mer ut pa att vi tanker pa framtid for bjorn och att vi faktist trivs sa bra har',
|
||||
'Go more into how we are thinking about the future.',
|
||||
])
|
||||
def test_implicit_revision_followups_route_to_active_document_only(prompt):
|
||||
assert requested_capabilities(prompt, active_document=True) == {'documents'}
|
||||
assert requested_capabilities(prompt, active_document=False) == frozenset()
|
||||
assert targets_bound_editor_request(prompt)
|
||||
|
||||
|
||||
def test_source_review_still_requires_tool_completion_after_research():
|
||||
schemas = [s for s in FUNCTION_TOOL_SCHEMAS if s['function']['name'] in {'web_search', 'suggest_document'}]
|
||||
assert required_active_editor_tool_choice(active_editor_target=True, suggestion_target=True,
|
||||
|
||||
@@ -26,7 +26,7 @@ class _FakeSock:
|
||||
class _FakeIMAP:
|
||||
calls = []
|
||||
|
||||
def __init__(self, host, port, timeout=None, **_policy):
|
||||
def __init__(self, host, port, timeout=None):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.timeout = timeout
|
||||
@@ -60,8 +60,8 @@ def test_open_imap_connection_uses_shared_timeout_for_implicit_ssl(monkeypatch):
|
||||
import routes.email_helpers as helpers
|
||||
|
||||
_FakeIMAP.calls = []
|
||||
monkeypatch.setattr(helpers, "_PolicyIMAP4", _FakeIMAP)
|
||||
monkeypatch.setattr(helpers, "_PolicyIMAP4_SSL", _FakeIMAPSSL)
|
||||
monkeypatch.setattr(helpers.imaplib, "IMAP4", _FakeIMAP)
|
||||
monkeypatch.setattr(helpers.imaplib, "IMAP4_SSL", _FakeIMAPSSL)
|
||||
|
||||
conn = _open_imap_connection("imap.one.com", 993, starttls=False)
|
||||
|
||||
@@ -75,8 +75,8 @@ def test_open_imap_connection_supports_starttls(monkeypatch):
|
||||
import routes.email_helpers as helpers
|
||||
|
||||
_FakeIMAP.calls = []
|
||||
monkeypatch.setattr(helpers, "_PolicyIMAP4", _FakeIMAP)
|
||||
monkeypatch.setattr(helpers, "_PolicyIMAP4_SSL", _FakeIMAPSSL)
|
||||
monkeypatch.setattr(helpers.imaplib, "IMAP4", _FakeIMAP)
|
||||
monkeypatch.setattr(helpers.imaplib, "IMAP4_SSL", _FakeIMAPSSL)
|
||||
|
||||
_open_imap_connection("imap.local", 143, starttls=True)
|
||||
|
||||
@@ -99,7 +99,7 @@ async def test_account_config_uses_shared_imap_timeout(monkeypatch):
|
||||
def logout(self):
|
||||
captured["logout"] = True
|
||||
|
||||
def fake_open(host, port, *, starttls, timeout, **_policy):
|
||||
def fake_open(host, port, *, starttls, timeout):
|
||||
captured["open"] = (host, port, starttls, timeout)
|
||||
return _Conn()
|
||||
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
"""Email list stale-while-revalidate (SWR) regression tests.
|
||||
|
||||
A plain list view must never block ~2-5 s on IMAP while we already hold a
|
||||
recent snapshot or a durable index: it should return the snapshot immediately
|
||||
(with ``sync.stale``) and refresh in the background. Blocking IMAP only remains
|
||||
for a genuine cold start (no snapshot AND no indexed rows).
|
||||
"""
|
||||
import asyncio
|
||||
import pathlib
|
||||
|
||||
from routes.email import email_routes as er
|
||||
|
||||
SRC = pathlib.Path(er.__file__).read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def _list_endpoint():
|
||||
router = er.setup_email_routes()
|
||||
for route in router.routes:
|
||||
if getattr(route, "name", "") == "list_emails":
|
||||
return route.endpoint
|
||||
raise AssertionError("list_emails route not found")
|
||||
|
||||
|
||||
def _call(endpoint, **over):
|
||||
kwargs = dict(
|
||||
folder="INBOX", limit=50, offset=0, filter="all", from_addr=None,
|
||||
account_id=None, has_attachments=0, date_from=None, date_to=None,
|
||||
cached_only=0, cache_bust=None, refresh=0, owner="pewds",
|
||||
)
|
||||
kwargs.update(over)
|
||||
return asyncio.run(endpoint(**kwargs))
|
||||
|
||||
|
||||
def test_list_returns_index_snapshot_with_a_stale_marker_without_blocking_on_imap(monkeypatch):
|
||||
endpoint = _list_endpoint()
|
||||
monkeypatch.setattr(
|
||||
er, "_email_index_list",
|
||||
lambda owner, account_id, folder, filter_, limit, offset, has_att: (
|
||||
[{"uid": "11", "subject": "Indexed", "is_read": False}], 1, "2026-10-11T00:00:00Z",
|
||||
),
|
||||
)
|
||||
# The synchronous path must never open IMAP; if it did the request would
|
||||
# raise instead of returning the index snapshot immediately.
|
||||
monkeypatch.setattr(er, "_imap_connect", lambda *a, **k: (_ for _ in ()).throw(
|
||||
AssertionError("SWR must not open IMAP when an index snapshot exists")))
|
||||
|
||||
result = _call(endpoint, cache_bust="123")
|
||||
assert result["total"] == 1
|
||||
assert [e["uid"] for e in result["emails"]] == ["11"]
|
||||
assert result["sync"]["stale"] is True
|
||||
assert result["sync"]["source"] in ("index", "cache_stale")
|
||||
|
||||
|
||||
def test_cold_start_with_no_snapshot_and_no_index_is_not_marked_stale(monkeypatch):
|
||||
# With neither a snapshot nor an index, SWR must NOT fire: the request
|
||||
# falls through to the real (blocking) path, so the response is not a
|
||||
# stale-marked snapshot.
|
||||
endpoint = _list_endpoint()
|
||||
monkeypatch.setattr(er, "_email_index_list", lambda *a, **k: ([], 0, None))
|
||||
monkeypatch.setattr(er, "_imap_connect", lambda *a, **k: (_ for _ in ()).throw(
|
||||
RuntimeError("offline test")))
|
||||
try:
|
||||
result = _call(endpoint)
|
||||
except Exception:
|
||||
return # the cold path legitimately attempted real IMAP
|
||||
assert isinstance(result, dict)
|
||||
assert not (result.get("sync") or {}).get("stale")
|
||||
|
||||
|
||||
def test_swr_contract_is_wired_into_the_list_path():
|
||||
assert "_LIST_STALE" in SRC
|
||||
assert "def _list_stale_put" in SRC
|
||||
assert "def _schedule_list_refresh" in SRC
|
||||
assert 'sync_meta["source"] = "cache_stale"' in SRC
|
||||
assert 'sync_meta["stale"] = True' in SRC
|
||||
# The manual-refresh contract the UI relies on must stay intact.
|
||||
assert "manual_refresh = bool(refresh)" in SRC
|
||||
assert "_invalidate_list_cache(account_id, folder)" in SRC
|
||||
assert "cached_rows = {} if refresh else _email_index_rows" in SRC
|
||||
@@ -64,7 +64,7 @@ def test_imap_connect_with_host_still_connects(monkeypatch):
|
||||
def login(self, user, password):
|
||||
opened["login"] = (user, password)
|
||||
|
||||
def _fake_open(host, port, *, starttls, timeout, **_policy):
|
||||
def _fake_open(host, port, *, starttls, timeout):
|
||||
opened["host"] = host
|
||||
return _FakeConn()
|
||||
|
||||
|
||||
@@ -12,7 +12,7 @@ from routes.email_helpers import _send_smtp_message
|
||||
class _FakeSMTP:
|
||||
calls = []
|
||||
|
||||
def __init__(self, host, port, timeout=None, **_policy):
|
||||
def __init__(self, host, port, timeout=None):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.timeout = timeout
|
||||
@@ -54,8 +54,8 @@ def test_send_smtp_message_supports_plain_smtp(monkeypatch):
|
||||
import routes.email_helpers as helpers
|
||||
|
||||
_FakeSMTP.calls = []
|
||||
monkeypatch.setattr(helpers, "_PolicySMTP", _FakeSMTP)
|
||||
monkeypatch.setattr(helpers, "_PolicySMTP_SSL", _FakeSMTPSSL)
|
||||
monkeypatch.setattr(helpers.smtplib, "SMTP", _FakeSMTP)
|
||||
monkeypatch.setattr(helpers.smtplib, "SMTP_SSL", _FakeSMTPSSL)
|
||||
|
||||
_send_smtp_message(_cfg("none"), "from@example.com", ["to@example.com"], "hello")
|
||||
|
||||
@@ -68,8 +68,8 @@ def test_send_smtp_message_supports_explicit_starttls(monkeypatch):
|
||||
import routes.email_helpers as helpers
|
||||
|
||||
_FakeSMTP.calls = []
|
||||
monkeypatch.setattr(helpers, "_PolicySMTP", _FakeSMTP)
|
||||
monkeypatch.setattr(helpers, "_PolicySMTP_SSL", _FakeSMTPSSL)
|
||||
monkeypatch.setattr(helpers.smtplib, "SMTP", _FakeSMTP)
|
||||
monkeypatch.setattr(helpers.smtplib, "SMTP_SSL", _FakeSMTPSSL)
|
||||
|
||||
_send_smtp_message(_cfg("starttls", port=2525), "from@example.com", ["to@example.com"], "hello")
|
||||
|
||||
@@ -82,8 +82,8 @@ def test_send_smtp_message_defaults_587_to_starttls(monkeypatch):
|
||||
import routes.email_helpers as helpers
|
||||
|
||||
_FakeSMTP.calls = []
|
||||
monkeypatch.setattr(helpers, "_PolicySMTP", _FakeSMTP)
|
||||
monkeypatch.setattr(helpers, "_PolicySMTP_SSL", _FakeSMTPSSL)
|
||||
monkeypatch.setattr(helpers.smtplib, "SMTP", _FakeSMTP)
|
||||
monkeypatch.setattr(helpers.smtplib, "SMTP_SSL", _FakeSMTPSSL)
|
||||
|
||||
cfg = _cfg("", port=587)
|
||||
_send_smtp_message(cfg, "from@example.com", ["to@example.com"], "hello")
|
||||
@@ -96,8 +96,8 @@ def test_send_smtp_message_uses_ssl_when_configured(monkeypatch):
|
||||
import routes.email_helpers as helpers
|
||||
|
||||
_FakeSMTP.calls = []
|
||||
monkeypatch.setattr(helpers, "_PolicySMTP", _FakeSMTP)
|
||||
monkeypatch.setattr(helpers, "_PolicySMTP_SSL", _FakeSMTPSSL)
|
||||
monkeypatch.setattr(helpers.smtplib, "SMTP", _FakeSMTP)
|
||||
monkeypatch.setattr(helpers.smtplib, "SMTP_SSL", _FakeSMTPSSL)
|
||||
|
||||
_send_smtp_message(_cfg("ssl", port=465), "from@example.com", ["to@example.com"], "hello")
|
||||
|
||||
|
||||
@@ -89,8 +89,8 @@ async def test_test_connection_oauth_account_uses_xoauth2_for_imap_and_smtp():
|
||||
|
||||
with mock.patch("core.database.SessionLocal", Factory), \
|
||||
mock.patch("routes.email_routes._open_imap_connection", return_value=mock_imap_conn), \
|
||||
mock.patch("routes.email_routes._PolicySMTP", return_value=mock_smtp_conn), \
|
||||
mock.patch("routes.email_routes._PolicySMTP_SSL", return_value=mock_smtp_conn), \
|
||||
mock.patch("routes.email_routes.smtplib.SMTP", return_value=mock_smtp_conn), \
|
||||
mock.patch("routes.email_routes.smtplib.SMTP_SSL", return_value=mock_smtp_conn), \
|
||||
mock.patch("routes.email_routes._get_valid_google_token", return_value="ya29.live") as token_getter:
|
||||
result = await test_conn(req=_FakeReq(), owner="alice")
|
||||
|
||||
@@ -141,8 +141,8 @@ async def test_test_connection_password_account_still_uses_login():
|
||||
|
||||
with mock.patch("core.database.SessionLocal", Factory), \
|
||||
mock.patch("routes.email_routes._open_imap_connection", return_value=mock_imap_conn), \
|
||||
mock.patch("routes.email_routes._PolicySMTP", return_value=mock_smtp_conn), \
|
||||
mock.patch("routes.email_routes._PolicySMTP_SSL", return_value=mock_smtp_conn):
|
||||
mock.patch("routes.email_routes.smtplib.SMTP", return_value=mock_smtp_conn), \
|
||||
mock.patch("routes.email_routes.smtplib.SMTP_SSL", return_value=mock_smtp_conn):
|
||||
result = await test_conn(req=_FakeReq(), owner="alice")
|
||||
|
||||
assert result["ok"] is True
|
||||
@@ -187,8 +187,8 @@ async def test_test_connection_rejects_non_google_hosts_before_oauth_auth():
|
||||
|
||||
with mock.patch("core.database.SessionLocal", Factory), \
|
||||
mock.patch("routes.email_routes._open_imap_connection") as open_imap, \
|
||||
mock.patch("routes.email_routes._PolicySMTP") as open_smtp, \
|
||||
mock.patch("routes.email_routes._PolicySMTP_SSL") as open_smtp_ssl, \
|
||||
mock.patch("routes.email_routes.smtplib.SMTP") as open_smtp, \
|
||||
mock.patch("routes.email_routes.smtplib.SMTP_SSL") as open_smtp_ssl, \
|
||||
mock.patch("routes.email_routes._get_valid_google_token") as token_getter:
|
||||
result = await test_conn(req=_FakeReq(), owner="alice")
|
||||
|
||||
@@ -238,8 +238,8 @@ async def test_test_connection_rejects_insecure_oauth_transports_before_auth():
|
||||
|
||||
with mock.patch("core.database.SessionLocal", Factory), \
|
||||
mock.patch("routes.email_routes._open_imap_connection") as open_imap, \
|
||||
mock.patch("routes.email_routes._PolicySMTP") as open_smtp, \
|
||||
mock.patch("routes.email_routes._PolicySMTP_SSL") as open_smtp_ssl, \
|
||||
mock.patch("routes.email_routes.smtplib.SMTP") as open_smtp, \
|
||||
mock.patch("routes.email_routes.smtplib.SMTP_SSL") as open_smtp_ssl, \
|
||||
mock.patch("routes.email_routes._get_valid_google_token") as token_getter:
|
||||
result = await test_conn(req=_FakeReq(), owner="alice")
|
||||
|
||||
@@ -336,10 +336,10 @@ async def test_test_connection_verifies_imap_tls_before_loading_oauth_token(
|
||||
"routes.email_routes.ssl.create_default_context",
|
||||
return_value=context,
|
||||
), mock.patch(
|
||||
"routes.email_helpers._PolicyIMAP4",
|
||||
"routes.email_helpers.imaplib.IMAP4",
|
||||
return_value=starttls_conn,
|
||||
) as imap_cls, mock.patch(
|
||||
"routes.email_helpers._PolicyIMAP4_SSL",
|
||||
"routes.email_helpers.imaplib.IMAP4_SSL",
|
||||
side_effect=ssl.SSLCertVerificationError("untrusted certificate"),
|
||||
) as imap_ssl_cls, mock.patch(
|
||||
"routes.email_routes._get_valid_google_token"
|
||||
@@ -412,10 +412,10 @@ async def test_test_connection_verifies_smtp_tls_before_loading_oauth_token(
|
||||
"routes.email_routes.ssl.create_default_context",
|
||||
return_value=context,
|
||||
), mock.patch(
|
||||
"routes.email_routes._PolicySMTP",
|
||||
"routes.email_routes.smtplib.SMTP",
|
||||
return_value=starttls_smtp,
|
||||
) as smtp_cls, mock.patch(
|
||||
"routes.email_routes._PolicySMTP_SSL",
|
||||
"routes.email_routes.smtplib.SMTP_SSL",
|
||||
side_effect=ssl.SSLCertVerificationError("untrusted certificate"),
|
||||
) as smtp_ssl_cls, mock.patch(
|
||||
"routes.email_routes._get_valid_google_token"
|
||||
|
||||
@@ -76,3 +76,158 @@ def test_mixed_endpoint_requires_configured_image_model(monkeypatch):
|
||||
assert lookup(Query(), endpoint.base_url, 'alice', 'text-model') is None
|
||||
assert lookup(Query(), endpoint.base_url, 'alice', 'openai/gpt-image-1') is None
|
||||
assert lookup(Query(), 'https://unregistered.test/v1', 'alice', 'openai/gpt-5-image') is None
|
||||
|
||||
|
||||
def _form_fields(content, content_type):
|
||||
"""Parse simple multipart form fields (text parts only) from a request body."""
|
||||
boundary = content_type.split('boundary=')[1].strip().strip('"')
|
||||
fields = {}
|
||||
for part in content.split(('--' + boundary).encode()):
|
||||
if b'name="' not in part:
|
||||
continue
|
||||
head, _, body = part.partition(b'\r\n\r\n')
|
||||
import re as _re
|
||||
match = _re.search(rb'name="([^"]+)"', head)
|
||||
if not match:
|
||||
continue
|
||||
fields[match.group(1).decode()] = body.rstrip(b'\r\n-').decode('latin1')
|
||||
return fields
|
||||
|
||||
|
||||
def _self_hosted_inpaint(monkeypatch, image, mask, body_extra, respond, include_mask=True):
|
||||
base_url = 'http://100.69.120.65:8188/v1'
|
||||
endpoint = SimpleNamespace(base_url=base_url, api_key='local', owner='alice')
|
||||
monkeypatch.setattr(gallery, 'SessionLocal', lambda: SimpleNamespace(close=lambda: None))
|
||||
monkeypatch.setattr(gallery, 'require_privilege', lambda *a: 'alice')
|
||||
monkeypatch.setattr(gallery, '_visible_image_endpoint_for_base', lambda *a: endpoint)
|
||||
monkeypatch.setattr('src.url_safety.check_outbound_url', lambda *a, **k: (True, ''))
|
||||
_real_client = httpx.AsyncClient
|
||||
monkeypatch.setattr(httpx, 'AsyncClient',
|
||||
lambda **kw: _real_client(transport=httpx.MockTransport(respond), **kw))
|
||||
|
||||
class Request:
|
||||
async def json(self):
|
||||
payload = {'_endpoint': base_url, 'image': png(image), 'prompt': 'do a thing'}
|
||||
if include_mask:
|
||||
payload['mask'] = png(mask)
|
||||
payload.update(body_extra)
|
||||
return payload
|
||||
|
||||
route = next(r for r in gallery.setup_gallery_routes().routes if r.path == '/api/image/inpaint')
|
||||
return asyncio.run(route.endpoint(Request()))
|
||||
|
||||
|
||||
def test_self_hosted_inpaint_caps_oversize_and_restores_dimensions(monkeypatch):
|
||||
seen = {}
|
||||
|
||||
def respond(request):
|
||||
assert str(request.url) == 'http://100.69.120.65:8188/v1/images/edits'
|
||||
seen.update(_form_fields(request.content, request.headers['content-type']))
|
||||
return httpx.Response(200, json={'data': [{'b64_json': png(Image.new('RGB', (768, 576), 'blue'))}]})
|
||||
|
||||
out = _self_hosted_inpaint(monkeypatch, Image.new('RGB', (1600, 1200), 'red'),
|
||||
Image.new('L', (1600, 1200), 255),
|
||||
{'steps': 8}, respond)
|
||||
# Working size is capped to a model-friendly size (<=768 long side, /16).
|
||||
assert seen['size'] == '768x576'
|
||||
# A requested step budget passes through up to the cap (default 8).
|
||||
assert seen['steps'] == '8'
|
||||
# The returned image keeps the original dimensions.
|
||||
assert Image.open(io.BytesIO(base64.b64decode(out['image']))).size == (1600, 1200)
|
||||
|
||||
|
||||
def test_self_hosted_inpaint_clamps_steps_to_the_cap(monkeypatch):
|
||||
seen = {}
|
||||
|
||||
def respond(request):
|
||||
seen.update(_form_fields(request.content, request.headers['content-type']))
|
||||
return httpx.Response(200, json={'data': [{'b64_json': png(Image.new('RGB', (640, 480), 'blue'))}]})
|
||||
|
||||
_self_hosted_inpaint(monkeypatch, Image.new('RGB', (640, 480), 'red'),
|
||||
Image.new('L', (640, 480), 255), {'steps': 20}, respond)
|
||||
# Over-cap requests are clamped to the configured maximum (default 8).
|
||||
assert seen['steps'] == '8'
|
||||
|
||||
|
||||
def test_self_hosted_inpaint_leaves_small_images_unchanged(monkeypatch):
|
||||
seen = {}
|
||||
|
||||
def respond(request):
|
||||
seen.update(_form_fields(request.content, request.headers['content-type']))
|
||||
return httpx.Response(200, json={'data': [{'b64_json': png(Image.new('RGB', (640, 480), 'green'))}]})
|
||||
|
||||
out = _self_hosted_inpaint(monkeypatch, Image.new('RGB', (640, 480), 'red'),
|
||||
Image.new('L', (640, 480), 255), {}, respond)
|
||||
# Under the cap: dimensions pass through untouched; with no requested count
|
||||
# sent the moderate default (4) is used.
|
||||
assert seen['size'] == '640x480'
|
||||
assert seen['steps'] == '4'
|
||||
assert Image.open(io.BytesIO(base64.b64decode(out['image']))).size == (640, 480)
|
||||
|
||||
|
||||
def test_self_hosted_inpaint_without_mask_uses_a_full_white_mask(monkeypatch):
|
||||
"""A whole-image instruction edit may omit the mask; it must not raise
|
||||
KeyError -> 400 (the health-audit regression)."""
|
||||
seen = {}
|
||||
mask_bytes = {}
|
||||
|
||||
def respond(request):
|
||||
seen.update(_form_fields(request.content, request.headers['content-type']))
|
||||
# The provider must receive a (full-white) mask, not a missing one.
|
||||
import re as _re
|
||||
parts = request.content.split(b'--')
|
||||
for part in parts:
|
||||
if b'name="mask"' in part:
|
||||
head, _, body = part.partition(b'\r\n\r\n')
|
||||
mask_bytes['png'] = body.rstrip(b'\r\n-')
|
||||
return httpx.Response(200, json={'data': [{'b64_json': png(Image.new('RGB', (32, 32), 'blue'))}]})
|
||||
|
||||
out = _self_hosted_inpaint(monkeypatch, Image.new('RGB', (32, 32), 'red'),
|
||||
None, {}, respond, include_mask=False)
|
||||
assert 'image' in out, out
|
||||
assert seen.get('size') == '32x32'
|
||||
# The synthesized mask is image-sized and all-white (whole image editable).
|
||||
mask_img = Image.open(io.BytesIO(mask_bytes['png'])).convert('L')
|
||||
assert mask_img.size == (32, 32)
|
||||
assert mask_img.getextrema() == (255, 255)
|
||||
|
||||
|
||||
def test_openai_inpaint_without_mask_is_a_clean_400(monkeypatch):
|
||||
"""OpenAI edits need a mask; a missing one must be a clean 400, not a
|
||||
KeyError traceback."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
endpoint = SimpleNamespace(base_url='https://api.openai.com/v1', api_key='sk-x', owner='alice')
|
||||
monkeypatch.setattr(gallery, 'SessionLocal', lambda: SimpleNamespace(close=lambda: None))
|
||||
monkeypatch.setattr(gallery, 'require_privilege', lambda *a: 'alice')
|
||||
monkeypatch.setattr(gallery, '_visible_image_endpoint_for_base', lambda *a: endpoint)
|
||||
monkeypatch.setattr('src.url_safety.check_outbound_url', lambda *a, **k: (True, ''))
|
||||
|
||||
class Request:
|
||||
async def json(self):
|
||||
return {'_endpoint': 'https://api.openai.com/v1',
|
||||
'image': png(Image.new('RGB', (32, 32), 'red')), 'prompt': 'x'}
|
||||
|
||||
route = next(r for r in gallery.setup_gallery_routes().routes if r.path == '/api/image/inpaint')
|
||||
try:
|
||||
asyncio.run(route.endpoint(Request()))
|
||||
raise AssertionError('expected a 400 for a missing mask')
|
||||
except HTTPException as exc:
|
||||
assert exc.status_code == 400
|
||||
assert 'mask' in str(exc.detail).lower()
|
||||
|
||||
|
||||
def test_self_hosted_inpaint_forwards_request_id(monkeypatch):
|
||||
"""A client-supplied request_id is forwarded to the provider so the client
|
||||
can poll GET /api/image/edit-progress for step progress."""
|
||||
seen = {}
|
||||
|
||||
def respond(request):
|
||||
seen.update(_form_fields(request.content, request.headers.get('content-type', '')))
|
||||
return httpx.Response(200, json={'model': 'qwen-image-edit-2511',
|
||||
'data': [{'b64_json': png(Image.new('RGBA', (32, 32), 'red'))}]})
|
||||
|
||||
out = _self_hosted_inpaint(monkeypatch, Image.new('RGB', (32, 32), 'red'),
|
||||
Image.new('L', (32, 32), 255), {'request_id': 'job-777'}, respond)
|
||||
assert 'image' in out, out
|
||||
assert seen.get('request_id') == 'job-777'
|
||||
|
||||
@@ -0,0 +1,120 @@
|
||||
"""Tests for the image-op progress contract (HARNESS-owned).
|
||||
|
||||
The 7011 side accepts an optional client-supplied ``request_id`` on heavy image
|
||||
ops and exposes ``GET /api/image/edit-progress`` which proxies the provider's
|
||||
``/v1/images/progress/{request_id}`` route. When the provider has no such route
|
||||
the endpoint returns ``{"available": false}`` so the client keeps an
|
||||
indeterminate bar.
|
||||
"""
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
|
||||
import httpx
|
||||
|
||||
import routes.gallery_routes as gallery
|
||||
|
||||
|
||||
def _route(path):
|
||||
return next(r for r in gallery.setup_gallery_routes().routes if r.path == path)
|
||||
|
||||
|
||||
class _Req:
|
||||
def __init__(self, query):
|
||||
self.query_params = query
|
||||
|
||||
|
||||
def _fake_db():
|
||||
class Db:
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
def query(self, *args, **kwargs):
|
||||
return self
|
||||
|
||||
def filter(self, *args, **kwargs):
|
||||
return self
|
||||
|
||||
def first(self):
|
||||
return None
|
||||
|
||||
return Db()
|
||||
|
||||
|
||||
def test_edit_progress_requires_request_id(monkeypatch):
|
||||
monkeypatch.setattr(gallery, "require_privilege", lambda *a: "alice")
|
||||
route = _route("/api/image/edit-progress")
|
||||
try:
|
||||
asyncio.run(route.endpoint(_Req({})))
|
||||
except Exception as exc: # FastAPI HTTPException
|
||||
assert getattr(exc, "status_code", None) == 400
|
||||
else:
|
||||
raise AssertionError("expected a 400 for a missing request_id")
|
||||
|
||||
|
||||
def test_edit_progress_rejects_an_unsafe_request_id(monkeypatch):
|
||||
monkeypatch.setattr(gallery, "require_privilege", lambda *a: "alice")
|
||||
monkeypatch.setattr(gallery, "SessionLocal", lambda: _fake_db())
|
||||
route = _route("/api/image/edit-progress")
|
||||
try:
|
||||
asyncio.run(route.endpoint(_Req({"request_id": "../secret"})))
|
||||
except Exception as exc:
|
||||
assert getattr(exc, "status_code", None) == 400
|
||||
else:
|
||||
raise AssertionError("expected a 400 for an unsafe request_id")
|
||||
|
||||
|
||||
def test_edit_progress_returns_provider_payload(monkeypatch):
|
||||
endpoint = SimpleNamespace(base_url="http://100.69.120.65:8188/v1", api_key="k")
|
||||
monkeypatch.setattr(gallery, "SessionLocal", lambda: _fake_db())
|
||||
monkeypatch.setattr(gallery, "require_privilege", lambda *a: "alice")
|
||||
monkeypatch.setattr(gallery, "_first_visible_image_endpoint", lambda *a: endpoint)
|
||||
seen = {}
|
||||
|
||||
def respond(request):
|
||||
seen["url"] = str(request.url)
|
||||
# Live seneca contract stages: loading | editing | compositing | done | error.
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={"stage": "editing", "step": 3, "steps": 8, "eta_s": 42.0, "progress": 0.375},
|
||||
)
|
||||
|
||||
client = httpx.AsyncClient
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda **k: client(transport=httpx.MockTransport(respond), **k))
|
||||
out = asyncio.run(_route("/api/image/edit-progress").endpoint(_Req({"request_id": "abc-123"})))
|
||||
assert out["available"] is True
|
||||
assert out["stage"] == "editing" # passed through unchanged
|
||||
assert out["step"] == 3 and out["steps"] == 8
|
||||
assert seen["url"].endswith("/images/progress/abc-123")
|
||||
|
||||
|
||||
def test_edit_progress_available_false_when_provider_has_no_route(monkeypatch):
|
||||
endpoint = SimpleNamespace(base_url="http://100.69.120.65:8188/v1", api_key="k")
|
||||
monkeypatch.setattr(gallery, "SessionLocal", lambda: _fake_db())
|
||||
monkeypatch.setattr(gallery, "require_privilege", lambda *a: "alice")
|
||||
monkeypatch.setattr(gallery, "_first_visible_image_endpoint", lambda *a: endpoint)
|
||||
|
||||
client = httpx.AsyncClient
|
||||
monkeypatch.setattr(
|
||||
httpx, "AsyncClient",
|
||||
lambda **k: client(transport=httpx.MockTransport(lambda req: httpx.Response(404)), **k),
|
||||
)
|
||||
out = asyncio.run(_route("/api/image/edit-progress").endpoint(_Req({"request_id": "abc"})))
|
||||
assert out["available"] is False
|
||||
|
||||
|
||||
def test_edit_progress_available_false_when_unreachable(monkeypatch):
|
||||
endpoint = SimpleNamespace(base_url="http://100.69.120.65:8188/v1", api_key="k")
|
||||
monkeypatch.setattr(gallery, "SessionLocal", lambda: _fake_db())
|
||||
monkeypatch.setattr(gallery, "require_privilege", lambda *a: "alice")
|
||||
monkeypatch.setattr(gallery, "_first_visible_image_endpoint", lambda *a: endpoint)
|
||||
|
||||
class Boom:
|
||||
async def __aenter__(self):
|
||||
raise httpx.ConnectError("nope")
|
||||
|
||||
async def __aexit__(self, *a):
|
||||
return False
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda **k: Boom())
|
||||
out = asyncio.run(_route("/api/image/edit-progress").endpoint(_Req({"request_id": "abc"})))
|
||||
assert out["available"] is False
|
||||
@@ -0,0 +1,201 @@
|
||||
"""Heavy in-process image ops must not load GB-scale models by default.
|
||||
|
||||
remove-bg / SAM mask / Real-ESRGAN upscale+denoise used to load rembg / torch
|
||||
models inside the web process and OOM-killed the host. They now default to a
|
||||
graceful 503, proxy to the seneca remote ops service when
|
||||
``ODYSSEUS_IMAGE_OPS_BASE`` is set (translate to the /v1/images/* contract), and
|
||||
only run in-process when explicitly enabled via ODYSSEUS_ALLOW_LOCAL_IMAGE_MODELS.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from routes.gallery import gallery_routes as gr
|
||||
|
||||
|
||||
HEAVY_ROUTES = [
|
||||
("/api/image/remove-bg", {"image": "x"}, "Background removal"),
|
||||
("/api/image/mask", {"image": "x"}, "Smart mask"),
|
||||
("/api/image/upscale-local", {"image": "x"}, "AI upscale"),
|
||||
("/api/image/denoise", {"image": "x"}, "AI denoise"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clean_env(monkeypatch):
|
||||
monkeypatch.setattr(gr, "require_privilege", lambda *a, **k: None)
|
||||
monkeypatch.delenv("ODYSSEUS_IMAGE_OPS_BASE", raising=False)
|
||||
monkeypatch.delenv("ODYSSEUS_ALLOW_LOCAL_IMAGE_MODELS", raising=False)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def client():
|
||||
app = FastAPI()
|
||||
app.include_router(gr.setup_gallery_routes())
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def test_local_image_models_disabled_by_default():
|
||||
assert gr._local_image_models_enabled() is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["1", "true", "Yes", "ON"])
|
||||
def test_local_image_models_enabled_env(monkeypatch, value):
|
||||
monkeypatch.setenv("ODYSSEUS_ALLOW_LOCAL_IMAGE_MODELS", value)
|
||||
assert gr._local_image_models_enabled() is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["", "0", "no", "off", "maybe"])
|
||||
def test_local_image_models_disabled_env(monkeypatch, value):
|
||||
monkeypatch.setenv("ODYSSEUS_ALLOW_LOCAL_IMAGE_MODELS", value)
|
||||
assert gr._local_image_models_enabled() is False
|
||||
|
||||
|
||||
def test_rss_mb_is_positive():
|
||||
assert gr._rss_mb() > 0
|
||||
|
||||
|
||||
def test_route_heavy_image_op_returns_none_when_enabled(monkeypatch):
|
||||
monkeypatch.setenv("ODYSSEUS_ALLOW_LOCAL_IMAGE_MODELS", "1")
|
||||
assert asyncio.run(gr._route_heavy_image_op("/api/image/denoise", {}, label="x")) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path,body,label", HEAVY_ROUTES)
|
||||
def test_heavy_ops_503_when_disabled(client, path, body, label):
|
||||
resp = client.post(path, json=body)
|
||||
assert resp.status_code == 503
|
||||
detail = resp.json().get("detail", "")
|
||||
assert "temporarily unavailable" in detail
|
||||
assert label in detail
|
||||
|
||||
|
||||
def test_remove_bg_disabled_does_not_import_rembg(monkeypatch):
|
||||
import sys
|
||||
|
||||
monkeypatch.setitem(sys.modules, "rembg", None)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(gr.setup_gallery_routes())
|
||||
client = TestClient(app)
|
||||
resp = client.post("/api/image/remove-bg", json={"image": "x"})
|
||||
assert resp.status_code == 503
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status_code=200):
|
||||
self._payload = payload
|
||||
self.status_code = status_code
|
||||
|
||||
def json(self):
|
||||
return self._payload
|
||||
|
||||
|
||||
class _FakeAsyncClient:
|
||||
calls = []
|
||||
response = _FakeResponse({})
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return False
|
||||
|
||||
async def post(self, url, json=None, files=None, data=None):
|
||||
_FakeAsyncClient.calls.append((url, {"json": json, "files": files, "data": data}))
|
||||
return _FakeAsyncClient.response
|
||||
|
||||
|
||||
def _client():
|
||||
app = FastAPI()
|
||||
app.include_router(gr.setup_gallery_routes())
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def test_remove_bg_proxies_to_remote_contract(monkeypatch):
|
||||
monkeypatch.setenv("ODYSSEUS_IMAGE_OPS_BASE", "http://seneca.test:8188")
|
||||
monkeypatch.setattr(httpx, "AsyncClient", _FakeAsyncClient)
|
||||
_FakeAsyncClient.calls = []
|
||||
_FakeAsyncClient.response = _FakeResponse(
|
||||
{"data": [{"b64_json": "RB", "mask_b64": "M", "bbox": [0, 0, 1, 1]}], "model": "birefnet"}
|
||||
)
|
||||
|
||||
resp = _client().post("/api/image/remove-bg", json={"image": "AAAA"})
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["image"] == "RB"
|
||||
url, kw = _FakeAsyncClient.calls[-1]
|
||||
assert url == "http://seneca.test:8188/v1/images/remove-background"
|
||||
assert "image" in kw["files"]
|
||||
|
||||
|
||||
def test_mask_proxies_to_remote_contract(monkeypatch):
|
||||
monkeypatch.setenv("ODYSSEUS_IMAGE_OPS_BASE", "http://seneca.test:8188")
|
||||
monkeypatch.setattr(httpx, "AsyncClient", _FakeAsyncClient)
|
||||
_FakeAsyncClient.calls = []
|
||||
_FakeAsyncClient.response = _FakeResponse(
|
||||
{"masks": [{"b64_json": "MASK", "bbox": [10, 20, 30, 40], "score": 0.9, "label": "person"}],
|
||||
"model": "sam+grounding"}
|
||||
)
|
||||
|
||||
resp = _client().post("/api/image/mask", json={"image": "AAAA", "text": "person"})
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["mask"] == "MASK"
|
||||
# remote [x,y,w,h] -> app [x1,y1,x2,y2]
|
||||
assert body["bbox"] == [10, 20, 40, 60]
|
||||
url, kw = _FakeAsyncClient.calls[-1]
|
||||
assert url == "http://seneca.test:8188/v1/images/mask"
|
||||
assert kw["data"]["prompt"] == "person"
|
||||
|
||||
|
||||
def test_mask_picks_top_score_not_first(monkeypatch):
|
||||
monkeypatch.setenv("ODYSSEUS_IMAGE_OPS_BASE", "http://seneca.test:8188")
|
||||
monkeypatch.setattr(httpx, "AsyncClient", _FakeAsyncClient)
|
||||
_FakeAsyncClient.response = _FakeResponse(
|
||||
{"masks": [
|
||||
{"b64_json": "LOW", "bbox": [0, 0, 1, 1], "score": 0.2},
|
||||
{"b64_json": "HIGH", "bbox": [0, 0, 2, 2], "score": 0.95},
|
||||
{"b64_json": "MID", "bbox": [0, 0, 3, 3], "score": 0.5},
|
||||
], "model": "sam+grounding"}
|
||||
)
|
||||
resp = _client().post("/api/image/mask", json={"image": "AAAA", "text": "person"})
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["mask"] == "HIGH"
|
||||
|
||||
|
||||
def test_mask_box_is_sent_as_xywh(monkeypatch):
|
||||
monkeypatch.setenv("ODYSSEUS_IMAGE_OPS_BASE", "http://seneca.test:8188")
|
||||
monkeypatch.setattr(httpx, "AsyncClient", _FakeAsyncClient)
|
||||
_FakeAsyncClient.calls = []
|
||||
_FakeAsyncClient.response = _FakeResponse({"masks": [{"b64_json": "M", "bbox": [0, 0, 1, 1], "score": 1.0}]})
|
||||
resp = _client().post("/api/image/mask", json={"image": "AAAA", "box": [10, 20, 40, 60]})
|
||||
assert resp.status_code == 200
|
||||
_, kw = _FakeAsyncClient.calls[-1]
|
||||
assert kw["data"]["box"] == "10,20,30,40"
|
||||
|
||||
|
||||
def test_mask_remote_empty_masks_is_200_nothing_selected(monkeypatch):
|
||||
monkeypatch.setenv("ODYSSEUS_IMAGE_OPS_BASE", "http://seneca.test:8188")
|
||||
monkeypatch.setattr(httpx, "AsyncClient", _FakeAsyncClient)
|
||||
_FakeAsyncClient.response = _FakeResponse({"masks": [], "model": "sam+grounding"})
|
||||
resp = _client().post("/api/image/mask", json={"image": "AAAA", "text": "person"})
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["mask"] is None
|
||||
assert body["masks"] == []
|
||||
assert "Nothing selected" in body["error"]
|
||||
|
||||
|
||||
def test_remote_proxy_surfaces_error_status(monkeypatch):
|
||||
monkeypatch.setenv("ODYSSEUS_IMAGE_OPS_BASE", "http://seneca.test:8188")
|
||||
monkeypatch.setattr(httpx, "AsyncClient", _FakeAsyncClient)
|
||||
_FakeAsyncClient.response = _FakeResponse({}, status_code=500)
|
||||
resp = _client().post("/api/image/remove-bg", json={"image": "AAAA"})
|
||||
assert resp.status_code == 503
|
||||
assert "temporarily unavailable" in resp.json().get("detail", "")
|
||||
@@ -0,0 +1,124 @@
|
||||
"""Image endpoint fallback: image models on mixed (non image-only) endpoints.
|
||||
|
||||
Chat image generation resolves its model by name across every enabled endpoint,
|
||||
so a ``qwen-image-edit`` / ``gpt-image-1`` model registered on a shared
|
||||
OpenAI-compatible gateway works there. The editor's generative ops used a
|
||||
stricter image-only lookup and reported "No image generation endpoint
|
||||
configured" on those same setups. These tests pin the fallback that makes the
|
||||
editor find the same endpoint.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import routes.gallery_routes as gallery_routes
|
||||
|
||||
|
||||
class _Predicate:
|
||||
def __init__(self, check):
|
||||
self._check = check
|
||||
|
||||
def __call__(self, row):
|
||||
return self._check(row)
|
||||
|
||||
def __or__(self, other):
|
||||
return _Predicate(lambda row: self(row) or other(row))
|
||||
|
||||
|
||||
class _Column:
|
||||
def __init__(self, name):
|
||||
self.name = name
|
||||
|
||||
def __eq__(self, value):
|
||||
return _Predicate(lambda row: getattr(row, self.name) == value)
|
||||
|
||||
|
||||
class _ModelEndpoint:
|
||||
base_url = _Column("base_url")
|
||||
model_type = _Column("model_type")
|
||||
is_enabled = _Column("is_enabled")
|
||||
owner = _Column("owner")
|
||||
|
||||
|
||||
class _Query:
|
||||
def __init__(self, rows):
|
||||
self._rows = list(rows)
|
||||
|
||||
def filter(self, *predicates):
|
||||
self._rows = [row for row in self._rows if all(pred(row) for pred in predicates)]
|
||||
return self
|
||||
|
||||
def all(self):
|
||||
return list(self._rows)
|
||||
|
||||
|
||||
class _DB:
|
||||
def __init__(self, rows):
|
||||
self._rows = rows
|
||||
|
||||
def query(self, model):
|
||||
assert model is _ModelEndpoint
|
||||
return _Query(self._rows)
|
||||
|
||||
|
||||
def _ep(base_url, owner=None, *, enabled=True, model_type="chat", models=None):
|
||||
return SimpleNamespace(
|
||||
base_url=base_url,
|
||||
owner=owner,
|
||||
is_enabled=enabled,
|
||||
model_type=model_type,
|
||||
api_key="sk-secret",
|
||||
cached_models=models or [],
|
||||
pinned_models=[],
|
||||
)
|
||||
|
||||
|
||||
def _patch_model(monkeypatch):
|
||||
monkeypatch.setattr(gallery_routes, "ModelEndpoint", _ModelEndpoint)
|
||||
|
||||
|
||||
def test_mixed_endpoint_with_image_model_is_found(monkeypatch):
|
||||
_patch_model(monkeypatch)
|
||||
rows = [_ep("https://gateway.example.com/v1", models=["qwen-image-edit", "llama-3"])]
|
||||
|
||||
ep = gallery_routes._first_visible_image_endpoint(_DB(rows), None)
|
||||
|
||||
assert ep is not None
|
||||
assert ep.base_url == "https://gateway.example.com/v1"
|
||||
|
||||
|
||||
def test_mixed_endpoint_without_image_model_is_ignored(monkeypatch):
|
||||
_patch_model(monkeypatch)
|
||||
rows = [_ep("https://gateway.example.com/v1", models=["llama-3", "mistral"])]
|
||||
|
||||
assert gallery_routes._first_visible_image_endpoint(_DB(rows), None) is None
|
||||
|
||||
|
||||
def test_image_only_endpoint_wins_over_mixed(monkeypatch):
|
||||
_patch_model(monkeypatch)
|
||||
rows = [
|
||||
_ep("https://gateway.example.com/v1", models=["gpt-image-1"]),
|
||||
_ep("https://images.example.com/v1", model_type="image"),
|
||||
]
|
||||
|
||||
ep = gallery_routes._first_visible_image_endpoint(_DB(rows), None)
|
||||
|
||||
assert ep is not None
|
||||
assert ep.model_type == "image"
|
||||
|
||||
|
||||
def test_unreachable_endpoint_reports_502():
|
||||
"""The inpaint proxy must distinguish an unreachable endpoint from a provider error.
|
||||
|
||||
502 (gateway) is reserved for "the image server could not be reached"; 503 is
|
||||
reserved for "generation is switched off on the server", which the editor maps
|
||||
to its own fixed message.
|
||||
"""
|
||||
from pathlib import Path
|
||||
|
||||
src = (
|
||||
Path(__file__).resolve().parents[1] / "routes" / "gallery" / "gallery_routes.py"
|
||||
).read_text()
|
||||
assert "Image server unreachable" in src
|
||||
assert "except httpx.TransportError as exc:" in src
|
||||
assert "provider HTTP" in src
|
||||
assert 'raise HTTPException(503, f"Image endpoint unreachable' not in src
|
||||
@@ -0,0 +1,27 @@
|
||||
"""Source checks for the gallery Settings-tab image-model picker.
|
||||
|
||||
The Quick Edit model picker lives in the editor; HARNESS mirrors it into the
|
||||
gallery Settings tab, fed by GET /api/editor/image-models and sharing the
|
||||
same localStorage key (`ge-ai-command-model`) so the two stay in sync.
|
||||
"""
|
||||
from pathlib import Path
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
GALLERY_JS = (ROOT / "static/js/gallery.js").read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def test_settings_tab_has_image_model_picker():
|
||||
assert 'id="gallery-image-model"' in GALLERY_JS
|
||||
assert "/api/editor/image-models" in GALLERY_JS
|
||||
assert "ge-ai-command-model" in GALLERY_JS
|
||||
|
||||
|
||||
def test_image_model_picker_lives_in_the_settings_container():
|
||||
settings_idx = GALLERY_JS.index('id="gallery-settings-container"')
|
||||
picker_idx = GALLERY_JS.index('id="gallery-image-model"')
|
||||
assert picker_idx > settings_idx
|
||||
|
||||
|
||||
def test_picker_persists_with_the_shared_localstorage_key():
|
||||
assert "localStorage.setItem(_GE_MODEL_KEY" in GALLERY_JS
|
||||
assert "localStorage.getItem(_GE_MODEL_KEY)" in GALLERY_JS
|
||||
@@ -0,0 +1,25 @@
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
GALLERY = ROOT / "static/js/gallery.js"
|
||||
|
||||
|
||||
def test_editor_landing_has_an_import_button_next_to_new_canvas():
|
||||
"""The editor landing must offer a direct "Import" action.
|
||||
|
||||
It sits in the same action row as "New canvas..." / "Browse gallery" and,
|
||||
on file pick, opens the editor on the chosen raster image (``openEditor``
|
||||
takes an image data URL as its first argument).
|
||||
"""
|
||||
src = GALLERY.read_text(encoding="utf-8")
|
||||
assert 'id="gallery-editor-import"' in src, "landing Import button missing"
|
||||
# Button order: New canvas → Import → Browse gallery.
|
||||
new_at = src.index('id="gallery-editor-new"')
|
||||
import_at = src.index('id="gallery-editor-import"')
|
||||
pick_at = src.index('id="gallery-editor-pick"')
|
||||
assert new_at < import_at < pick_at, "Import must sit next to New canvas"
|
||||
# Wiring opens the editor on the selected file's data URL.
|
||||
assert "getElementById('gallery-editor-import')?.addEventListener('click'" in src
|
||||
assert "readAsDataURL(file)" in src
|
||||
assert "await openEditor(reader.result, null, null, name)" in src
|
||||
@@ -0,0 +1,187 @@
|
||||
"""Tests for the seneca single-image op proxies (HARNESS-owned).
|
||||
|
||||
POST /api/image/replace-background -> provider /v1/images/replace-background
|
||||
POST /api/image/outpaint -> provider /v1/images/outpaint (alias /extend)
|
||||
|
||||
Both forward a client-supplied JPEG/PNG as multipart with the prompt/params plus
|
||||
``request_id`` (so GET /api/image/edit-progress lights the progress bar) and
|
||||
return the provider image as ``{image}``.
|
||||
"""
|
||||
import asyncio
|
||||
import base64
|
||||
from types import SimpleNamespace
|
||||
|
||||
import httpx
|
||||
|
||||
import routes.gallery_routes as gallery
|
||||
|
||||
_IMG_B64 = base64.b64encode(b"abc-not-a-real-png").decode()
|
||||
|
||||
|
||||
def _route(path):
|
||||
return next(r for r in gallery.setup_gallery_routes().routes if r.path == path)
|
||||
|
||||
|
||||
class _Req:
|
||||
def __init__(self, body):
|
||||
self._body = body
|
||||
|
||||
async def json(self):
|
||||
return self._body
|
||||
|
||||
|
||||
def _fake_db():
|
||||
class Db:
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
def query(self, *args, **kwargs):
|
||||
return self
|
||||
|
||||
def filter(self, *args, **kwargs):
|
||||
return self
|
||||
|
||||
def first(self):
|
||||
return None
|
||||
|
||||
return Db()
|
||||
|
||||
|
||||
def _setup(monkeypatch, respond):
|
||||
endpoint = SimpleNamespace(base_url="http://100.69.120.65:8188/v1", api_key="k")
|
||||
monkeypatch.setattr(gallery, "SessionLocal", lambda: _fake_db())
|
||||
monkeypatch.setattr(gallery, "require_privilege", lambda *a: "alice")
|
||||
monkeypatch.setattr(gallery, "_first_visible_image_endpoint", lambda *a: endpoint)
|
||||
seen = {}
|
||||
|
||||
def _respond(request):
|
||||
seen["url"] = str(request.url)
|
||||
seen["body"] = bytes(request.content)
|
||||
return respond(request)
|
||||
|
||||
client = httpx.AsyncClient
|
||||
monkeypatch.setattr(
|
||||
httpx, "AsyncClient",
|
||||
lambda **k: client(transport=httpx.MockTransport(_respond), **k),
|
||||
)
|
||||
return seen
|
||||
|
||||
|
||||
def _ok(_req):
|
||||
return httpx.Response(200, json={"data": [{"b64_json": "AAAA"}]})
|
||||
|
||||
|
||||
def test_replace_background_forwards_multipart_and_returns_image(monkeypatch):
|
||||
seen = _setup(monkeypatch, _ok)
|
||||
out = asyncio.run(_route("/api/image/replace-background").endpoint(_Req({
|
||||
"image": _IMG_B64, "prompt": "put her on a beach", "steps": 8,
|
||||
"harmonise": True, "shadow": False, "request_id": "rb-123",
|
||||
})))
|
||||
assert out["image"] == "AAAA"
|
||||
assert seen["url"].endswith("/images/replace-background")
|
||||
body = seen["body"]
|
||||
assert b"put her on a beach" in body
|
||||
assert b"rb-123" in body # request_id forwarded for progress
|
||||
assert b"harmonise" in body and b"shadow" in body
|
||||
assert b"8" in body # steps
|
||||
|
||||
|
||||
def test_outpaint_forwards_ratio_size_context_and_returns_image(monkeypatch):
|
||||
seen = _setup(monkeypatch, _ok)
|
||||
out = asyncio.run(_route("/api/image/outpaint").endpoint(_Req({
|
||||
"image": _IMG_B64, "prompt": "extend the sky", "ratio": "16:9",
|
||||
"size": "1920x1080", "steps": 4, "context": "calm blue sky",
|
||||
"request_id": "op-9",
|
||||
})))
|
||||
assert out["image"] == "AAAA"
|
||||
assert seen["url"].endswith("/images/outpaint")
|
||||
body = seen["body"]
|
||||
for token in (b"16:9", b"1920x1080", b"calm blue sky", b"op-9"):
|
||||
assert token in body, token
|
||||
|
||||
|
||||
def test_outpaint_extend_alias_uses_the_same_provider_route(monkeypatch):
|
||||
seen = _setup(monkeypatch, _ok)
|
||||
out = asyncio.run(_route("/api/image/extend").endpoint(_Req({
|
||||
"image": _IMG_B64, "prompt": "extend", "ratio": "9:16",
|
||||
})))
|
||||
assert out["image"] == "AAAA"
|
||||
assert seen["url"].endswith("/images/outpaint")
|
||||
|
||||
|
||||
def test_outpaint_reads_openai_style_url_result(monkeypatch):
|
||||
def _respond(_req):
|
||||
return httpx.Response(200, json={"data": [{"url": "http://100.69.120.65:8188/v1/images/x.png"}]})
|
||||
|
||||
_setup(monkeypatch, _respond)
|
||||
monkeypatch.setattr(gallery, "_fetch_result_image_b64", lambda url: _async_return("ZZZZ"))
|
||||
out = asyncio.run(_route("/api/image/outpaint").endpoint(_Req({"image": _IMG_B64, "prompt": "x"})))
|
||||
assert out["image"] == "ZZZZ"
|
||||
|
||||
|
||||
def test_replace_background_502_when_provider_returns_no_image(monkeypatch):
|
||||
_setup(monkeypatch, lambda _req: httpx.Response(200, json={"data": []}))
|
||||
try:
|
||||
asyncio.run(_route("/api/image/replace-background").endpoint(_Req({"image": _IMG_B64})))
|
||||
except Exception as exc:
|
||||
assert getattr(exc, "status_code", None) == 502
|
||||
else:
|
||||
raise AssertionError("expected a 502 when the provider returns no image")
|
||||
|
||||
|
||||
async def _async_return(value):
|
||||
return value
|
||||
|
||||
|
||||
def _capture_forward(monkeypatch):
|
||||
"""Replace _forward_seneca_image_op with a fake that records kwargs."""
|
||||
captured = {}
|
||||
|
||||
async def _fake_forward(base, api_key, provider_path, image_bytes, fields, *, label, timeout=300.0):
|
||||
captured["label"] = label
|
||||
captured["provider_path"] = provider_path
|
||||
captured["timeout"] = timeout
|
||||
return {"image": "AAAA"}
|
||||
|
||||
monkeypatch.setattr(gallery, "_forward_seneca_image_op", _fake_forward)
|
||||
monkeypatch.setattr(gallery, "SessionLocal", lambda: _fake_db())
|
||||
monkeypatch.setattr(gallery, "require_privilege", lambda *a: "alice")
|
||||
monkeypatch.setattr(
|
||||
gallery, "_first_visible_image_endpoint",
|
||||
lambda *a: SimpleNamespace(base_url="http://100.69.120.65:8188/v1", api_key="k"),
|
||||
)
|
||||
return captured
|
||||
|
||||
|
||||
def test_outpaint_proxy_uses_a_long_timeout(monkeypatch):
|
||||
captured = _capture_forward(monkeypatch)
|
||||
out = asyncio.run(_route("/api/image/outpaint").endpoint(_Req({"image": _IMG_B64, "prompt": "extend to 9:16"})))
|
||||
assert out["image"] == "AAAA"
|
||||
assert captured["provider_path"] == "/images/outpaint"
|
||||
assert captured["timeout"] >= 900, captured
|
||||
|
||||
|
||||
def test_outpaint_extend_alias_also_uses_the_long_timeout(monkeypatch):
|
||||
captured = _capture_forward(monkeypatch)
|
||||
asyncio.run(_route("/api/image/extend").endpoint(_Req({"image": _IMG_B64, "prompt": "extend"})))
|
||||
assert captured["timeout"] >= 900, captured
|
||||
|
||||
|
||||
def test_replace_background_keeps_the_short_default_timeout(monkeypatch):
|
||||
captured = _capture_forward(monkeypatch)
|
||||
asyncio.run(_route("/api/image/replace-background").endpoint(_Req({"image": _IMG_B64, "prompt": "beach"})))
|
||||
assert captured["timeout"] == 300.0, captured
|
||||
|
||||
|
||||
def test_outpaint_timeout_env_override(monkeypatch):
|
||||
monkeypatch.setenv("ODYSSEUS_IMAGE_OUTPAINT_TIMEOUT", "720")
|
||||
assert gallery._outpaint_timeout() == 720.0
|
||||
monkeypatch.delenv("ODYSSEUS_IMAGE_OUTPAINT_TIMEOUT", raising=False)
|
||||
assert gallery._outpaint_timeout() == 900.0
|
||||
|
||||
|
||||
def test_graceful_shutdown_window_covers_a_long_outpaint():
|
||||
from pathlib import Path
|
||||
|
||||
src = (Path(__file__).resolve().parent.parent / "app.py").read_text()
|
||||
assert '"ODYSSEUS_GRACEFUL_SHUTDOWN_TIMEOUT", "900"' in src
|
||||
@@ -1,11 +1,5 @@
|
||||
import pytest
|
||||
|
||||
from services.hwfit.fit import rank_models
|
||||
from services.hwfit.models import get_models, is_prequantized
|
||||
from tests.hwfit_publication_fixtures import publication_catalog # noqa: F401
|
||||
|
||||
# Rank authored inputs rather than publication catalog snapshots.
|
||||
pytestmark = pytest.mark.usefixtures("publication_catalog")
|
||||
|
||||
|
||||
def _8gb_vram_system():
|
||||
@@ -20,63 +14,38 @@ def _8gb_vram_system():
|
||||
}
|
||||
|
||||
|
||||
GEMMA = "google/gemma-4-12B-it"
|
||||
def test_gemma4_12b_in_catalog():
|
||||
catalog = {m["name"]: m for m in get_models()}
|
||||
assert "google/gemma-4-12B-it" in catalog, "gemma-4-12B-it missing from catalog"
|
||||
|
||||
|
||||
def _input_row(rows, name):
|
||||
return next(r for r in rows if r["name"] == name)
|
||||
|
||||
|
||||
def _qat_rows(rows):
|
||||
return [r for r in rows if r["name"].startswith(GEMMA + "-qat-")]
|
||||
|
||||
|
||||
def test_gemma4_12b_user_catalog_row_wins_the_merge(publication_catalog, monkeypatch):
|
||||
"""A dynamic cache listing the same repo must not replace or duplicate the
|
||||
user-catalog row; rows only the cache knows are still merged in."""
|
||||
from services.hwfit import hf_discovery
|
||||
|
||||
monkeypatch.setattr(hf_discovery, "load_cached_hf_collection_models", lambda: [
|
||||
{"name": GEMMA, "quantization": "F16", "gguf_sources": [{"repo": "test/cache-GGUF", "file": "cache.gguf"}]},
|
||||
{"name": "test/cache-only", "quantization": "F16", "gguf_sources": []},
|
||||
])
|
||||
merged = get_models()
|
||||
names = [m["name"] for m in merged]
|
||||
|
||||
assert names.count(GEMMA) == 1
|
||||
assert "test/cache-only" in names
|
||||
entry = next(m for m in merged if m["name"] == GEMMA)
|
||||
assert entry["gguf_sources"] == _input_row(publication_catalog, GEMMA)["gguf_sources"]
|
||||
|
||||
|
||||
def test_gemma4_12b_ranked_row_carries_its_gguf_source(publication_catalog):
|
||||
hit = next(r for r in rank_models(_8gb_vram_system(), search="gemma-4-12B-it", limit=20) if r["name"] == GEMMA)
|
||||
assert hit["gguf_sources"] == _input_row(publication_catalog, GEMMA)["gguf_sources"]
|
||||
# The F16 input does not fit 8 GB, so the fit falls back to a GGUF quant.
|
||||
assert hit["quant"] == "Q4_K_M"
|
||||
def test_gemma4_12b_has_gguf_source():
|
||||
catalog = {m["name"]: m for m in get_models()}
|
||||
entry = catalog["google/gemma-4-12B-it"]
|
||||
assert entry.get("gguf_sources"), "gemma-4-12B-it has no gguf_sources"
|
||||
repos = [s["repo"] for s in entry["gguf_sources"]]
|
||||
assert "unsloth/gemma-4-12B-it-GGUF" in repos
|
||||
|
||||
|
||||
def test_gemma4_12b_rank_models_returns_it_for_8gb_vram():
|
||||
results = rank_models(_8gb_vram_system(), search="gemma-4-12B-it", limit=20)
|
||||
names = [r["name"] for r in results]
|
||||
assert GEMMA in names, "rank_models did not return gemma-4-12B-it for 8 GB VRAM"
|
||||
assert "google/gemma-4-12B-it" in names, "rank_models did not return gemma-4-12B-it for 8 GB VRAM"
|
||||
|
||||
|
||||
def test_gemma4_12b_qat_entries_rank_with_their_native_quant(publication_catalog):
|
||||
qat = _qat_rows(publication_catalog)
|
||||
assert len(qat) == 2
|
||||
ranked = {r["name"]: r for r in rank_models(_8gb_vram_system(), search="gemma-4-12B-it-qat", limit=20)}
|
||||
for row in qat:
|
||||
assert ranked[row["name"]]["quant"] == row["quantization"]
|
||||
|
||||
|
||||
def test_gemma4_12b_qat_entries_are_prequantized(publication_catalog):
|
||||
def test_gemma4_12b_qat_entries_in_catalog():
|
||||
catalog = {m["name"]: m for m in get_models()}
|
||||
for row in _qat_rows(publication_catalog):
|
||||
assert is_prequantized(catalog[row["name"]])
|
||||
assert "google/gemma-4-12B-it-qat-int4" in catalog
|
||||
assert "google/gemma-4-12B-it-qat-int8" in catalog
|
||||
|
||||
|
||||
def test_gemma4_12b_qat_entries_have_no_gguf(publication_catalog):
|
||||
ranked = {r["name"]: r for r in rank_models(_8gb_vram_system(), search="gemma-4-12B-it-qat", limit=20)}
|
||||
for row in _qat_rows(publication_catalog):
|
||||
assert ranked[row["name"]]["gguf_sources"] == []
|
||||
def test_gemma4_12b_qat_entries_are_prequantized():
|
||||
catalog = {m["name"]: m for m in get_models()}
|
||||
assert is_prequantized(catalog["google/gemma-4-12B-it-qat-int4"])
|
||||
assert is_prequantized(catalog["google/gemma-4-12B-it-qat-int8"])
|
||||
|
||||
|
||||
def test_gemma4_12b_qat_entries_have_no_gguf():
|
||||
catalog = {m["name"]: m for m in get_models()}
|
||||
assert catalog["google/gemma-4-12B-it-qat-int4"]["gguf_sources"] == []
|
||||
assert catalog["google/gemma-4-12B-it-qat-int8"]["gguf_sources"] == []
|
||||
|
||||
@@ -3,33 +3,8 @@
|
||||
The handler did `n = int(gpu_count)` with no guard, so `?gpu_count=abc` (or any
|
||||
non-integer) raised ValueError -> HTTP 500. A malformed count is now ignored,
|
||||
matching how the neighbouring gpu_group param is already parsed.
|
||||
|
||||
The count parsing runs only after the empty-catalog early return, so these
|
||||
tests rank the authored publication catalog against a fixed two-GPU system and
|
||||
spy on rank_models to prove the parsing path actually ran.
|
||||
"""
|
||||
import functools
|
||||
from copy import deepcopy
|
||||
|
||||
import pytest
|
||||
|
||||
from routes.hwfit_routes import setup_hwfit_routes
|
||||
from tests.hwfit_publication_fixtures import publication_catalog # noqa: F401
|
||||
|
||||
SYSTEM = {
|
||||
"has_gpu": True,
|
||||
"backend": "cuda",
|
||||
"gpu_name": "Synthetic GPU",
|
||||
"gpu_vram_gb": 24.0,
|
||||
"gpu_count": 2,
|
||||
"gpus": [
|
||||
{"index": 0, "name": "Synthetic GPU", "vram_gb": 12.0},
|
||||
{"index": 1, "name": "Synthetic GPU", "vram_gb": 12.0},
|
||||
],
|
||||
"gpu_groups": [{"name": "Synthetic GPU", "vram_each": 12.0, "count": 2, "indices": [0, 1], "vram_total": 24.0}],
|
||||
"available_ram_gb": 32.0,
|
||||
"total_ram_gb": 32.0,
|
||||
}
|
||||
|
||||
|
||||
def _get_models():
|
||||
@@ -40,80 +15,24 @@ def _get_models():
|
||||
raise AssertionError("hwfit /models route not found")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ranked_systems(publication_catalog, monkeypatch):
|
||||
"""Fix detection and record every system the handler hands to rank_models."""
|
||||
from services.hwfit import fit, hardware
|
||||
|
||||
monkeypatch.setattr(hardware, "detect_system", lambda **_: deepcopy(SYSTEM))
|
||||
calls = []
|
||||
real_rank = fit.rank_models
|
||||
|
||||
@functools.wraps(real_rank)
|
||||
def spy(system, **kwargs):
|
||||
calls.append(deepcopy(system))
|
||||
return real_rank(system, **kwargs)
|
||||
|
||||
monkeypatch.setattr(fit, "rank_models", spy)
|
||||
return calls
|
||||
|
||||
|
||||
def _assert_ranked(result, ranked_systems):
|
||||
assert isinstance(result, dict)
|
||||
assert "error" not in result
|
||||
assert isinstance(result["models"], list) and result["models"]
|
||||
# Reaching rank_models means the empty-catalog return was not taken and the
|
||||
# manual-hardware and gpu_count parsing ran first.
|
||||
assert len(ranked_systems) == 1
|
||||
assert ranked_systems[0] == result["system"]
|
||||
return result["system"]
|
||||
|
||||
|
||||
def test_non_numeric_gpu_count_does_not_raise(ranked_systems):
|
||||
def test_non_numeric_gpu_count_does_not_raise():
|
||||
handler = _get_models()
|
||||
# Previously raised ValueError (HTTP 500); now degrades to a normal ranking.
|
||||
system = _assert_ranked(handler(gpu_count="abc"), ranked_systems)
|
||||
# Ignored like an omitted count: the auto pool keeps both GPUs and no
|
||||
# explicit-count GPU-only pin is applied.
|
||||
assert system["detected_gpu_count"] == 2
|
||||
assert system["active_group"]["use_count"] == 2
|
||||
assert system["gpu_count"] == 2
|
||||
assert "gpu_only" not in system
|
||||
result = handler(gpu_count="abc")
|
||||
assert isinstance(result, dict)
|
||||
|
||||
|
||||
def test_numeric_gpu_count_still_accepted(ranked_systems):
|
||||
def test_numeric_gpu_count_still_accepted():
|
||||
handler = _get_models()
|
||||
system = _assert_ranked(handler(gpu_count="0"), ranked_systems)
|
||||
# 0 switches to RAM-only ranking.
|
||||
assert system["detected_gpu_count"] == 2
|
||||
assert system["has_gpu"] is False
|
||||
assert system["gpu_count"] == 0
|
||||
assert system["gpu_only"] is False
|
||||
assert "active_group" not in system
|
||||
result = handler(gpu_count="0")
|
||||
assert isinstance(result, dict)
|
||||
|
||||
|
||||
def test_non_numeric_manual_gpu_count_does_not_raise(ranked_systems):
|
||||
def test_non_numeric_manual_gpu_count_does_not_raise():
|
||||
# manual_gpu_count is the other count param on this endpoint (the hardware
|
||||
# simulator in _apply_manual_hardware). A non-numeric value must also degrade
|
||||
# (default to 1) rather than 500, so the endpoint's count parsing is fully
|
||||
# covered.
|
||||
handler = _get_models()
|
||||
system = _assert_ranked(handler(manual_mode="gpu", manual_gpu_count="abc"), ranked_systems)
|
||||
assert system["manual_hardware"] is True
|
||||
assert system["gpu_name"] == "Simulated CUDA GPU"
|
||||
assert system["gpu_count"] == 1
|
||||
assert system["gpu_groups"][0]["count"] == 1
|
||||
assert system["detected_gpu_count"] == 1
|
||||
|
||||
|
||||
def test_empty_catalog_returns_before_count_parsing(ranked_systems, monkeypatch):
|
||||
# Control: without catalog rows the handler stops before the parsing above,
|
||||
# which is why the tests in this module must rank the authored catalog.
|
||||
from services.hwfit import models
|
||||
|
||||
monkeypatch.setattr(models, "get_models", lambda: [])
|
||||
result = _get_models()(gpu_count="abc", manual_mode="gpu", manual_gpu_count="abc")
|
||||
assert result["models"] == []
|
||||
assert "Model catalog is empty" in result["error"]
|
||||
assert ranked_systems == []
|
||||
assert "detected_gpu_count" not in result["system"]
|
||||
result = handler(manual_mode="gpu", manual_gpu_count="abc")
|
||||
assert isinstance(result, dict)
|
||||
|
||||
+10
-17
@@ -6,15 +6,9 @@ guarantee that non-macOS (Linux/Windows) detection is unchanged.
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from services.hwfit import hardware
|
||||
from services.hwfit.fit import rank_models
|
||||
from services.hwfit.models import get_models
|
||||
from tests.hwfit_publication_fixtures import publication_catalog # noqa: F401
|
||||
|
||||
# Rank authored inputs rather than publication catalog snapshots.
|
||||
pytestmark = pytest.mark.usefixtures("publication_catalog")
|
||||
|
||||
|
||||
def _metal_system(ram_gb=16.0, vram_gb=10.7):
|
||||
@@ -84,20 +78,19 @@ def test_only_gguf_or_mlx_models_recommended_on_metal():
|
||||
assert unservable == [], f"{len(unservable)} non-servable models on Metal, e.g. {unservable[:3]}"
|
||||
|
||||
|
||||
def test_qwen_gguf_rows_rank_with_their_gguf_downloads_on_metal(publication_catalog):
|
||||
"""Qwen GGUF-looking Cookbook rows must download the GGUF repo/file their
|
||||
catalog row names, not the base safetensors repository."""
|
||||
def test_qwen_catalog_entries_point_at_verified_gguf_repos():
|
||||
"""Qwen GGUF-looking Cookbook rows must download GGUF repos, not the base
|
||||
safetensors repositories."""
|
||||
catalog = {m["name"]: m for m in get_models()}
|
||||
expected = {
|
||||
r["name"]: r["gguf_sources"]
|
||||
for r in publication_catalog
|
||||
if r["name"].startswith("Qwen/Qwen3.") and r["gguf_sources"]
|
||||
"Qwen/Qwen3.5-9B": ("unsloth/Qwen3.5-9B-GGUF", "Qwen3.5-9B-Q4_K_M.gguf"),
|
||||
"Qwen/Qwen3.6-27B": ("unsloth/Qwen3.6-27B-GGUF", "Qwen3.6-27B-Q4_K_M.gguf"),
|
||||
"Qwen/Qwen3.6-35B-A3B": ("unsloth/Qwen3.6-35B-A3B-GGUF", "Qwen3.6-35B-A3B-UD-Q4_K_M.gguf"),
|
||||
}
|
||||
assert len(expected) == 3
|
||||
ranked = {r["name"]: r for r in rank_models(_metal_system(), search="Qwen/Qwen3.", limit=50)}
|
||||
|
||||
for model_name, sources in expected.items():
|
||||
assert ranked[model_name]["gguf_sources"] == sources
|
||||
assert all(src["repo"] != model_name and src["file"].endswith(".gguf") for src in ranked[model_name]["gguf_sources"])
|
||||
for model_name, (repo, filename) in expected.items():
|
||||
sources = catalog[model_name].get("gguf_sources") or []
|
||||
assert any(src.get("repo") == repo and src.get("file") == filename for src in sources)
|
||||
|
||||
|
||||
def test_safetensors_models_still_recommended_on_cuda():
|
||||
|
||||
@@ -1,15 +1,9 @@
|
||||
import pytest
|
||||
|
||||
from services.hwfit.fit import analyze_model, rank_models
|
||||
from services.hwfit.models import (
|
||||
get_models,
|
||||
infer_quantization_from_name,
|
||||
is_prequantized,
|
||||
)
|
||||
from tests.hwfit_publication_fixtures import publication_catalog # noqa: F401
|
||||
|
||||
# Rank authored inputs rather than publication catalog snapshots.
|
||||
pytestmark = pytest.mark.usefixtures("publication_catalog")
|
||||
|
||||
|
||||
def _dual_5060ti_system():
|
||||
@@ -41,7 +35,7 @@ def test_infers_native_hf_quant_formats_from_repo_names():
|
||||
|
||||
def test_nvfp4_catalog_quant_is_preserved():
|
||||
catalog = {m["name"]: m for m in get_models()}
|
||||
model = catalog["test/native-NVFP4"]
|
||||
model = catalog["txn545/Qwen3.5-122B-A10B-NVFP4"]
|
||||
|
||||
assert model["quantization"] == "NVFP4"
|
||||
assert is_prequantized(model)
|
||||
@@ -49,7 +43,7 @@ def test_nvfp4_catalog_quant_is_preserved():
|
||||
|
||||
def test_nvfp4_search_result_is_not_gguf_or_cpu_offload():
|
||||
catalog = {m["name"]: m for m in get_models()}
|
||||
model = catalog["test/native-NVFP4"]
|
||||
model = catalog["txn545/Qwen3.5-122B-A10B-NVFP4"]
|
||||
|
||||
fit = analyze_model(model, _dual_5060ti_system())
|
||||
assert fit["quant"] == "NVFP4"
|
||||
@@ -57,10 +51,10 @@ def test_nvfp4_search_result_is_not_gguf_or_cpu_offload():
|
||||
|
||||
results = rank_models(
|
||||
_dual_5060ti_system(),
|
||||
search="native-NVFP4",
|
||||
search="Qwen3.5-122B-A10B-NVFP4",
|
||||
limit=10,
|
||||
)
|
||||
hit = next(r for r in results if r["name"] == "test/native-NVFP4")
|
||||
hit = next(r for r in results if r["name"] == "txn545/Qwen3.5-122B-A10B-NVFP4")
|
||||
assert hit["quant"] == "NVFP4"
|
||||
assert hit["run_mode"] != "cpu_offload"
|
||||
|
||||
|
||||
@@ -120,8 +120,9 @@ def test_disabled_generation_is_not_restored():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('incognito', [False, True])
|
||||
@pytest.mark.parametrize('exit_code', [None, 0])
|
||||
async def test_generation_dispatch_uses_owner_aware_backend(monkeypatch, exit_code):
|
||||
async def test_generation_dispatch_uses_owner_aware_backend(monkeypatch, exit_code, incognito):
|
||||
from types import SimpleNamespace
|
||||
from src import ai_interaction, tool_execution
|
||||
from src.agent_runtime.journal import ActionJournal, bind_journal
|
||||
@@ -145,6 +146,7 @@ async def test_generation_dispatch_uses_owner_aware_backend(monkeypatch, exit_co
|
||||
_, denied = await execute(block, owner='pewds', session_id='fixture',
|
||||
disabled_tools={'generate_image'}, security_context=tool_execution.NO_TOOL_SECURITY_CONTEXT)
|
||||
_, result = await execute(block, owner='pewds', session_id='fixture',
|
||||
client_runtime_context={'web_ui_state': {'incognito': incognito}},
|
||||
security_context=tool_execution.NO_TOOL_SECURITY_CONTEXT)
|
||||
assert denied['exit_code'] != 0
|
||||
assert journal.actions[0].execution_id is None
|
||||
@@ -152,7 +154,7 @@ async def test_generation_dispatch_uses_owner_aware_backend(monkeypatch, exit_co
|
||||
assert journal.actions[1].execution_id is not None
|
||||
assert journal.actions[1].outcome['authoritative'] is (exit_code == 0)
|
||||
assert result['image_id'] == 'test'
|
||||
assert calls == [(block.content, {'owner': 'pewds', 'session_id': 'fixture'})]
|
||||
assert calls == [(block.content, {'owner': 'pewds', 'session_id': 'fixture', **({'incognito': True} if incognito else {})})]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
"""The heavy image-edit CPU work must run off the event loop, and the app must
|
||||
expose in-flight heavy-op counts + a graceful shutdown so a deploy restart does
|
||||
not drop a 60-150 s edit."""
|
||||
import base64
|
||||
import io
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from routes.gallery import gallery_routes as gr
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def _png_b64(w, h, color=(120, 160, 210)):
|
||||
buf = io.BytesIO()
|
||||
Image.new("RGBA", (w, h), color).save(buf, format="PNG")
|
||||
return base64.b64encode(buf.getvalue()).decode("ascii")
|
||||
|
||||
|
||||
def test_prepare_local_edit_caps_large_input():
|
||||
src, mask, mw, mh, orig = gr._prepare_local_edit(_png_b64(4000, 3000), None, 768)
|
||||
assert orig == (4000, 3000)
|
||||
assert max(mw, mh) <= 768 and mw % 16 == 0 and mh % 16 == 0
|
||||
assert mask # a full-white mask is synthesised when none is supplied
|
||||
assert Image.open(io.BytesIO(src)).size == (mw, mh)
|
||||
|
||||
|
||||
def test_prepare_local_edit_leaves_small_input_unchanged():
|
||||
_src, _mask, mw, mh, orig = gr._prepare_local_edit(_png_b64(96, 96), _png_b64(96, 96, (255, 255, 255)), 768)
|
||||
assert orig == (96, 96) and (mw, mh) == (96, 96)
|
||||
|
||||
|
||||
def test_restore_local_b64_resizes_back_to_original():
|
||||
raw = _png_b64(768, 576)
|
||||
out = gr._restore_local_b64(raw, (4000, 3000), (768, 576))
|
||||
assert Image.open(io.BytesIO(base64.b64decode(out))).size == (4000, 3000)
|
||||
# a same-size result is returned untouched
|
||||
assert gr._restore_local_b64(raw, (768, 576), (768, 576)) == raw
|
||||
|
||||
|
||||
def test_inpaint_offloads_cpu_work_to_a_thread():
|
||||
src = (ROOT / "routes/gallery/gallery_routes.py").read_text()
|
||||
assert "await asyncio.to_thread(\n _prepare_local_edit" in src
|
||||
assert "asyncio.to_thread(\n _restore_local_b64" in src
|
||||
|
||||
|
||||
def test_app_exposes_busy_count_and_graceful_shutdown():
|
||||
src = (ROOT / "app.py").read_text()
|
||||
assert "busy_image_ops" in src
|
||||
assert '@app.get("/api/busy")' in src
|
||||
assert "_HEAVY_IMAGE_OP_PREFIXES" in src
|
||||
assert "timeout_graceful_shutdown" in src
|
||||
assert "activity.begin_image_op" in src and "activity.end_image_op" in src
|
||||
@@ -33,6 +33,38 @@ def test_native_generation_respects_admin_disable(monkeypatch):
|
||||
assert 'disabled' in result['error']
|
||||
|
||||
|
||||
@pytest.mark.parametrize('editing', [False, True])
|
||||
def test_incognito_images_never_write_gallery_or_files(image_transport, editing):
|
||||
configure, source, requests, rows = image_transport
|
||||
configure('https://openrouter.ai/api/v1')
|
||||
if editing:
|
||||
result = asyncio.run(ai_interaction.do_edit_image(
|
||||
'Private edit', str(source), model_spec='openai/gpt-5-image',
|
||||
owner='alice', session_id='private-chat', incognito=True))
|
||||
else:
|
||||
result = asyncio.run(ai_interaction.do_generate_image(
|
||||
'Private image\nopenai/gpt-5-image', owner='alice',
|
||||
session_id='private-chat', incognito=True))
|
||||
assert 'error' not in result
|
||||
assert result['image_url'].startswith('data:image/png;base64,')
|
||||
assert base64.b64decode(result['image_url'].split(',', 1)[1]) == source.read_bytes()
|
||||
assert not result['image_id']
|
||||
assert not rows
|
||||
assert not (source.parent / 'generated').exists()
|
||||
|
||||
|
||||
def test_incognito_url_image_is_not_persisted(image_transport, monkeypatch):
|
||||
configure, source, requests, rows = image_transport
|
||||
configure('https://api.openai.com/v1', body={'data': [{'url': 'https://example.com/output.png'}]})
|
||||
monkeypatch.setattr(httpx, 'get', lambda *args, **kwargs: httpx.Response(200, content=source.read_bytes()))
|
||||
monkeypatch.setattr('src.url_safety.check_outbound_url', lambda *args, **kwargs: (True, ''))
|
||||
result = asyncio.run(ai_interaction.do_generate_image(
|
||||
'Private image\nopenai/gpt-5-image', owner='alice', incognito=True))
|
||||
assert result['image_url'].startswith('data:image/png;base64,')
|
||||
assert not rows
|
||||
assert not (source.parent / 'generated').exists()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def image_transport(monkeypatch, tmp_path):
|
||||
rows = []
|
||||
|
||||
@@ -281,7 +281,7 @@ def test_open_imap_connection_shuts_down_on_starttls_failure(monkeypatch):
|
||||
))
|
||||
conn.starttls = MagicMock(side_effect=RuntimeError("STARTTLS rejected"))
|
||||
|
||||
monkeypatch.setattr(helpers, "_PolicyIMAP4", lambda *a, **kw: conn)
|
||||
monkeypatch.setattr(helpers.imaplib, "IMAP4", lambda *a, **kw: conn)
|
||||
|
||||
raised = False
|
||||
try:
|
||||
|
||||
@@ -0,0 +1,503 @@
|
||||
"""Scoped ``photo_editor`` capability family: routing, tools, and the
|
||||
Quick Edit agent endpoint.
|
||||
|
||||
The family is context-only: it is offered to the model ONLY when an editor
|
||||
session / edit-this-photo request explicitly selects it, never inferred from
|
||||
message text and never part of the general tool list.
|
||||
"""
|
||||
import inspect
|
||||
|
||||
import pytest
|
||||
|
||||
from src.turn_contract import (
|
||||
FAMILY_TOOLS,
|
||||
requested_capabilities,
|
||||
resolve_full_inventory_contract,
|
||||
resolve_turn_contract,
|
||||
)
|
||||
from src.tool_schemas import FUNCTION_TOOL_SCHEMAS
|
||||
from src.tool_policy import ToolPolicy
|
||||
from src.tool_types import TOOL_TAGS
|
||||
from src.agent_tools import TOOL_HANDLERS
|
||||
from src.tool_index import BUILTIN_TOOL_DESCRIPTIONS
|
||||
from src.tool_capabilities import KNOWN_CAPABILITY_TOOLS
|
||||
from src.tool_security import PLAN_MODE_READONLY_TOOLS, _PLAN_MODE_KNOWN_MUTATORS
|
||||
|
||||
PHOTO_TOOLS = frozenset(
|
||||
{"plan_edit", "run_editor_command", "get_editor_screenshot", "ai_edit_image"}
|
||||
)
|
||||
FAMILY = FAMILY_TOOLS["photo_editor"]
|
||||
CORE_NAMES = frozenset({"bash", "python", "web_search", "web_fetch", "write_file"})
|
||||
|
||||
|
||||
def _schema_names():
|
||||
return {s["function"]["name"] for s in FUNCTION_TOOL_SCHEMAS}
|
||||
|
||||
|
||||
def test_family_contains_scoped_tools_only():
|
||||
assert FAMILY == PHOTO_TOOLS | {"generate_image"}
|
||||
|
||||
|
||||
def test_contract_offers_only_editor_tools_and_ask_user():
|
||||
contract = resolve_turn_contract(
|
||||
capabilities={"photo_editor"},
|
||||
schemas=FUNCTION_TOOL_SCHEMAS,
|
||||
policy=ToolPolicy(),
|
||||
required_capabilities={"photo_editor"},
|
||||
)
|
||||
assert contract.offered == FAMILY | {"ask_user", "update_plan"}
|
||||
assert not (CORE_NAMES & contract.offered)
|
||||
|
||||
|
||||
def test_photo_editor_is_never_inferred_from_message_text():
|
||||
for message in (
|
||||
"make it black and white",
|
||||
"generate an image of a cat",
|
||||
"remove the person on the left",
|
||||
"add a full moon in the sky",
|
||||
"blur the background behind the person",
|
||||
):
|
||||
caps = requested_capabilities(message)
|
||||
assert "photo_editor" not in caps, message
|
||||
|
||||
|
||||
def test_generate_image_keeps_its_own_family_without_leaking_photo_editor():
|
||||
caps = requested_capabilities("Create a note and generate an image of a cat")
|
||||
assert "image_generation" in caps
|
||||
assert "photo_editor" not in caps
|
||||
|
||||
|
||||
def test_editor_tools_are_registered_everywhere():
|
||||
schema_names = _schema_names()
|
||||
for name in PHOTO_TOOLS:
|
||||
assert name in TOOL_TAGS, name
|
||||
assert name in TOOL_HANDLERS, name
|
||||
assert name in BUILTIN_TOOL_DESCRIPTIONS, name
|
||||
assert name in KNOWN_CAPABILITY_TOOLS, name
|
||||
assert name in schema_names, name
|
||||
assert "generate_image" in FAMILY
|
||||
|
||||
|
||||
def test_editor_tools_plan_mode_partition():
|
||||
assert "get_editor_screenshot" in PLAN_MODE_READONLY_TOOLS
|
||||
for name in ("plan_edit", "run_editor_command", "ai_edit_image"):
|
||||
assert name in _PLAN_MODE_KNOWN_MUTATORS, name
|
||||
|
||||
|
||||
def test_full_inventory_contains_the_editor_tools():
|
||||
preview = resolve_full_inventory_contract(
|
||||
schemas=FUNCTION_TOOL_SCHEMAS, policy=ToolPolicy()
|
||||
)
|
||||
assert FAMILY <= set(preview.offered)
|
||||
|
||||
|
||||
def test_scope_preview_contract_keeps_only_the_editor_family():
|
||||
from dataclasses import replace
|
||||
|
||||
from src.clean_agent_preview import scope_preview_contract
|
||||
|
||||
preview = resolve_full_inventory_contract(
|
||||
schemas=FUNCTION_TOOL_SCHEMAS, policy=ToolPolicy()
|
||||
)
|
||||
routed = resolve_turn_contract(
|
||||
capabilities={"photo_editor"},
|
||||
schemas=FUNCTION_TOOL_SCHEMAS,
|
||||
policy=ToolPolicy(),
|
||||
required_capabilities={"photo_editor"},
|
||||
)
|
||||
scoped = scope_preview_contract(
|
||||
replace(preview, selection_mode="full_compact_experiment"),
|
||||
routed,
|
||||
frozenset({"photo_editor"}),
|
||||
extra_tools=frozenset({"ask_user"}),
|
||||
)
|
||||
assert FAMILY <= set(scoped.offered)
|
||||
assert "ask_user" in scoped.offered
|
||||
assert not (CORE_NAMES & set(scoped.offered))
|
||||
|
||||
|
||||
def test_editor_agent_endpoint_reports_scoped_tools():
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import routes.editor_agent_routes as eaar
|
||||
|
||||
eaar.set_editor_planner(lambda messages, owner: {"reply": "ok", "steps": [], "done": True})
|
||||
try:
|
||||
app = FastAPI()
|
||||
app.include_router(eaar.setup_editor_agent_routes())
|
||||
client = TestClient(app)
|
||||
resp = client.post(
|
||||
"/api/editor/agent",
|
||||
json={"prompt": "hello there", "session": "s1"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["ok"] is True
|
||||
assert PHOTO_TOOLS <= set(data["tools_offered"])
|
||||
assert "generate_image" in data["tools_offered"]
|
||||
assert not (CORE_NAMES & set(data["tools_offered"]))
|
||||
assert data["trace"][0]["type"] == "session"
|
||||
finally:
|
||||
eaar.set_editor_planner(None)
|
||||
|
||||
|
||||
async def _route_contract(monkeypatch, message, *, attachment_meta=()):
|
||||
from routes import chat_routes
|
||||
from src import tool_security
|
||||
from tests.test_foreground_model_routing import _RouteRequest, _chat_stream_endpoint
|
||||
|
||||
with monkeypatch.context() as route_patch:
|
||||
endpoint = _chat_stream_endpoint(
|
||||
route_patch, "agent", {},
|
||||
session_model="odysseus-qwen3.5-tools-pre-heretic",
|
||||
attachment_meta=attachment_meta,
|
||||
)
|
||||
route_patch.setattr(
|
||||
chat_routes, "coerce_message_and_session",
|
||||
lambda *args, **kwargs: (message, "session-1"),
|
||||
)
|
||||
route_patch.setattr(
|
||||
tool_security, "owner_is_admin_or_single_user", lambda owner: True,
|
||||
)
|
||||
observed = []
|
||||
|
||||
async def capture_agent(*args, **kwargs):
|
||||
observed.append(kwargs["turn_contract"])
|
||||
yield 'data: {"delta":"ok"}\n\n'
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
route_patch.setattr(chat_routes, "stream_agent_loop", capture_agent)
|
||||
request = _RouteRequest("agent")
|
||||
request._form.update({
|
||||
"message": message, "compare_mode": "false",
|
||||
"allow_web_search": "false", "allow_bash": "false",
|
||||
})
|
||||
response = await endpoint(request)
|
||||
async for _ in response.body_iterator:
|
||||
pass
|
||||
assert len(observed) == 1
|
||||
return observed[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_normal_chat_turn_offers_no_editor_tools(monkeypatch):
|
||||
for message in ("what is the capital of France", "write a haiku"):
|
||||
contract = await _route_contract(monkeypatch, message)
|
||||
assert "photo_editor" not in contract.capabilities, message
|
||||
assert not (PHOTO_TOOLS & set(contract.offered)), message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_attached_photo_edit_selects_photo_editor_family(monkeypatch):
|
||||
contract = await _route_contract(
|
||||
monkeypatch,
|
||||
"make this black and white",
|
||||
attachment_meta=[{"id": "att1", "mime": "image/png", "name": "x.png"}],
|
||||
)
|
||||
assert contract.capabilities == {"photo_editor"}
|
||||
assert set(contract.offered) == FAMILY | {"ask_user"}
|
||||
assert not (CORE_NAMES & set(contract.offered))
|
||||
|
||||
|
||||
def test_chat_context_guides_open_or_absent_editor():
|
||||
from pathlib import Path
|
||||
|
||||
src = (Path(__file__).resolve().parent.parent / "routes" / "chat_routes.py").read_text(
|
||||
encoding="utf-8"
|
||||
)
|
||||
assert "The image editor is open (session " in src
|
||||
assert "no image editor is open" in src
|
||||
assert "the attached image in the editor (Quick Edit)" in src
|
||||
|
||||
|
||||
def test_generative_intent_classification():
|
||||
from routes.editor_agent_routes import _generative_kind
|
||||
|
||||
assert _generative_kind("remove the person on the left") == "remove"
|
||||
assert _generative_kind("erase the sign in the background") == "remove"
|
||||
assert _generative_kind("add a full moon in the sky") == "add"
|
||||
assert _generative_kind("put a bird on the fence") == "add"
|
||||
assert _generative_kind("replace the sky with a sunset") == "replace"
|
||||
assert _generative_kind("change the background to night") == "replace"
|
||||
# Non-generative edits must stay deterministic.
|
||||
assert _generative_kind("make it black and white") is None
|
||||
assert _generative_kind("make it brighter") is None
|
||||
assert _generative_kind("add a vignette") is None
|
||||
assert _generative_kind("remove the background") is None
|
||||
assert _generative_kind("sharpen the image") is None
|
||||
|
||||
|
||||
def test_generative_step_uses_ai_content_tools():
|
||||
from routes.editor_agent_routes import _generative_step
|
||||
|
||||
assert _generative_step("remove", "remove the person")["op"] == "ai_edit_image"
|
||||
assert _generative_step("replace", "replace the sky")["op"] == "ai_edit_image"
|
||||
# Adding content to an existing photo is an instruction edit, not generate_image.
|
||||
assert _generative_step("add", "add a moon")["op"] == "ai_edit_image"
|
||||
assert _generative_step("add", "add a moon")["params"]["mode"] == "instruction"
|
||||
# generate_image only for a brand-new image.
|
||||
assert _generative_step("add", "create a new image of a moon")["op"] == "generate_image"
|
||||
|
||||
|
||||
def test_generative_mode_instruction_vs_masked():
|
||||
from routes.editor_agent_routes import _generative_mode
|
||||
|
||||
# Masked only for an explicit regional removal of a clearly identified object.
|
||||
assert _generative_mode("remove", "remove the person on the left") == "masked"
|
||||
# Diffuse cleanup, replacement, season/style and additions use instruction mode.
|
||||
assert _generative_mode("remove", "remove the distractions") == "instruction"
|
||||
assert _generative_mode("replace", "make the background nicer") == "instruction"
|
||||
assert _generative_mode("add", "add a full moon in the sky") == "instruction"
|
||||
|
||||
|
||||
def test_ensure_generative_replaces_deterministic_approximations():
|
||||
from routes.editor_agent_routes import _ensure_generative
|
||||
|
||||
bad = [{"op": "selection.rectangle", "params": {}}, {"op": "adjustments.apply", "params": {}}]
|
||||
removal = _ensure_generative("remove the person on the left", list(bad))
|
||||
# A mask selection step now precedes the masked AI edit.
|
||||
assert removal[-1]["op"] == "ai_edit_image"
|
||||
assert removal[-1]["params"]["mode"] == "masked"
|
||||
assert all(s["op"].startswith("selection.") for s in removal[:-1])
|
||||
# Adding an object to an existing photo is an instruction edit (no selection).
|
||||
addition = _ensure_generative("add a full moon in the sky", list(bad))
|
||||
assert addition[-1]["op"] == "ai_edit_image"
|
||||
assert addition[-1]["params"]["mode"] == "instruction"
|
||||
# Non-generative plans are left untouched.
|
||||
assert _ensure_generative("make it black and white", list(bad)) == bad
|
||||
|
||||
|
||||
def test_fallback_plan_routes_generative_intents():
|
||||
import routes.editor_agent_routes as eaar
|
||||
|
||||
removal = eaar._fallback_plan("remove the person on the left")
|
||||
assert removal["steps"][-1]["op"] == "ai_edit_image"
|
||||
assert removal["steps"][-1]["params"]["mode"] == "masked"
|
||||
addition = eaar._fallback_plan("add a full moon in the sky")
|
||||
assert addition["steps"][-1]["op"] == "ai_edit_image"
|
||||
assert addition["steps"][-1]["params"]["mode"] == "instruction"
|
||||
plain = eaar._fallback_plan("make it black and white")
|
||||
assert plain["steps"][0]["op"] == "adjustments.apply"
|
||||
|
||||
|
||||
def test_provider_message_is_actionable():
|
||||
from src.ai_interaction import IMAGE_PROVIDER_REQUIRED_MESSAGE
|
||||
|
||||
assert "image generation provider" in IMAGE_PROVIDER_REQUIRED_MESSAGE
|
||||
assert "Settings > Models" in IMAGE_PROVIDER_REQUIRED_MESSAGE
|
||||
|
||||
|
||||
def test_image_provider_respects_disabled_setting(monkeypatch):
|
||||
import src.settings as settings
|
||||
|
||||
monkeypatch.setattr(settings, "load_settings", lambda: {"image_gen_enabled": False})
|
||||
from src.ai_interaction import image_editing_provider_available
|
||||
|
||||
assert image_editing_provider_available(None) is False
|
||||
|
||||
|
||||
def _agent_client():
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import routes.editor_agent_routes as eaar
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(eaar.setup_editor_agent_routes())
|
||||
return eaar, TestClient(app)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"prompt",
|
||||
["remove the person on the left", "add a full moon in the sky"],
|
||||
)
|
||||
def test_editor_agent_blocks_generative_without_provider(monkeypatch, prompt):
|
||||
"""No image provider -> zero canvas changes + a plain provider message."""
|
||||
eaar, client = _agent_client()
|
||||
# Planner proposes an unrelated deterministic edit, as the buggy run did.
|
||||
eaar.set_editor_planner(
|
||||
lambda messages, owner: {"reply": "done", "steps": [
|
||||
{"op": "selection.rectangle", "params": {}},
|
||||
{"op": "adjustments.apply", "params": {}},
|
||||
]}
|
||||
)
|
||||
monkeypatch.setattr(eaar, "_image_provider_available", lambda owner: False)
|
||||
try:
|
||||
resp = client.post(
|
||||
"/api/editor/agent", json={"prompt": prompt, "session": "s1"}
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["needs_provider"] is True
|
||||
assert data["blocked"] is True
|
||||
assert "image generation provider" in data["message"]
|
||||
assert data["image_provider"] is False
|
||||
assert data["results"] == []
|
||||
# Nothing was executed through the bridge (zero canvas changes).
|
||||
assert not any(t.get("type") == "step" for t in data["trace"])
|
||||
assert any(t.get("reason") == "image_provider_unavailable" for t in data["trace"])
|
||||
# The proposed deterministic plan was replaced by the correct AI op.
|
||||
ops = [s["op"] for s in data["steps"]]
|
||||
assert ops[-1] in ("ai_edit_image", "generate_image")
|
||||
finally:
|
||||
eaar.set_editor_planner(None)
|
||||
|
||||
|
||||
def test_editor_agent_prefers_ai_op_over_deterministic_when_provider_present(monkeypatch):
|
||||
eaar, client = _agent_client()
|
||||
eaar.set_editor_planner(
|
||||
lambda messages, owner: {"reply": "done", "steps": [
|
||||
{"op": "selection.rectangle", "params": {}},
|
||||
{"op": "adjustments.apply", "params": {}},
|
||||
]}
|
||||
)
|
||||
monkeypatch.setattr(eaar, "_image_provider_available", lambda owner: True)
|
||||
|
||||
async def _instant_wait_for(awaitable, *args, **kwargs):
|
||||
# The plan now starts with a mask-selection bridge step; resolve it
|
||||
# immediately instead of waiting for a (non-existent) editor subscriber.
|
||||
try:
|
||||
awaitable.close()
|
||||
except Exception:
|
||||
pass
|
||||
return {"ok": True, "result": {"width": 100, "height": 100}}
|
||||
|
||||
monkeypatch.setattr(eaar.asyncio, "wait_for", _instant_wait_for)
|
||||
|
||||
calls = []
|
||||
|
||||
class _FakeAiEdit:
|
||||
async def execute(self, content, ctx):
|
||||
calls.append(content)
|
||||
return "ai_edit_image: fake", {"output": "ok", "exit_code": 0}
|
||||
|
||||
class _FakeGen:
|
||||
async def execute(self, content, ctx):
|
||||
calls.append(content)
|
||||
return "generate_image: fake", {"output": "ok", "exit_code": 0}
|
||||
|
||||
import src.agent_tools.editor_tools as et
|
||||
monkeypatch.setattr(et, "AiEditImageTool", _FakeAiEdit)
|
||||
monkeypatch.setattr(et, "GenerateImageLayerTool", _FakeGen)
|
||||
try:
|
||||
resp = client.post(
|
||||
"/api/editor/agent",
|
||||
json={"prompt": "remove the person on the left", "session": "s1"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["ok"] is True
|
||||
ops = [s["op"] for s in data["steps"]]
|
||||
assert ops[-1] == "ai_edit_image"
|
||||
assert any(r.get("op") == "ai_edit_image" for r in data["results"])
|
||||
assert calls, "the AI edit tool should have been invoked"
|
||||
assert "adjustments.apply" not in ops
|
||||
finally:
|
||||
eaar.set_editor_planner(None)
|
||||
|
||||
|
||||
def test_editor_prompts_advertise_canonical_shapes():
|
||||
import routes.editor_agent_routes as eaar
|
||||
|
||||
canonical = eaar._CANONICAL_GUIDANCE
|
||||
assert "brightness-contrast" in canonical and "render-clouds" in canonical
|
||||
assert canonical in eaar._SYSTEM_PROMPT
|
||||
assert canonical in eaar._CHECK_SYSTEM_PROMPT
|
||||
assert "_CANONICAL_GUIDANCE" in inspect.getsource(eaar.setup_editor_agent_routes)
|
||||
|
||||
|
||||
def test_run_editor_command_schema_advertises_canonical_shapes():
|
||||
from src.tool_schemas import FUNCTION_TOOL_SCHEMAS
|
||||
|
||||
schema = next(s for s in FUNCTION_TOOL_SCHEMAS
|
||||
if s["function"]["name"] == "run_editor_command")
|
||||
desc = schema["function"]["description"]
|
||||
assert "brightness-contrast" in desc and "render-clouds" in desc
|
||||
assert "vignette" in desc # names the banned invented ids
|
||||
|
||||
|
||||
def test_photo_editor_chat_without_canvas_asks_to_open_editor():
|
||||
from src import clean_agent_preview
|
||||
|
||||
source = inspect.getsource(clean_agent_preview.stream_preview)
|
||||
assert "not open on the editor canvas" in source
|
||||
assert "'photo_editor' in (getattr(turn_contract, 'active_capabilities'" in source
|
||||
|
||||
|
||||
def test_chat_photo_edit_without_open_editor_short_circuits(monkeypatch):
|
||||
"""A photo-edit chat turn with no open editor must not run the tool loop.
|
||||
|
||||
Regression: offering the editor tools with no canvas made the model loop
|
||||
on plan_edit/run_editor_command until the round cap and end with an empty
|
||||
answer. The runtime now answers deterministically and asks the user to
|
||||
open the image in the editor.
|
||||
"""
|
||||
import asyncio
|
||||
import json as _json
|
||||
from dataclasses import replace
|
||||
|
||||
import src.clean_agent_preview as module
|
||||
|
||||
contract = replace(
|
||||
resolve_full_inventory_contract(schemas=FUNCTION_TOOL_SCHEMAS, policy=ToolPolicy()),
|
||||
active_capabilities=frozenset({"photo_editor"}),
|
||||
)
|
||||
provider_calls = []
|
||||
|
||||
class _Client:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
pass
|
||||
|
||||
def stream(self, *args, **kwargs):
|
||||
provider_calls.append(kwargs.get("json"))
|
||||
raise AssertionError("provider must not be called when no editor is open")
|
||||
|
||||
monkeypatch.setattr(module.httpx, "AsyncClient", _Client)
|
||||
|
||||
async def _run():
|
||||
return [
|
||||
chunk
|
||||
async for chunk in module.stream_preview(
|
||||
endpoint_url="http://test",
|
||||
model="test",
|
||||
headers={},
|
||||
messages=[{"role": "user", "content": "make this black and white"}],
|
||||
turn_contract=contract,
|
||||
session_id="test",
|
||||
owner="test",
|
||||
disabled_tools=set(),
|
||||
tool_policy=ToolPolicy(),
|
||||
client_runtime_context={"web_ui_state": {"photo_editor_open": False}},
|
||||
)
|
||||
]
|
||||
|
||||
chunks = asyncio.run(_run())
|
||||
events = [
|
||||
_json.loads(chunk[6:])
|
||||
for chunk in chunks
|
||||
if chunk.startswith("data: ") and not chunk.startswith("data: [DONE]")
|
||||
]
|
||||
|
||||
assert provider_calls == []
|
||||
contract_event = next(e for e in events if e.get("type") == "turn_contract")
|
||||
assert "photo_editor" in contract_event["active_capabilities"]
|
||||
final = next(e for e in events if e.get("type") == "final_response")
|
||||
lowered = final["content"].lower()
|
||||
assert "editor" in lowered and "open" in lowered
|
||||
# One-click affordance contract for the chat UI.
|
||||
assert final["photo_editor_action"] == {
|
||||
"kind": "open_and_apply",
|
||||
"prompt": "make this black and white",
|
||||
}
|
||||
assert not any(e.get("type") == "tool_start" for e in events)
|
||||
|
||||
|
||||
if __name__ == "__main__": # pragma: no cover
|
||||
raise SystemExit(pytest.main([__file__, "-q"]))
|
||||
@@ -0,0 +1,642 @@
|
||||
"""Unit tests for the shared photo-editor style guide and its server wiring.
|
||||
|
||||
The guide (``src/editor_style_guide.py``) is the single source of truth for
|
||||
intent -> ops routing, calibrated strengths and the low-key / temperature /
|
||||
portrait / ambiguity rules. It feeds both the server planner and (via
|
||||
``GET /api/editor/style-guide``) the browser fallback.
|
||||
"""
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import routes.editor_agent_routes as eaar
|
||||
from src import editor_style_guide as guide
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_planner():
|
||||
eaar.set_editor_planner(None)
|
||||
yield
|
||||
eaar.set_editor_planner(None)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def client():
|
||||
app = FastAPI()
|
||||
app.include_router(eaar.setup_editor_agent_routes())
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def _ops(steps):
|
||||
return [s["op"] for s in steps]
|
||||
|
||||
|
||||
def test_new_ops_are_in_the_server_op_catalog():
|
||||
# If a new retouch/selection op is missing from OP_CATALOG, _parse_plan
|
||||
# filters a valid plan step out and the edit silently disappears.
|
||||
for name in guide.NEW_OP_NAMES:
|
||||
assert name in eaar.OP_CATALOG
|
||||
for name in ("photo.autoTone", "retouch.skin", "retouch.teeth",
|
||||
"selection.face", "selection.background"):
|
||||
assert name in eaar._OP_SET
|
||||
|
||||
|
||||
def test_portrait_skin_intent_is_a_multi_step_plan():
|
||||
plan = guide.fallback_plan("can you fix the skin a bit, look my best")
|
||||
assert plan is not None
|
||||
# No leading whole-image auto tone: it washes out the face. Go straight to a
|
||||
# gentle face retouch, then clear the selection.
|
||||
assert _ops(plan["steps"]) == [
|
||||
"selection.face", "retouch.skin", "selection.deselect",
|
||||
]
|
||||
skin = plan["steps"][1]["params"]
|
||||
assert 0 < skin["smooth"] <= 14
|
||||
assert 0 < skin["opacity"] <= 0.45
|
||||
# No global lift anywhere in the portrait chain (client never brightens skin).
|
||||
assert not any("autoTone" in op or "brightness" in op or "exposure" in op
|
||||
for op in _ops(plan["steps"]))
|
||||
|
||||
|
||||
def test_portrait_intent_does_not_lead_with_auto_tone():
|
||||
for prompt in ("make me look better", "fix the skin", "fix my teeth",
|
||||
"smooth my complexion", "portrait touch up"):
|
||||
plan = guide.fallback_plan(prompt)
|
||||
assert plan is not None, prompt
|
||||
assert "photo.autoTone" not in _ops(plan["steps"]), prompt
|
||||
|
||||
|
||||
def test_teeth_intent_selects_face_then_whitens():
|
||||
plan = guide.fallback_plan("fix my teeth")
|
||||
assert _ops(plan["steps"]) == ["selection.face", "retouch.teeth"]
|
||||
assert plan["steps"][1]["params"]["desat"] > 0
|
||||
|
||||
|
||||
def test_background_intent_is_selection_then_masked_blur():
|
||||
plan = guide.fallback_plan("blur the background / make the background nicer")
|
||||
ops = _ops(plan["steps"])
|
||||
assert ops == ["selection.background", "filters.applyMasked"]
|
||||
assert plan["steps"][1]["params"]["id"] == "blur-gaussian"
|
||||
|
||||
|
||||
def test_temperature_sign_is_cool_negative_warm_positive():
|
||||
cool = guide.fallback_plan("make it moodier")
|
||||
assert cool["steps"][0]["params"] == {
|
||||
"type": "white-balance", "params": {"temperature": -25}}
|
||||
warm = guide.fallback_plan("make it warmer, golden summer feel")
|
||||
assert warm["steps"][0]["params"]["params"]["temperature"] == 25
|
||||
|
||||
|
||||
def test_lighting_intent_uses_auto_tone_then_single_shadow_lift():
|
||||
plan = guide.fallback_plan("fix the lighting, it's too dark")
|
||||
assert _ops(plan["steps"]) == ["photo.autoTone", "adjustments.apply"]
|
||||
inner = plan["steps"][1]["params"]
|
||||
assert inner["type"] == "shadows-highlights"
|
||||
assert inner["params"]["highlights"] <= 0
|
||||
# The sanctioned low-key pair, never a stack of unspecified lifts.
|
||||
assert len(plan["steps"]) == 2
|
||||
|
||||
|
||||
def test_ambiguous_prompt_asks_the_user_with_concrete_options():
|
||||
plan = guide.fallback_plan("clean this up")
|
||||
assert plan["steps"] == []
|
||||
ask = plan["ask_user"]
|
||||
assert ask["question"]
|
||||
assert 2 <= len(ask["options"]) <= 3
|
||||
assert all(o.get("label") for o in ask["options"])
|
||||
assert guide.is_ambiguous("clean this up") is True
|
||||
# A concrete defect is not ambiguous.
|
||||
assert guide.is_ambiguous("fix the lighting") is False
|
||||
|
||||
|
||||
def test_collapse_global_lifts_replaces_a_stack():
|
||||
# With an auto tone present, extra brightness/exposure lifts are dropped
|
||||
# (the auto tone is the single sanctioned lift).
|
||||
with_auto = [
|
||||
{"op": "photo.autoTone", "params": {}},
|
||||
{"op": "adjustments.apply",
|
||||
"params": {"type": "brightness-contrast", "params": {"brightness": 30}}},
|
||||
{"op": "adjustments.apply",
|
||||
"params": {"type": "exposure", "params": {"exposure": 20}}},
|
||||
]
|
||||
out = guide.collapse_global_lifts(with_auto)
|
||||
assert len(out) == 1
|
||||
assert out[0]["op"] == "photo.autoTone"
|
||||
# Without an auto tone, a stack collapses to one safe shadows-highlights.
|
||||
no_auto = [
|
||||
{"op": "adjustments.apply",
|
||||
"params": {"type": "brightness-contrast", "params": {"brightness": 30}}},
|
||||
{"op": "adjustments.apply",
|
||||
"params": {"type": "exposure", "params": {"exposure": 20}}},
|
||||
]
|
||||
out2 = guide.collapse_global_lifts(no_auto)
|
||||
assert len(out2) == 1
|
||||
assert out2[0]["params"]["type"] == "shadows-highlights"
|
||||
assert out2[0]["params"]["params"]["highlights"] <= 0
|
||||
# The sanctioned low-key pair (auto tone + one shadows-highlights) is kept.
|
||||
pair = [
|
||||
{"op": "photo.autoTone", "params": {}},
|
||||
{"op": "adjustments.apply",
|
||||
"params": {"type": "shadows-highlights",
|
||||
"params": {"shadows": 30, "highlights": -15}}},
|
||||
]
|
||||
assert guide.collapse_global_lifts(pair) == pair
|
||||
|
||||
|
||||
def test_collapse_global_lifts_clamps_a_single_brightness_lift():
|
||||
steps = [{"op": "adjustments.apply",
|
||||
"params": {"type": "brightness-contrast", "params": {"brightness": 80}}}]
|
||||
out = guide.collapse_global_lifts(steps)
|
||||
assert out[0]["params"]["params"]["brightness"] <= 25
|
||||
|
||||
|
||||
def test_llm_plan_parse_collapses_lifts_and_keeps_new_ops():
|
||||
raw = (
|
||||
'{"reply": "x", "steps": ['
|
||||
'{"op": "retouch.skin", "params": {"smooth": 20}},'
|
||||
'{"op": "photo.autoTone", "params": {}},'
|
||||
'{"op": "adjustments.apply", "params": {"type": "exposure", "params": {"exposure": 25}}}'
|
||||
']}'
|
||||
)
|
||||
parsed = eaar._parse_plan(raw)
|
||||
ops = _ops(parsed["steps"])
|
||||
assert "retouch.skin" in ops
|
||||
# Every global lift that trails retouch.skin is dropped (the client guard
|
||||
# "never brighten skin"), so nothing but the retouch step remains.
|
||||
assert ops == ["retouch.skin"]
|
||||
|
||||
|
||||
def test_editor_agent_returns_ask_user_for_ambiguous_prompt(client):
|
||||
eaar.set_editor_planner(lambda messages, owner: '{"reply": "ok", "steps": []}')
|
||||
resp = client.post("/api/editor/agent", json={"prompt": "clean this up"})
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body.get("ask_user_pending") is True
|
||||
assert body["steps"] == []
|
||||
assert 2 <= len(body["ask_user"]["options"]) <= 3
|
||||
|
||||
|
||||
def test_editor_plan_returns_ask_user_for_ambiguous_prompt(client):
|
||||
eaar.set_editor_planner(lambda messages, owner: None)
|
||||
resp = client.post("/api/editor/plan", json={"prompt": "clean this up"})
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["steps"] == []
|
||||
assert body.get("ask_user")
|
||||
|
||||
|
||||
def test_style_guide_endpoint_serves_the_shared_source(client):
|
||||
resp = client.get("/api/editor/style-guide")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["version"] == guide.STYLE_GUIDE_VERSION
|
||||
assert {o["op"] for o in data["new_ops"]} == set(guide.NEW_OP_NAMES)
|
||||
assert data["intents"]
|
||||
assert "temperature" in data["rules"]["temperature_sign"].lower() or \
|
||||
"negative" in data["rules"]["temperature_sign"]
|
||||
# The prompt guidance states the sign rule explicitly.
|
||||
assert "NEGATIVE" in data["guidance"] and "POSITIVE" in data["guidance"]
|
||||
|
||||
|
||||
def test_prompts_embed_the_style_guidance():
|
||||
assert guide.guidance() in eaar._SYSTEM_PROMPT
|
||||
assert guide.guidance() in eaar._CHECK_SYSTEM_PROMPT
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prompt", [
|
||||
"make it pretty", "make it pop", "professional", "film look",
|
||||
])
|
||||
def test_generic_look_intent_is_a_decisive_multi_step_plan(prompt):
|
||||
plan = guide.fallback_plan(prompt)
|
||||
assert plan is not None
|
||||
ops = _ops(plan["steps"])
|
||||
assert ops[0] == "photo.autoTone"
|
||||
assert len(ops) >= 3
|
||||
# A visible vibrance lift and a local-contrast filter are present.
|
||||
assert any(s["op"] == "adjustments.apply" and s["params"].get("type") == "vibrance"
|
||||
for s in plan["steps"])
|
||||
assert any(s["op"] == "filters.apply" and s["params"].get("id") == "clarity"
|
||||
for s in plan["steps"])
|
||||
|
||||
|
||||
def test_restore_intent_never_large_exposure_lift():
|
||||
plan = guide.fallback_plan("restore this old photo")
|
||||
assert plan is not None
|
||||
ops = _ops(plan["steps"])
|
||||
assert ops[0] == "photo.autoTone"
|
||||
assert "filters.apply" in ops # sharpen
|
||||
# Colour recovery present.
|
||||
assert any(s["op"] == "adjustments.apply" and s["params"].get("type") == "white-balance"
|
||||
for s in plan["steps"])
|
||||
# No exposure lift at all, and the contrast lift keeps brightness neutral.
|
||||
assert not any(s["params"].get("type") == "exposure" for s in plan["steps"])
|
||||
for s in plan["steps"]:
|
||||
if s["params"].get("type") == "brightness-contrast":
|
||||
assert s["params"]["params"].get("brightness", 0) == 0
|
||||
|
||||
|
||||
def test_strengthen_promotes_a_timid_plan_to_the_recipe():
|
||||
timid = [{"op": "adjustments.apply",
|
||||
"params": {"type": "brightness-contrast", "params": {"brightness": 8}}}]
|
||||
promoted = guide.strengthen("make it pop", timid)
|
||||
assert promoted != timid
|
||||
assert _ops(promoted)[0] == "photo.autoTone"
|
||||
assert len(promoted) >= 3
|
||||
# A prompt with no recipe leaves the plan untouched.
|
||||
assert guide.strengthen("nonsense prompt zzz", timid) == timid
|
||||
|
||||
|
||||
def test_enforce_guards_removes_trailing_lift_after_skin():
|
||||
steps = [
|
||||
{"op": "selection.face", "params": {}},
|
||||
{"op": "retouch.skin", "params": {"smooth": 24}},
|
||||
{"op": "photo.autoTone", "params": {}},
|
||||
]
|
||||
guarded = guide.enforce_guards(steps)
|
||||
assert _ops(guarded) == ["selection.face", "retouch.skin"]
|
||||
stacked = [
|
||||
{"op": "adjustments.apply",
|
||||
"params": {"type": "brightness-contrast", "params": {"brightness": 20}}},
|
||||
{"op": "adjustments.apply",
|
||||
"params": {"type": "exposure", "params": {"exposure": 20}}},
|
||||
]
|
||||
collapsed = guide.enforce_guards(stacked)
|
||||
assert len(collapsed) == 1
|
||||
assert collapsed[0]["params"]["type"] == "shadows-highlights"
|
||||
|
||||
|
||||
def test_plan_route_follows_the_recipe_for_generic_look(client):
|
||||
import json as _json
|
||||
eaar.set_editor_planner(lambda messages, owner: _json.dumps({
|
||||
"reply": "ok",
|
||||
"steps": [{"op": "adjustments.apply",
|
||||
"params": {"type": "brightness-contrast", "params": {"brightness": 8}}}],
|
||||
}))
|
||||
resp = client.post("/api/editor/plan", json={"prompt": "make it pop"})
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["steps"][0]["op"] == "photo.autoTone"
|
||||
assert len(data["steps"]) >= 3
|
||||
assert data["planner"].endswith("+guide")
|
||||
|
||||
|
||||
def test_guidance_documents_look_and_restore_recipes():
|
||||
text = guide.guidance()
|
||||
low = text.lower()
|
||||
assert "clarity" in low and "vignette" in low
|
||||
assert "vibrance" in low
|
||||
assert "restore" in low and "exposure" in low
|
||||
assert "autoTone" in text
|
||||
|
||||
|
||||
def test_plan_special_format_crops_and_resizes():
|
||||
special = eaar._plan_special(
|
||||
"make a 1080x1350 Instagram post from this", {"width": 900, "height": 600}, False)
|
||||
assert special is not None
|
||||
steps, reply, note = special
|
||||
assert _ops(steps) == ["document.crop", "document.resizeCanvas"]
|
||||
assert "1080" in reply and "1350" in reply
|
||||
assert set(steps[0]["params"]) >= {"x", "y", "width", "height"}
|
||||
assert steps[1]["params"]["width"] == 1080 and steps[1]["params"]["height"] == 1350
|
||||
|
||||
|
||||
def test_plan_special_text_add_with_shadow_is_not_generative():
|
||||
assert eaar._generative_kind("add a title TOKYO with a drop shadow") is None
|
||||
special = eaar._plan_special("add a title TOKYO with a drop shadow", {}, False)
|
||||
assert special is not None
|
||||
steps, reply, note = special
|
||||
assert _ops(steps) == ["text.add", "layers.addStyle"]
|
||||
assert steps[0]["params"]["text"] == "TOKYO"
|
||||
assert steps[1]["params"]["type"] == "drop-shadow"
|
||||
assert "drop shadow" in reply.lower()
|
||||
|
||||
|
||||
def test_plan_special_extend_resizes_without_provider_generative_with():
|
||||
steps, reply, note = eaar._plan_special(
|
||||
"extend the canvas to 16:9", {"width": 900, "height": 600}, False)
|
||||
assert _ops(steps) == ["document.resizeCanvas"]
|
||||
assert "provider" in note.lower()
|
||||
steps2, _, note2 = eaar._plan_special(
|
||||
"extend the canvas to 16:9", {"width": 900, "height": 600}, True)
|
||||
assert _ops(steps2) == ["ai.generativeExpand"]
|
||||
assert note2 == ""
|
||||
|
||||
|
||||
def test_plan_special_psd_export():
|
||||
special = eaar._plan_special("save as PSD with layers", {}, False)
|
||||
assert special is not None and _ops(special[0]) == ["document.savePsd"]
|
||||
|
||||
|
||||
def test_background_replacement_is_generative_not_blur():
|
||||
assert eaar._generative_kind("put her on a beach") == "replace"
|
||||
assert eaar._generative_kind("remove the background and put her on a beach") == "replace"
|
||||
|
||||
|
||||
def test_verified_reply_corrects_a_false_size_claim():
|
||||
corrected = eaar._verified_reply(
|
||||
"make a 1080x1350 Instagram post",
|
||||
[{"op": "history.mark", "params": {}}],
|
||||
[{"ok": True}],
|
||||
{"width": 600, "height": 600},
|
||||
"Exported your 900x600 as a 1080x1350 Instagram post.",
|
||||
)
|
||||
assert "600" in corrected
|
||||
assert "still" in corrected.lower() or "didn't apply" in corrected.lower()
|
||||
passed = eaar._verified_reply(
|
||||
"make a 1080x1350 Instagram post",
|
||||
[{"op": "document.crop", "params": {}}, {"op": "document.resizeCanvas", "params": {}}],
|
||||
[{"ok": True}, {"ok": True}],
|
||||
{"width": 1080, "height": 1350},
|
||||
"Setting the canvas to 1080\u00d71350.",
|
||||
)
|
||||
assert "1080" in passed and "1350" in passed
|
||||
|
||||
|
||||
def test_plan_endpoint_returns_deterministic_format_steps(client, monkeypatch):
|
||||
eaar.set_editor_planner(lambda messages, owner: {"reply": "Exporting.", "steps": []})
|
||||
monkeypatch.setattr(eaar, "_image_provider_available", lambda owner: False)
|
||||
resp = client.post("/api/editor/plan", json={
|
||||
"prompt": "make a 1080x1350 Instagram post",
|
||||
"state": {"width": 900, "height": 600},
|
||||
})
|
||||
data = resp.json()
|
||||
assert data["planner"] == "deterministic"
|
||||
assert _ops(data["steps"])[:2] == ["document.crop", "document.resizeCanvas"]
|
||||
|
||||
|
||||
def test_style_guide_serves_format_text_psd_rules():
|
||||
rules = guide.as_json()["rules"]
|
||||
assert "format" in rules and "text" in rules and "psd" in rules
|
||||
assert guide.STYLE_GUIDE_VERSION >= 3
|
||||
|
||||
|
||||
# ── Quick Edit coverage: structural ops (server planner gaps) ──────────
|
||||
# Every prompt that failed the live Quick Edit coverage audit because the
|
||||
# server planner emitted a generic adjustments/transform step instead of the
|
||||
# exact structural op the browser supports.
|
||||
_STRUCTURAL_PROMPTS = [
|
||||
("fill the white area with red", "fill.flood"), # 1
|
||||
("flood fill the sky with blue", "fill.flood"), # 2
|
||||
("make the text bigger", "text.edit"), # 3
|
||||
("make the title red", "text.edit"), # 4
|
||||
("add a colour overlay", "layers.addStyle"), # 8
|
||||
("align the logo to the centre", "layers.align"), # 9
|
||||
("align the layer to the left", "layers.align"), # 10
|
||||
("rotate only this layer 90 degrees", "layer.rotate"), # 11
|
||||
("flip this layer horizontally", "layer.flip"), # 12
|
||||
("rotate the whole image 90 degrees", "document.rotate"), # 13
|
||||
("flip the canvas vertically", "document.flip"), # 14
|
||||
("copy the selection", "selection.copy"), # 19
|
||||
("paste", "selection.paste"), # 20
|
||||
("flatten the image", "layers.flatten"), # 40
|
||||
("merge the layer down", "layers.mergeDown"), # 41
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prompt,expected_op", _STRUCTURAL_PROMPTS)
|
||||
def test_structural_plan_routes_to_the_exact_op(prompt, expected_op):
|
||||
plan = guide.structural_plan(prompt)
|
||||
assert plan is not None, prompt
|
||||
assert plan["planner"] == "deterministic"
|
||||
assert expected_op in _ops(plan["steps"]), (prompt, _ops(plan["steps"]))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prompt,expected_op", _STRUCTURAL_PROMPTS)
|
||||
def test_plan_special_emits_structural_ops(prompt, expected_op):
|
||||
# The server override path (_plan_special) must return the structural step,
|
||||
# otherwise a wrong model plan suppresses the browser fallback.
|
||||
special = eaar._plan_special(prompt, {"width": 900, "height": 600}, True)
|
||||
assert special is not None, prompt
|
||||
steps, reply, _ = special
|
||||
assert expected_op in _ops(steps), (prompt, _ops(steps))
|
||||
assert reply
|
||||
|
||||
|
||||
def test_structural_ops_are_in_the_server_op_catalog():
|
||||
for name in ("fill.flood", "selection.fill", "text.edit", "layers.align",
|
||||
"layer.rotate", "layer.flip", "document.rotate", "document.flip",
|
||||
"selection.copy", "selection.paste", "layers.flatten",
|
||||
"layers.mergeDown"):
|
||||
assert name in eaar._OP_SET, name
|
||||
|
||||
|
||||
def test_structural_plan_ignores_non_structural_intents():
|
||||
for prompt in ("make it brighter", "blur the background", "make it punchy",
|
||||
"restore this old photo", "smooth the skin", "sharpen it",
|
||||
"add the title \"Summer\"", "duplicate this layer"):
|
||||
assert guide.structural_plan(prompt) is None, prompt
|
||||
|
||||
|
||||
def test_structural_params_match_the_client():
|
||||
# Canonical params must match static/js/editor/galleryEditor.js control ops.
|
||||
fill = guide.structural_plan("fill the white area with red")["steps"][0]
|
||||
assert fill["op"] == "fill.flood" and fill["params"]["color"] == "#ff0000"
|
||||
sel = guide.structural_plan("fill the selection with blue")["steps"][0]
|
||||
assert sel["op"] == "selection.fill" and sel["params"]["color"] == "#0000ff"
|
||||
align = guide.structural_plan("align the layer to the left")["steps"][0]
|
||||
assert align["params"]["mode"] == "left"
|
||||
rot = guide.structural_plan("rotate only this layer 90 degrees")["steps"][0]
|
||||
assert rot["params"]["degrees"] == 90
|
||||
flip = guide.structural_plan("flip the canvas vertically")["steps"][0]
|
||||
assert flip["op"] == "document.flip" and flip["params"]["axis"] == "v"
|
||||
overlay = guide.structural_plan("add a colour overlay")["steps"][0]
|
||||
assert overlay["params"]["type"] == "color-overlay"
|
||||
edit = guide.structural_plan("make the text bigger")["steps"][0]
|
||||
assert edit["op"] == "text.edit" and edit["params"]["size"] > 0
|
||||
|
||||
|
||||
def test_plan_route_returns_structural_steps(client):
|
||||
resp = client.post("/api/editor/plan", json={
|
||||
"prompt": "flatten the image",
|
||||
"state": {"width": 900, "height": 600},
|
||||
})
|
||||
data = resp.json()
|
||||
assert data["planner"] == "deterministic"
|
||||
assert "layers.flatten" in _ops(data["steps"])
|
||||
|
||||
|
||||
def test_style_guide_serves_structural_rule():
|
||||
rules = guide.as_json()["rules"]
|
||||
assert "structural" in rules
|
||||
assert guide.STYLE_GUIDE_VERSION >= 5
|
||||
|
||||
|
||||
# ── Targeted recolour + named reference looks ──────────────────────────
|
||||
|
||||
_RECOLOR_CASES = [
|
||||
("make only the red car blue", "red"),
|
||||
("recolour the red rock to teal", "red"),
|
||||
("change the yellow car to red", "yellow"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prompt,source", _RECOLOR_CASES)
|
||||
def test_targeted_recolor_selects_object_then_masked_hue(prompt, source):
|
||||
plan = guide.recolor_plan(prompt)
|
||||
assert plan is not None
|
||||
assert _ops(plan["steps"]) == ["selection.samText", "adjustments.applyMasked"]
|
||||
sel, adj = plan["steps"]
|
||||
assert sel["params"]["text"] # object description for the mask
|
||||
assert sel["params"]["fallbackColor"] == guide.color_for_name(source)
|
||||
assert adj["params"]["mode"] == "hue-saturation"
|
||||
assert "hue" in adj["params"] and isinstance(adj["params"]["hue"], int)
|
||||
|
||||
|
||||
def test_recolor_never_emits_selection_fill_or_a_global_tint():
|
||||
for prompt, _ in _RECOLOR_CASES:
|
||||
ops = _ops(guide.recolor_plan(prompt)["steps"])
|
||||
# A flat fill destroys shading; a whole-image tint is not targeted.
|
||||
assert "selection.fill" not in ops
|
||||
assert "adjustments.apply" not in ops # exact op, not applyMasked
|
||||
|
||||
|
||||
def test_targeted_recolor_hue_delta_is_source_to_target():
|
||||
# red (hue 0) -> blue (hue 240): shortest rotation is -120 degrees.
|
||||
plan = guide.recolor_plan("make only the red car blue")
|
||||
assert plan["steps"][1]["params"]["hue"] == -120
|
||||
|
||||
|
||||
def test_recolor_masked_hue_carries_source_hue():
|
||||
# red has hue 0; the client uses sourceHue as the rotation center so
|
||||
# neighbouring colours (chrome, windows, highlights) keep their own hue.
|
||||
plan = guide.recolor_plan("make only the red car blue")
|
||||
assert plan["steps"][1]["params"]["sourceHue"] == 0
|
||||
|
||||
|
||||
def test_enforce_guards_drops_global_lift_after_masked_recolor():
|
||||
# A recolour keeps shading; a later whole-image tone lift would tint the
|
||||
# rest of the photo, so enforce_guards must drop global lifts after it.
|
||||
steps = [
|
||||
{"op": "selection.samText", "params": {"text": "red car"}},
|
||||
{"op": "adjustments.applyMasked",
|
||||
"params": {"mode": "hue-saturation", "hue": -120, "sourceHue": 0}},
|
||||
{"op": "adjustments.apply",
|
||||
"params": {"type": "brightness-contrast", "params": {"brightness": 20}}},
|
||||
{"op": "photo.autoTone", "params": {}},
|
||||
]
|
||||
assert _ops(guide.enforce_guards(steps)) == ["selection.samText", "adjustments.applyMasked"]
|
||||
|
||||
|
||||
def test_enforce_guards_drops_whole_image_grade_after_masked_recolor():
|
||||
# Not just tone "lifts": ANY whole-image adjustment (vibrance/saturation/
|
||||
# autoTone) after a masked recolour re-grades the photo and undoes the
|
||||
# "only this object" intent, so enforce_guards must drop them all.
|
||||
steps = [
|
||||
{"op": "selection.samText", "params": {"text": "red car"}},
|
||||
{"op": "adjustments.applyMasked",
|
||||
"params": {"mode": "hue-saturation", "hue": -120, "sourceHue": 0}},
|
||||
{"op": "adjustments.apply",
|
||||
"params": {"type": "vibrance", "params": {"vibrance": 20}}},
|
||||
{"op": "adjustments.apply",
|
||||
"params": {"type": "hue-saturation", "params": {"saturation": 12}}},
|
||||
{"op": "photo.autoTone", "params": {}},
|
||||
]
|
||||
assert _ops(guide.enforce_guards(steps)) == ["selection.samText", "adjustments.applyMasked"]
|
||||
|
||||
|
||||
def test_enforce_guards_keeps_grades_when_there_is_no_recolor():
|
||||
# A recipe/look (no masked recolour) still keeps its whole-image grades.
|
||||
steps = [
|
||||
{"op": "photo.autoTone", "params": {}},
|
||||
{"op": "adjustments.apply", "params": {"type": "vibrance", "params": {"vibrance": 18}}},
|
||||
{"op": "adjustments.apply",
|
||||
"params": {"type": "brightness-contrast", "params": {"brightness": 6, "contrast": -10}}},
|
||||
]
|
||||
assert _ops(guide.enforce_guards(steps)) == ["photo.autoTone", "adjustments.apply"]
|
||||
|
||||
|
||||
def test_recolor_to_white_uses_the_image_model_not_a_hue_shift():
|
||||
# A white/black/grey target cannot be reached by a hue rotation.
|
||||
plan = guide.recolor_plan("make the red car white")
|
||||
assert plan is not None
|
||||
assert _ops(plan["steps"]) == ["selection.samText", "ai_edit_image"]
|
||||
assert plan["steps"][1]["params"]["mode"] == "masked"
|
||||
|
||||
|
||||
def test_recolor_with_single_colour_uses_masked_image_model():
|
||||
plan = guide.recolor_plan("make the flowers bright pink")
|
||||
assert plan is not None
|
||||
assert _ops(plan["steps"]) == ["selection.samText", "ai_edit_image"]
|
||||
assert plan["steps"][0]["params"]["text"] == "flowers"
|
||||
params = plan["steps"][1]["params"]
|
||||
assert params["mode"] == "masked"
|
||||
assert params["steps"] == 8
|
||||
|
||||
|
||||
def test_recolor_does_not_hijack_unrelated_colour_prompts():
|
||||
for prompt in ("make it blue", "fill the white area with red",
|
||||
"make the title red", "make this black and white",
|
||||
"blur the background"):
|
||||
assert guide.recolor_plan(prompt) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prompt", [
|
||||
"make it look like a Wes Anderson film",
|
||||
"make it look like a film noir",
|
||||
"make it look like vintage Kodachrome",
|
||||
])
|
||||
def test_named_recipe_looks_are_adjustment_stacks(prompt):
|
||||
plan = guide.look_plan(prompt)
|
||||
assert plan is not None
|
||||
ops = _ops(plan["steps"])
|
||||
assert ops and "ai_edit_image" not in ops
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prompt", [
|
||||
"make it look like cyberpunk",
|
||||
"make it look like a Studio Ghibli film",
|
||||
])
|
||||
def test_strong_named_looks_use_instruction_mode(prompt):
|
||||
plan = guide.look_plan(prompt)
|
||||
assert plan is not None
|
||||
assert _ops(plan["steps"]) == ["ai_edit_image"]
|
||||
assert plan["steps"][0]["params"]["mode"] == "instruction"
|
||||
assert plan["steps"][0]["params"]["steps"] == 8
|
||||
|
||||
|
||||
def test_structural_plan_dispatches_looks_and_recolor():
|
||||
wes = guide.structural_plan("make it look like a Wes Anderson film")
|
||||
assert "adjustments.apply" in _ops(wes["steps"])
|
||||
recolor = guide.structural_plan("make only the red car blue")
|
||||
assert _ops(recolor["steps"]) == ["selection.samText", "adjustments.applyMasked"]
|
||||
|
||||
|
||||
def test_named_looks_do_not_hijack_generative_lookalikes():
|
||||
for prompt in ("make it look snowy", "make it look like summer"):
|
||||
assert guide.look_plan(prompt) is None
|
||||
|
||||
|
||||
def test_plan_route_returns_targeted_recolor(client):
|
||||
resp = client.post("/api/editor/plan", json={
|
||||
"prompt": "make only the red car blue",
|
||||
"state": {"width": 900, "height": 600},
|
||||
})
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["planner"] == "deterministic"
|
||||
assert _ops(data["steps"]) == ["selection.samText", "adjustments.applyMasked"]
|
||||
assert "selection.fill" not in _ops(data["steps"])
|
||||
|
||||
|
||||
def test_plan_route_returns_named_look(client):
|
||||
resp = client.post("/api/editor/plan", json={
|
||||
"prompt": "make it look like a Wes Anderson film",
|
||||
"state": {"width": 900, "height": 600},
|
||||
})
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["planner"] == "deterministic"
|
||||
assert "adjustments.apply" in _ops(data["steps"])
|
||||
|
||||
|
||||
def test_style_guide_serves_recolor_and_looks_rules():
|
||||
body = guide.as_json()
|
||||
assert "recolor" in body["rules"]
|
||||
assert "looks" in body["rules"]
|
||||
assert {lk["id"] for lk in body["looks"]} >= {
|
||||
"wes_anderson", "film_noir", "cyberpunk", "vintage_kodachrome",
|
||||
"studio_ghibli",
|
||||
}
|
||||
|
||||
|
||||
if __name__ == "__main__": # pragma: no cover
|
||||
raise SystemExit(pytest.main([__file__, "-q"]))
|
||||
@@ -0,0 +1,51 @@
|
||||
"""Graceful-shutdown signalling tests (HARNESS-owned).
|
||||
|
||||
SSE / streaming responses never complete on their own, so uvicorn's graceful
|
||||
shutdown would wait for an open browser stream up to the whole
|
||||
``timeout_graceful_shutdown`` window. src.shutdown is a process-wide flag that
|
||||
the SIGTERM/SIGINT handler flips and every streaming generator checks, so a
|
||||
restart ends open streams promptly while ordinary in-flight requests still
|
||||
finish.
|
||||
"""
|
||||
from pathlib import Path
|
||||
|
||||
from src import shutdown
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
|
||||
|
||||
def test_request_shutdown_is_sticky_and_resettable():
|
||||
shutdown.reset()
|
||||
try:
|
||||
assert shutdown.is_shutting_down() is False
|
||||
shutdown.request_shutdown()
|
||||
assert shutdown.is_shutting_down() is True
|
||||
# idempotent
|
||||
shutdown.request_shutdown()
|
||||
assert shutdown.is_shutting_down() is True
|
||||
finally:
|
||||
shutdown.reset()
|
||||
assert shutdown.is_shutting_down() is False
|
||||
|
||||
|
||||
def test_app_installs_its_own_shutdown_handler():
|
||||
src = (ROOT / "app.py").read_text(encoding="utf-8")
|
||||
# This uvicorn version has no install_signal_handlers kwarg, so app.py wraps
|
||||
# server.handle_exit (which capture_signals wires SIGINT/SIGTERM to).
|
||||
assert "install_signal_handlers=False" not in src
|
||||
assert "server.handle_exit = _handle_shutdown_signal" in src
|
||||
assert "_shutdown.request_shutdown()" in src
|
||||
assert "server.run()" in src
|
||||
|
||||
|
||||
def test_streaming_generators_check_the_shutdown_flag():
|
||||
# editor bridge SSE + chat stream/resume + research + email events + shell.
|
||||
for rel in (
|
||||
"routes/editor_agent_routes.py",
|
||||
"routes/chat_routes.py",
|
||||
"routes/research/research_routes.py",
|
||||
"routes/email/email_routes.py",
|
||||
"routes/shell_routes.py",
|
||||
):
|
||||
src = (ROOT / rel).read_text(encoding="utf-8")
|
||||
assert "_shutdown.is_shutting_down()" in src, rel
|
||||
Reference in New Issue
Block a user