diff --git a/tools/fal_common.py b/tools/fal_common.py index 61feb31ed4..45055feff9 100644 --- a/tools/fal_common.py +++ b/tools/fal_common.py @@ -1,10 +1,9 @@ """Shared FAL.ai SDK plumbing: lazy import, managed-gateway sync client, small helpers. Stateful pieces (cache globals, ``_managed_fal_client*``, ``_submit_fal_request``) -intentionally stay on :mod:`tools.image_generation_tool`: it is the patch target -for the test suites and for ``plugins/image_gen/fal/``'s ``_it`` indirection, so -moving the caches here would silently defeat -``monkeypatch.setattr(image_tool, "_managed_fal_client", None)``. +intentionally stay on :mod:`tools.image_generation_tool`: it is the patch target for the +test suites and for ``plugins/image_gen/fal/``'s ``_it`` indirection, so moving the caches +here would silently defeat ``monkeypatch.setattr(image_tool, "_managed_fal_client", None)``. """ from __future__ import annotations @@ -16,9 +15,8 @@ from urllib.parse import urlencode def import_fal_client() -> Any: """Import ``fal_client`` (via ``lazy_deps`` when available); raises ImportError if unavailable. - Not imported at cold start (it cost ~64 ms per CLI invocation). Callers - cache the result on their own module global so tests can monkeypatch that - module's ``fal_client`` attribute and have it stick for its call sites. + Not imported at cold start (~64 ms per CLI invocation). Callers cache the result on their + own module global so tests can monkeypatch that module's ``fal_client`` attribute. """ try: from tools.lazy_deps import ensure as _lazy_ensure @@ -46,26 +44,25 @@ def _extract_http_status(exc: BaseException) -> Optional[int]: if isinstance(status, int): return status status = getattr(exc, "status_code", None) - if isinstance(status, int): - return status - return None + return status if isinstance(status, int) else None + + +def _require(value: Any, what: str) -> Any: + if value is None: + raise RuntimeError(f"{what} is required for managed FAL gateway mode") + return value class _ManagedFalSyncClient: """Per-instance wrapper driving a Nous-managed fal-queue gateway via ``fal_client.SyncClient`` primitives. - Carries its own ``fal_client`` reference instead of a module global so the - caller decides which module's (possibly test-patched) ``fal_client`` is used. + Carries its own ``fal_client`` reference instead of a module global so the caller decides + which module's (possibly test-patched) ``fal_client`` is used. """ def __init__(self, fal_client: Any, *, key: str, queue_run_origin: str): - sync_client_class = getattr(fal_client, "SyncClient", None) - if sync_client_class is None: - raise RuntimeError("fal_client.SyncClient is required for managed FAL gateway mode") - - client_module = getattr(fal_client, "client", None) - if client_module is None: - raise RuntimeError("fal_client.client is required for managed FAL gateway mode") + sync_client_class = _require(getattr(fal_client, "SyncClient", None), "fal_client.SyncClient") + client_module = _require(getattr(fal_client, "client", None), "fal_client.client") self._queue_url_format = _normalize_fal_queue_url_format(queue_run_origin) self._sync_client = sync_client_class(key=key) @@ -77,24 +74,15 @@ class _ManagedFalSyncClient: self._add_priority_header = getattr(client_module, "add_priority_header", None) self._add_timeout_header = getattr(client_module, "add_timeout_header", None) - if self._http_client is None: - raise RuntimeError("fal_client.SyncClient._client is required for managed FAL gateway mode") + _require(self._http_client, "fal_client.SyncClient._client") if self._maybe_retry_request is None or self._raise_for_status is None: raise RuntimeError("fal_client.client request helpers are required for managed FAL gateway mode") - if self._request_handle_class is None: - raise RuntimeError("fal_client.client.SyncRequestHandle is required for managed FAL gateway mode") + _require(self._request_handle_class, "fal_client.client.SyncRequestHandle") def submit( - self, - application: str, - arguments: Dict[str, Any], - *, - path: str = "", - hint: Optional[str] = None, - webhook_url: Optional[str] = None, - priority: Any = None, - headers: Optional[Dict[str, str]] = None, - start_timeout: Optional[Union[int, float]] = None, + self, application: str, arguments: Dict[str, Any], *, path: str = "", + hint: Optional[str] = None, webhook_url: Optional[str] = None, priority: Any = None, + headers: Optional[Dict[str, str]] = None, start_timeout: Optional[Union[int, float]] = None, ): url = self._queue_url_format + application if path: @@ -115,20 +103,13 @@ class _ManagedFalSyncClient: self._add_timeout_header(start_timeout, request_headers) response = self._maybe_retry_request( - self._http_client, - "POST", - url, - json=arguments, - timeout=getattr(self._sync_client, "default_timeout", 120.0), - headers=request_headers, + self._http_client, "POST", url, json=arguments, + timeout=getattr(self._sync_client, "default_timeout", 120.0), headers=request_headers, ) self._raise_for_status(response) data = response.json() return self._request_handle_class( - request_id=data["request_id"], - response_url=data["response_url"], - status_url=data["status_url"], - cancel_url=data["cancel_url"], - client=self._http_client, + request_id=data["request_id"], response_url=data["response_url"], + status_url=data["status_url"], cancel_url=data["cancel_url"], client=self._http_client, ) diff --git a/tools/image_generation_catalog.py b/tools/image_generation_catalog.py index d8e3c3660e..fc9e3ddfbe 100644 --- a/tools/image_generation_catalog.py +++ b/tools/image_generation_catalog.py @@ -1,52 +1,31 @@ """FAL image model catalog + upscaler constants for ``tools.image_generation_tool``. -Each entry declares how to translate the unified inputs (prompt + aspect_ratio) -into the model's native payload. ``size_style`` picks the family: -``"image_size_preset"`` (FAL preset enum), ``"aspect_ratio"`` (ratio enum), -``"gpt_literal"`` (literal "WxH" strings). ``supports`` / ``edit_supports`` are -whitelists — keys outside them are stripped so models never receive rejected -parameters. ``upscale`` (Clarity Upscaler chained after generation) is False -everywhere: Clarity redraws content (creativity 0.35) and degraded text/CJK/ -faces when default-on, so upscaling is strictly per-call opt-in. -Pricing strings are as-of-commit and allowed to drift. +Each entry translates the unified inputs (prompt + aspect_ratio) into the model's native +payload. ``size_style``: ``"image_size_preset"`` (FAL preset enum), ``"aspect_ratio"`` (ratio +enum), ``"gpt_literal"`` (literal "WxH"). ``supports`` / ``edit_supports`` are whitelists — +other keys are stripped so models never receive rejected parameters. ``upscale`` is False +everywhere: Clarity redraws content (creativity 0.35) and degraded text/CJK/faces when +default-on, so upscaling is strictly per-call opt-in. Pricing strings may drift. """ from typing import Any, Dict, Optional -_PRESET_SIZES = { - "landscape": "landscape_16_9", - "square": "square_hd", - "portrait": "portrait_16_9", -} +_PRESET_SIZES = {"landscape": "landscape_16_9", "square": "square_hd", "portrait": "portrait_16_9"} _ASPECT_SIZES = {"landscape": "16:9", "square": "1:1", "portrait": "9:16"} _DEFAULT_SIZES = {"image_size_preset": _PRESET_SIZES, "aspect_ratio": _ASPECT_SIZES} def _model( - display: str, - speed: str, - strengths: str, - price: str, - *, - style: str = "image_size_preset", - sizes: Optional[Dict[str, Any]] = None, - defaults: Dict[str, Any], - supports: set, - edit_endpoint: Optional[str] = None, - edit_supports: Optional[set] = None, + display: str, speed: str, strengths: str, price: str, *, style: str = "image_size_preset", + sizes: Optional[Dict[str, Any]] = None, defaults: Dict[str, Any], supports: set, + edit_endpoint: Optional[str] = None, edit_supports: Optional[set] = None, max_reference_images: Optional[int] = None, ) -> Dict[str, Any]: """Build one catalog entry; edit keys are present only for edit-capable models.""" entry: Dict[str, Any] = { - "display": display, - "speed": speed, - "strengths": strengths, - "price": price, - "size_style": style, - "sizes": sizes if sizes is not None else _DEFAULT_SIZES[style], - "defaults": defaults, - "supports": supports, - "upscale": False, + "display": display, "speed": speed, "strengths": strengths, "price": price, + "size_style": style, "sizes": sizes if sizes is not None else _DEFAULT_SIZES[style], + "defaults": defaults, "supports": supports, "upscale": False, } if edit_endpoint: entry["edit_endpoint"] = edit_endpoint diff --git a/tools/image_generation_tool.py b/tools/image_generation_tool.py index 11fde4e4ac..70f9ac7fb6 100644 --- a/tools/image_generation_tool.py +++ b/tools/image_generation_tool.py @@ -16,18 +16,16 @@ import uuid from typing import Any, Dict, Optional # Imported lazily by _load_fal_client(): the eager import cost ~64 ms on every CLI cold -# start (discover_builtin_tools() imports this module unconditionally). Tests that -# monkeypatch this attribute keep working because the loader short-circuits when truthy. +# start. Tests that monkeypatch this attribute keep working (loader short-circuits when truthy). fal_client: Any = None def _load_fal_client() -> Any: """Lazily import fal_client into the module global (idempotent; keeps a test-installed mock).""" global fal_client - if fal_client is not None: - return fal_client - from tools.fal_common import import_fal_client - fal_client = import_fal_client() + if fal_client is None: + from tools.fal_common import import_fal_client + fal_client = import_fal_client() return fal_client @@ -38,28 +36,14 @@ from tools.fal_common import ( _normalize_fal_queue_url_format, # noqa: F401 — re-exported for tests ) from tools.image_generation_catalog import ( # noqa: F401 — re-exported (plugins/tests/tools_config) - DEFAULT_ASPECT_RATIO, - DEFAULT_MODEL, - FAL_MODELS, - UPSCALER_CREATIVITY, - UPSCALER_DEFAULT_PROMPT, - UPSCALER_FACTOR, - UPSCALER_GUIDANCE_SCALE, - UPSCALER_MODEL, - UPSCALER_NEGATIVE_PROMPT, - UPSCALER_NUM_INFERENCE_STEPS, - UPSCALER_RESEMBLANCE, - UPSCALER_SAFETY_CHECKER, - VALID_ASPECT_RATIOS, + DEFAULT_ASPECT_RATIO, DEFAULT_MODEL, FAL_MODELS, UPSCALER_CREATIVITY, UPSCALER_DEFAULT_PROMPT, + UPSCALER_FACTOR, UPSCALER_GUIDANCE_SCALE, UPSCALER_MODEL, UPSCALER_NEGATIVE_PROMPT, + UPSCALER_NUM_INFERENCE_STEPS, UPSCALER_RESEMBLANCE, UPSCALER_SAFETY_CHECKER, VALID_ASPECT_RATIOS, ) from tools.managed_tool_gateway import resolve_managed_tool_gateway from tools.tool_backend_helpers import ( - NOUS_MANAGED_PROVIDER, - fal_key_is_configured, - managed_nous_tools_enabled, - nous_tool_gateway_unavailable_message, - read_selection, - selection_error, + NOUS_MANAGED_PROVIDER, fal_key_is_configured, managed_nous_tools_enabled, + nous_tool_gateway_unavailable_message, read_selection, selection_error, ) logger = logging.getLogger(__name__) @@ -70,9 +54,7 @@ _managed_fal_client_config = None _managed_fal_client_lock = threading.Lock() -# --------------------------------------------------------------------------- -# Managed FAL gateway (Nous Subscription) -# --------------------------------------------------------------------------- +# --- Managed FAL gateway (Nous Subscription) --- def _resolve_managed_fal_gateway(): """Managed gateway config for the stored `hermes tools` selection, or ``None`` for direct FAL. @@ -95,9 +77,7 @@ def _resolve_managed_fal_gateway(): raise ValueError(selection_error("image_gen", selected, "FAL_KEY is not set")) return None # Never-configured category: legacy credential autodetect (do NOT persist). - if fal_key_is_configured(): - return None - return resolve_managed_tool_gateway("fal-queue") + return None if fal_key_is_configured() else resolve_managed_tool_gateway("fal-queue") def _get_managed_fal_client(managed_gateway): @@ -106,17 +86,13 @@ def _get_managed_fal_client(managed_gateway): client_config = (managed_gateway.gateway_origin.rstrip("/"), managed_gateway.nous_user_token) with _managed_fal_client_lock: - if _managed_fal_client is not None and _managed_fal_client_config == client_config: - return _managed_fal_client - - # Resolved on this module so monkeypatching ``image_generation_tool.fal_client`` still applies. - _load_fal_client() - _managed_fal_client = _ManagedFalSyncClient( - fal_client, - key=managed_gateway.nous_user_token, - queue_run_origin=managed_gateway.gateway_origin, - ) - _managed_fal_client_config = client_config + if _managed_fal_client is None or _managed_fal_client_config != client_config: + # Resolved on this module so monkeypatching ``image_generation_tool.fal_client`` still applies. + _managed_fal_client = _ManagedFalSyncClient( + _load_fal_client(), key=managed_gateway.nous_user_token, + queue_run_origin=managed_gateway.gateway_origin, + ) + _managed_fal_client_config = client_config return _managed_fal_client @@ -163,9 +139,9 @@ def _submit_fal_request(model: str, arguments: Dict[str, Any]): if managed_gateway is None: return fal_client.submit(model, arguments=arguments, headers=request_headers) - managed_client = _get_managed_fal_client(managed_gateway) try: - return managed_client.submit(model, arguments=arguments, headers=request_headers) + return _get_managed_fal_client(managed_gateway).submit( + model, arguments=arguments, headers=request_headers) except Exception as exc: # A managed-gateway 4xx usually means the portal doesn't proxy this model # (allowlist miss, billing gate): give remediation instead of a raw httpx error. @@ -187,9 +163,7 @@ def _submit_fal_request(model: str, arguments: Dict[str, Any]): raise -# --------------------------------------------------------------------------- -# Config readers, model resolution + payload construction -# --------------------------------------------------------------------------- +# --- Config readers, model resolution + payload construction --- def _read_image_gen_key(key: str) -> Optional[str]: """Return the stripped ``image_gen.`` string from config.yaml, or None.""" try: @@ -212,23 +186,27 @@ def _read_configured_image_model(): def _read_configured_image_provider(): """``image_gen.provider`` from config.yaml, or None. - The plugin registry is consulted only when this is explicitly set — unset keeps - users on the in-tree FAL fallback even when other providers are registered (e.g. - OPENAI_API_KEY present for other features). ``"fal"`` routes through + The plugin registry is consulted only when this is explicitly set — unset keeps users on + the in-tree FAL fallback even when other providers are registered. ``"fal"`` routes through ``plugins/image_gen/fal/``, which delegates back here via call-time indirection. """ return _read_image_gen_key("provider") +def _plugin_provider_name() -> Optional[str]: + """Configured provider that must go through the plugin registry; None for unset/fal/nous.""" + configured = _read_configured_image_provider() + if not configured or configured in ("fal", NOUS_MANAGED_PROVIDER): + return None + return configured + + def _resolve_fal_model() -> tuple: """Return ``(model_id, meta)`` for the configured FAL model, falling back to DEFAULT_MODEL (warned) when unknown.""" # FAL_IMAGE_MODEL is an undocumented escape hatch (backward-compat for tests/scripts). model_id = _read_image_gen_key("model") or os.getenv("FAL_IMAGE_MODEL", "").strip() if model_id and model_id not in FAL_MODELS: - logger.warning( - "Unknown FAL model '%s' in config; falling back to %s", - model_id, DEFAULT_MODEL, - ) + logger.warning("Unknown FAL model '%s' in config; falling back to %s", model_id, DEFAULT_MODEL) model_id = None model_id = model_id or DEFAULT_MODEL return model_id, FAL_MODELS[model_id] @@ -237,14 +215,7 @@ def _resolve_fal_model() -> tuple: _SIZE_KEY_BY_STYLE = {"image_size_preset": "image_size", "gpt_literal": "image_size", "aspect_ratio": "aspect_ratio"} -def _build_payload( - model_id: str, - prompt: str, - aspect_ratio: str, - seed: Optional[int], - overrides: Optional[Dict[str, Any]], - image_urls: Optional[list] = None, -) -> Dict[str, Any]: +def _build_payload(model_id, prompt, aspect_ratio, seed, overrides, image_urls=None) -> Dict[str, Any]: """Text-to-image / edit payload (``image_urls`` selects edit mode): defaults + native size spec + overrides, filtered to the model whitelist. @@ -288,9 +259,7 @@ def _build_fal_edit_payload(model_id, prompt, image_urls, aspect_ratio=DEFAULT_A return _build_payload(model_id, prompt, aspect_ratio, seed, overrides, image_urls=image_urls) -# --------------------------------------------------------------------------- -# Upscaler -# --------------------------------------------------------------------------- +# --- Upscaler --- def _upscale_image(image_url: str, original_prompt: str) -> Optional[Dict[str, Any]]: """Upscale via FAL's Clarity Upscaler; None on failure (caller keeps the original).""" try: @@ -309,17 +278,14 @@ def _upscale_image(image_url: str, original_prompt: str) -> Optional[Dict[str, A result = _wait_fal_result(handler) if result and "image" in result: up = result["image"] - logger.info( - "Image upscaled successfully to %sx%s", - up.get("width", "unknown"), up.get("height", "unknown"), - ) + logger.info("Image upscaled successfully to %sx%s", + up.get("width", "unknown"), up.get("height", "unknown")) return { "url": up["url"], "width": up.get("width", 0), "height": up.get("height", 0), "upscaled": True, "upscale_factor": UPSCALER_FACTOR, } logger.error("Upscaler returned invalid response") return None - except ImageGenerationInterrupted: # A user interrupt must not degrade into a silent "use original" fallback. raise @@ -328,9 +294,14 @@ def _upscale_image(image_url: str, original_prompt: str) -> Optional[Dict[str, A return None -# --------------------------------------------------------------------------- -# Artifact path hinting for non-local terminal backends -# --------------------------------------------------------------------------- +# --- Artifact path hinting for non-local terminal backends --- +_CONTAINER_HOME_ENVS = {"DockerEnvironment", "SingularityEnvironment", "ModalEnvironment"} +# No environment yet: only backends with deterministic cache roots can be translated without +# side effects. SSH uses a shell-visible tilde path; its first sync uploads the cache file. +_CACHE_BASE_BY_BACKEND = {"docker": "/root/.hermes", "singularity": "/root/.hermes", + "modal": "/root/.hermes", "ssh": "~/.hermes"} + + def _looks_like_absolute_file_path(value: str) -> bool: if not value or not isinstance(value, str) or value.lower().startswith(("http://", "https://", "data:")): return False @@ -340,7 +311,6 @@ def _looks_like_absolute_file_path(value: str) -> bool: def _active_terminal_env(task_id: str | None): try: from tools.terminal_tool import get_active_env - return get_active_env(task_id or "default") except Exception as exc: # noqa: BLE001 - artifact hinting must not break generation logger.debug("Could not inspect active terminal environment: %s", exc) @@ -349,8 +319,7 @@ def _active_terminal_env(task_id: str | None): def _agent_cache_base_for_env(env: Any) -> str | None: if env is not None: - # Optional extension hook: an environment may expose its own agent-visible - # cache root. No backend defines it yet; the guards make it a safe no-op. + # Optional extension hook: an environment may expose its own agent-visible cache root. explicit = getattr(env, "agent_visible_cache_base", None) if callable(explicit): try: @@ -359,42 +328,13 @@ def _agent_cache_base_for_env(env: Any) -> str | None: return str(value).rstrip("/") except Exception as exc: # noqa: BLE001 logger.debug("active env agent_visible_cache_base failed: %s", exc) - remote_home = getattr(env, "_remote_home", None) if remote_home: return f"{str(remote_home).rstrip('/')}/.hermes" - if env.__class__.__name__ in {"DockerEnvironment", "SingularityEnvironment", "ModalEnvironment"}: + if env.__class__.__name__ in _CONTAINER_HOME_ENVS: return "/root/.hermes" - - # No environment yet: only backends with deterministic cache roots can be - # translated without side effects. SSH can use a shell-visible tilde path; - # its first environment sync uploads the cache file before the first command. backend = (os.getenv("TERMINAL_ENV") or "local").strip().lower() - return {"docker": "/root/.hermes", "singularity": "/root/.hermes", - "modal": "/root/.hermes", "ssh": "~/.hermes"}.get(backend) - - -def _agent_visible_cache_path(host_path: str, env: Any) -> str | None: - cache_base = _agent_cache_base_for_env(env) if _looks_like_absolute_file_path(host_path) else None - if not cache_base: - return None - try: - from tools.credential_files import map_cache_path_to_container - - return map_cache_path_to_container(host_path, container_base=cache_base) - except Exception as exc: # noqa: BLE001 - logger.debug("Could not translate image cache path for backend: %s", exc) - return None - - -def _force_artifact_sync(env: Any) -> None: - sync_manager = getattr(env, "_sync_manager", None) - if sync_manager is None: - return - try: - sync_manager.sync(force=True) - except Exception as exc: # noqa: BLE001 - keep generation success; log for operators - logger.warning("Could not force-sync generated image artifact: %s", exc) + return _CACHE_BASE_BY_BACKEND.get(backend) def _postprocess_image_generate_result(raw: str, task_id: str | None = None) -> str: @@ -411,27 +351,29 @@ def _postprocess_image_generate_result(raw: str, task_id: str | None = None) -> return raw env = _active_terminal_env(task_id) - agent_path = _agent_visible_cache_path(image, env) + cache_base = _agent_cache_base_for_env(env) + if not cache_base: + return raw + try: + from tools.credential_files import map_cache_path_to_container + agent_path = map_cache_path_to_container(image, container_base=cache_base) + except Exception as exc: # noqa: BLE001 + logger.debug("Could not translate image cache path for backend: %s", exc) + return raw if not agent_path or agent_path == image: return raw - if env is not None: - _force_artifact_sync(env) + sync_manager = getattr(env, "_sync_manager", None) + if sync_manager is not None: + try: + sync_manager.sync(force=True) + except Exception as exc: # noqa: BLE001 - keep generation success; log for operators + logger.warning("Could not force-sync generated image artifact: %s", exc) payload.setdefault("host_image", image) payload.setdefault("agent_visible_image", agent_path) return json.dumps(payload, ensure_ascii=False) -# --------------------------------------------------------------------------- -# Tool entry point -# --------------------------------------------------------------------------- -def _collect_source_images(image_url, reference_image_urls) -> list: - """Primary + reference source images as one ordered list of stripped, non-empty strings.""" - candidates = [image_url] - if isinstance(reference_image_urls, (list, tuple)): - candidates.extend(reference_image_urls) - return [c.strip() for c in candidates if isinstance(c, str) and c.strip()] - - +# --- Tool entry point --- def _format_images(images: list, should_upscale: bool, prompt: str) -> list: """Normalize FAL result images, optionally chaining the upscaler (falls back to the original on failure).""" formatted = [] @@ -451,14 +393,6 @@ def _format_images(images: list, should_upscale: bool, prompt: str) -> list: return formatted -def _finish_image_call(debug_call_data: Dict[str, Any], generation_time: float, response: Dict[str, Any]) -> str: - """Record generation time, log the debug entry and return the JSON result.""" - debug_call_data["generation_time"] = generation_time - _debug.log_call("image_generate_tool", debug_call_data) - _debug.save() - return json.dumps(response, indent=2, ensure_ascii=False) - - def _prepare_fal_request(model_id, meta, prompt, aspect_ratio, seed, overrides, source_images): """Validate inputs and return ``(endpoint, arguments)``; raises ValueError with the user-facing message.""" if not isinstance(prompt, str) or not prompt.strip(): @@ -490,12 +424,9 @@ def _prepare_fal_request(model_id, meta, prompt, aspect_ratio, seed, overrides, max_refs = int(meta.get("max_reference_images") or 1) clamped_sources = source_images[:max_refs] if max_refs > 0 else source_images arguments = _build_fal_edit_payload( - model_id, prompt, clamped_sources, aspect_lc, seed=seed, overrides=overrides, - ) - logger.info( - "Editing image with %s (%s) — %d source image(s), prompt: %s", - display, edit_endpoint, len(clamped_sources), prompt[:80], - ) + model_id, prompt, clamped_sources, aspect_lc, seed=seed, overrides=overrides) + logger.info("Editing image with %s (%s) — %d source image(s), prompt: %s", + display, edit_endpoint, len(clamped_sources), prompt[:80]) return edit_endpoint, arguments arguments = _build_fal_payload(model_id, prompt, aspect_lc, seed=seed, overrides=overrides) logger.info("Generating image with %s (%s) — prompt: %s", display, model_id, prompt[:80]) @@ -503,16 +434,11 @@ def _prepare_fal_request(model_id, meta, prompt, aspect_ratio, seed, overrides, def image_generate_tool( - prompt: str, - aspect_ratio: str = DEFAULT_ASPECT_RATIO, - num_inference_steps: Optional[int] = None, - guidance_scale: Optional[float] = None, - num_images: Optional[int] = None, - output_format: Optional[str] = None, - seed: Optional[int] = None, - image_url: Optional[str] = None, - reference_image_urls: Optional[list] = None, - upscale: Optional[bool] = None, + prompt: str, aspect_ratio: str = DEFAULT_ASPECT_RATIO, + num_inference_steps: Optional[int] = None, guidance_scale: Optional[float] = None, + num_images: Optional[int] = None, output_format: Optional[str] = None, + seed: Optional[int] = None, image_url: Optional[str] = None, + reference_image_urls: Optional[list] = None, upscale: Optional[bool] = None, ) -> str: """Generate (or, with source images + an ``edit_endpoint`` model, edit) an image via FAL. @@ -521,32 +447,34 @@ def image_generate_tool( model switches). Returns JSON ``{"success", "image", "modality", "error", "error_type"}``. """ model_id, meta = _resolve_fal_model() - source_images = _collect_source_images(image_url, reference_image_urls) + candidates = [image_url, *(reference_image_urls if isinstance(reference_image_urls, (list, tuple)) else [])] + source_images = [c.strip() for c in candidates if isinstance(c, str) and c.strip()] use_edit = bool(source_images) and bool(meta.get("edit_endpoint")) modality = "image" if use_edit else "text" - params = { - "prompt": prompt, "aspect_ratio": aspect_ratio, + overrides: Dict[str, Any] = { "num_inference_steps": num_inference_steps, "guidance_scale": guidance_scale, - "num_images": num_images, "output_format": output_format, "seed": seed, + "num_images": num_images, "output_format": output_format, } debug_call_data = { "model": model_id, - "parameters": {**params, "modality": modality, "source_images": len(source_images)}, + "parameters": {"prompt": prompt, "aspect_ratio": aspect_ratio, **overrides, "seed": seed, + "modality": modality, "source_images": len(source_images)}, "error": None, "success": False, "images_generated": 0, "generation_time": 0, } start_time = datetime.datetime.now() + def finish(generation_time: float, response: Dict[str, Any]) -> str: + debug_call_data["generation_time"] = generation_time + _debug.log_call("image_generate_tool", debug_call_data) + _debug.save() + return json.dumps(response, indent=2, ensure_ascii=False) + try: - overrides: Dict[str, Any] = { - k: params[k] for k in ("num_inference_steps", "guidance_scale", "num_images", "output_format") - if params[k] is not None - } endpoint, arguments = _prepare_fal_request( - model_id, meta, prompt, aspect_ratio, seed, overrides, source_images, - ) - handler = _submit_fal_request(endpoint, arguments=arguments) - result = _wait_fal_result(handler) + model_id, meta, prompt, aspect_ratio, seed, + {k: v for k, v in overrides.items() if v is not None}, source_images) + result = _wait_fal_result(_submit_fal_request(endpoint, arguments=arguments)) generation_time = (datetime.datetime.now() - start_time).total_seconds() if not result or "images" not in result: @@ -555,10 +483,9 @@ def image_generate_tool( if not images: raise ValueError("No images were generated") - # An explicit ``upscale`` wins over the catalog default, including for edits - # (an explicit request is intentional). The catalog default never upscales - # edits: Clarity is a text-to-image quality pass and must not silently alter - # edit compositions. + # An explicit ``upscale`` wins over the catalog default, including for edits. The + # catalog default never upscales edits: Clarity is a text-to-image quality pass and + # must not silently alter edit compositions. if upscale is not None: should_upscale = bool(upscale) else: @@ -569,31 +496,22 @@ def image_generate_tool( raise ValueError("No valid image URLs returned from API") upscaled_count = sum(1 for img in formatted_images if img.get("upscaled")) - logger.info( - "Generated %s image(s) in %.1fs (%s upscaled) via %s [%s]", - len(formatted_images), generation_time, upscaled_count, endpoint, - modality, - ) + logger.info("Generated %s image(s) in %.1fs (%s upscaled) via %s [%s]", + len(formatted_images), generation_time, upscaled_count, endpoint, modality) debug_call_data["success"] = True debug_call_data["images_generated"] = len(formatted_images) - return _finish_image_call(debug_call_data, generation_time, { + return finish(generation_time, { "success": True, "image": formatted_images[0]["url"], "modality": modality, "upscaled": bool(formatted_images[0].get("upscaled")), }) - except Exception as e: error_msg = f"Error generating image: {str(e)}" logger.error("%s", error_msg, exc_info=True) debug_call_data["error"] = error_msg generation_time = (datetime.datetime.now() - start_time).total_seconds() - return _finish_image_call(debug_call_data, generation_time, { - "success": False, - "image": None, - "error": str(e), - "error_type": type(e).__name__, - }) + return finish(generation_time, {"success": False, "image": None, "error": str(e), "error_type": type(e).__name__}) def check_fal_api_key() -> bool: @@ -640,12 +558,15 @@ def _build_no_backend_setup_message() -> str: return "\n".join(lines) -def _get_plugin_provider(name: str): +def _get_plugin_provider(name: str, *, force: bool = False): """Discover plugins (local import: importing this module must not trigger discovery) and return the named provider.""" from agent.image_gen_registry import get_provider from hermes_cli.plugins import _ensure_plugins_discovered - _ensure_plugins_discovered() + if force: + _ensure_plugins_discovered(force=True) + else: + _ensure_plugins_discovered() return get_provider(name) @@ -659,10 +580,9 @@ def check_image_generation_requirements() -> bool: except ImportError: pass - configured = _read_configured_image_provider() - if not configured or configured in ("fal", NOUS_MANAGED_PROVIDER): + configured = _plugin_provider_name() + if configured is None: return False - # Probe only the selected plugin: a cloud key alone must not opt a user into a paid backend. try: provider = _get_plugin_provider(configured) @@ -671,9 +591,7 @@ def check_image_generation_requirements() -> bool: return False -# --------------------------------------------------------------------------- -# Registry -# --------------------------------------------------------------------------- +# --- Registry --- from tools.registry import registry, tool_error IMAGE_GENERATE_SCHEMA = { @@ -710,21 +628,20 @@ IMAGE_GENERATE_SCHEMA = { } -# --------------------------------------------------------------------------- -# Plugin provider dispatch + managed-mode Krea routing -# --------------------------------------------------------------------------- +# --- Plugin provider dispatch + managed-mode Krea routing --- def _provider_error(error: str, error_type: str) -> str: """JSON error envelope shared by every provider-dispatch failure path.""" return json.dumps({"success": False, "image": None, "error": error, "error_type": error_type}) -def _add_provider_kwargs( - kwargs: Dict[str, Any], - image_url: Optional[str], - reference_image_urls: Optional[list], - upscale: Optional[bool], - model: Optional[str] = None, -) -> Dict[str, Any]: +def _provider_result(result, contract_error: str) -> str: + """JSON-encode a provider's dict result; anything else is a contract violation.""" + if not isinstance(result, dict): + return _provider_error(contract_error, "provider_contract") + return json.dumps(result) + + +def _add_provider_kwargs(kwargs, image_url, reference_image_urls, upscale, model=None) -> Dict[str, Any]: """Add the optional ``provider.generate(**kwargs)`` args in place (edit args only when supplied).""" if model: kwargs["model"] = model @@ -732,7 +649,6 @@ def _add_provider_kwargs( kwargs["image_url"] = image_url.strip() if reference_image_urls is not None: from agent.image_gen_provider import normalize_reference_images - norm_refs = normalize_reference_images(reference_image_urls) if norm_refs: kwargs["reference_image_urls"] = norm_refs @@ -742,11 +658,8 @@ def _add_provider_kwargs( def _dispatch_to_plugin_provider( - prompt: str, - aspect_ratio: str, - image_url: Optional[str] = None, - reference_image_urls: Optional[list] = None, - upscale: Optional[bool] = None, + prompt: str, aspect_ratio: str, image_url: Optional[str] = None, + reference_image_urls: Optional[list] = None, upscale: Optional[bool] = None, ): """JSON result from the selected plugin provider, or ``None`` to fall through to in-tree FAL. @@ -754,28 +667,21 @@ def _dispatch_to_plugin_provider( the legacy pipeline; ``"nous"`` via the managed fal-queue gateway). Edit args are forwarded for the backend's edit endpoint; providers without ``upscale`` ignore it via ``**kwargs``. """ - configured = _read_configured_image_provider() - if not configured or configured in ("fal", NOUS_MANAGED_PROVIDER): + configured = _plugin_provider_name() + if configured is None: return None try: - from hermes_cli.plugins import _ensure_plugins_discovered - provider = _get_plugin_provider(configured) except Exception as exc: logger.debug("image_gen plugin dispatch skipped: %s", exc) return None - if provider is None: # Long-lived sessions may have discovered plugins before a bundled backend # was patched in or config changed: retry once with a forced refresh. try: - from agent.image_gen_registry import get_provider - - _ensure_plugins_discovered(force=True) - provider = get_provider(configured) + provider = _get_plugin_provider(configured, force=True) except Exception as exc: logger.debug("image_gen plugin force-refresh skipped: %s", exc) - if provider is None: return _provider_error( f"image_gen.provider='{configured}' is set but no plugin " @@ -793,11 +699,8 @@ def _dispatch_to_plugin_provider( # generate() predating image_url support (third-party plugin not yet updated): # text-to-image keeps working; surface a clear note when an edit was requested. if "image_url" in kwargs or "reference_image_urls" in kwargs: - logger.warning( - "image_gen provider '%s' rejected image-to-image kwargs " - "(signature too narrow): %s", - pname, exc, - ) + logger.warning("image_gen provider '%s' rejected image-to-image kwargs " + "(signature too narrow): %s", pname, exc) return _provider_error( f"Provider '{pname}' does not " f"support image-to-image / editing (its generate() " @@ -811,9 +714,7 @@ def _dispatch_to_plugin_provider( except Exception as exc: logger.warning("Image gen provider '%s' raised: %s", pname, exc) return _provider_error(f"Provider '{pname}' error: {exc}", "provider_exception") - if not isinstance(result, dict): - return _provider_error("Provider returned a non-dict result", "provider_contract") - return json.dumps(result) + return _provider_result(result, "Provider returned a non-dict result") # Native ``krea-2-*`` ids are served by the Krea managed gateway (managed mode only — @@ -828,11 +729,8 @@ def _normalize_krea_model(model_id: Optional[str]) -> Optional[str]: def _maybe_route_managed_krea( - prompt: str, - aspect_ratio: str, - image_url: Optional[str] = None, - reference_image_urls: Optional[list] = None, - upscale: Optional[bool] = None, + prompt: str, aspect_ratio: str, image_url: Optional[str] = None, + reference_image_urls: Optional[list] = None, upscale: Optional[bool] = None, ) -> Optional[str]: """JSON result from the managed Krea gateway, or ``None`` to fall through. @@ -843,14 +741,11 @@ def _maybe_route_managed_krea( configured_provider = _read_configured_image_provider() if configured_provider is not None and configured_provider != NOUS_MANAGED_PROVIDER: return None - normalized = _normalize_krea_model(_read_configured_image_model()) if normalized is None: return None - try: from plugins.image_gen.krea import _resolve_managed_krea_gateway - if _resolve_managed_krea_gateway() is None: return None except Exception as exc: # noqa: BLE001 @@ -871,14 +766,10 @@ def _maybe_route_managed_krea( except Exception as exc: # noqa: BLE001 logger.warning("Managed Krea routing failed: %s", exc) return _provider_error(f"Managed Krea generation error: {exc}", "provider_exception") - if not isinstance(result, dict): - return _provider_error("Krea provider returned a non-dict result", "provider_contract") - return json.dumps(result) + return _provider_result(result, "Krea provider returned a non-dict result") -def _confine_source_images( - image_url, reference_image_urls, task_id, *, permitted: tuple = ("image",) -): +def _confine_source_images(image_url, reference_image_urls, task_id, *, permitted: tuple = ("image",)): """Resolve path-like sources to ``data:`` URLs under a non-local terminal backend. Goes through ``tools.image_source`` (in-sandbox exec-read, media-cache host reads, @@ -887,8 +778,7 @@ def _confine_source_images( sources. URLs/data: pass through; the local backend is a no-op (providers keep host reads). Returns ``(image_url, reference_image_urls, error_json_or_None)``. """ - backend = (os.getenv("TERMINAL_ENV") or "local").strip().lower() - if backend in ("", "local"): + if (os.getenv("TERMINAL_ENV") or "local").strip().lower() in ("", "local"): return image_url, reference_image_urls, None from model_tools import _run_async @@ -904,8 +794,7 @@ def _confine_source_images( reference_image_urls = [resolve(r) if isinstance(r, str) else r for r in reference_image_urls] except ImageResolutionError as exc: return image_url, reference_image_urls, _provider_error( - f"Could not read source image: {exc}", type(exc).__name__, - ) + f"Could not read source image: {exc}", type(exc).__name__) return image_url, reference_image_urls, None @@ -915,8 +804,6 @@ def _handle_image_generate(args, **kw): return tool_error("prompt is required for image generation") aspect_ratio = args.get("aspect_ratio", DEFAULT_ASPECT_RATIO) upscale = args.get("upscale") - if not isinstance(upscale, bool): - upscale = None task_id = kw.get("task_id") # Confinement chokepoint BEFORE any dispatch: plugin, managed Krea and in-tree FAL @@ -929,7 +816,8 @@ def _handle_image_generate(args, **kw): # Order matters: explicit plugin provider (incl. provider == "krea"), then # model-driven managed Krea interception (only when no provider is set, so # the BYO/direct FAL path stays untouched), then the in-tree FAL pipeline. - sources = dict(image_url=image_url, reference_image_urls=reference_image_urls, upscale=upscale) + sources = dict(image_url=image_url, reference_image_urls=reference_image_urls, + upscale=upscale if isinstance(upscale, bool) else None) raw = None for route in (_dispatch_to_plugin_provider, _maybe_route_managed_krea, image_generate_tool): raw = route(prompt, aspect_ratio, **sources) @@ -938,11 +826,10 @@ def _handle_image_generate(args, **kw): return _postprocess_image_generate_result(raw, task_id=task_id) -# --------------------------------------------------------------------------- -# Dynamic schema — reflect the active backend's image-to-image capability -# --------------------------------------------------------------------------- +# --- Dynamic schema — reflect the active backend's image-to-image capability --- # Telling the model up front whether it can edit saves a wasted turn. Memoized by # config.yaml mtime in model_tools.get_tool_definitions(), so it rebuilds on switch. +_NO_CAPABILITIES = {"modalities": ["text"], "max_reference_images": 0, "supports_upscale": False} def _active_image_capabilities() -> Dict[str, Any]: @@ -952,7 +839,7 @@ def _active_image_capabilities() -> Dict[str, Any]: FAL catalog. Fail-closed: an undeclared capability is advertised as absent (an under-declaring provider is that provider's bug, not a safety problem). """ - info: Dict[str, Any] = {"modalities": ["text"], "max_reference_images": 0, "supports_upscale": False} + info: Dict[str, Any] = dict(_NO_CAPABILITIES) configured_provider = _read_configured_image_provider() if configured_provider and configured_provider != "fal": @@ -988,7 +875,6 @@ def _active_image_capabilities() -> Dict[str, Any]: info["supports_upscale"] = True except Exception: # noqa: BLE001 pass - return info @@ -1022,11 +908,10 @@ def _build_dynamic_image_schema() -> Dict[str, Any]: "file path; reference it in your response using the current " "platform's file-delivery convention." ) - try: info = _active_image_capabilities() except Exception: # noqa: BLE001 - info = {"modalities": ["text"], "max_reference_images": 0, "supports_upscale": False} + info = dict(_NO_CAPABILITIES) max_refs = int(info.get("max_reference_images") or 0) can_edit = "image" in set(info.get("modalities") or ["text"]) @@ -1034,7 +919,6 @@ def _build_dynamic_image_schema() -> Dict[str, Any]: properties: Dict[str, Any] = { "prompt": static_props["prompt"], "aspect_ratio": static_props["aspect_ratio"], } - if can_edit: edit_clause = ", or edit / transform an existing image by passing image_url" properties["image_url"] = _IMAGE_URL_PARAM @@ -1051,17 +935,12 @@ def _build_dynamic_image_schema() -> Dict[str, Any]: } else: edit_clause = " (text-to-image only — the active model cannot edit existing images)" - if info.get("supports_upscale"): properties["upscale"] = _UPSCALE_PARAM return { "description": base_desc.format(edit_clause=edit_clause), - "parameters": { - "type": "object", - "properties": properties, - "required": ["prompt"], - }, + "parameters": {"type": "object", "properties": properties, "required": ["prompt"]}, } diff --git a/tools/image_source.py b/tools/image_source.py index 3158d32b1e..5d73c0a100 100644 --- a/tools/image_source.py +++ b/tools/image_source.py @@ -1,25 +1,13 @@ """Single resolver for every media source -> bytes + mime. All source handling (data:/http(s)/file/local/container) funnels through -:func:`resolve_image_source` so size and magic-byte checks are enforced exactly -once. Returns raw bytes (not a path): the downstream step is base64 -> data URL. +:func:`resolve_image_source` so size and magic-byte checks are enforced exactly once. +Images are the default; callers whose argument takes video opt in via ``permitted=("video",)``. -Images are the default. Callers whose argument takes video opt in via -``permitted=("video",)``: same confinement and credential-guard pipeline, only -the final type check differs (extension table + mp4 magic sniff). - -Security (terminal-backend confinement, GHSA-gpxw-6wxv-w3qq): under a non-local -backend the file tools are confined to the sandbox, so vision must be too: - - * local backend -> read any host path (chosen posture) - * non-local backend: - path in a media cache -> host-read (gateway/download caches live on the - host and are bind-mounted into the sandbox) - path anywhere else -> read *inside the sandbox* via exec-read (the - agent can already ``cat`` any container file) - -So a prompt-injected ``vision_analyze('/etc/passwd')`` under Docker reads the -container's file, never the host's, while container-only images stay deliverable. +Security (terminal-backend confinement, GHSA-gpxw-6wxv-w3qq): under a non-local backend +vision is confined like the file tools: local -> read any host path; non-local -> host-read +only inside a media cache (bind-mounted into the sandbox), anything else is exec-read +*inside the sandbox*, so an injected ``vision_analyze('/etc/passwd')`` never reads the host. """ from __future__ import annotations @@ -31,9 +19,8 @@ from dataclasses import dataclass from pathlib import Path from typing import Optional -# Raw-bytes INGEST budget: deliberately the 50MB download cap, NOT the 20MB -# provider payload cap — that one is enforced post-resize at the call sites, and -# a 20-50MB photo must still reach the resizer. +# Raw-bytes INGEST budget: deliberately the 50MB download cap, NOT the 20MB provider +# payload cap — that one is enforced post-resize at the call sites. _MAX_INGEST_BYTES = 50 * 1024 * 1024 @@ -43,24 +30,11 @@ class ImageResolutionError(Exception): self.src, self.origin = src, origin -class UnsupportedScheme(ImageResolutionError): - pass - - -class SourceUnsafe(ImageResolutionError): # SSRF / path-allowlist - pass - - -class SourceTooLarge(ImageResolutionError): - pass - - -class SourceNotFound(ImageResolutionError): - pass - - -class NotAnImage(ImageResolutionError): - pass +class UnsupportedScheme(ImageResolutionError): ... +class SourceUnsafe(ImageResolutionError): ... # SSRF / path-allowlist +class SourceTooLarge(ImageResolutionError): ... +class SourceNotFound(ImageResolutionError): ... +class NotAnImage(ImageResolutionError): ... @dataclass @@ -79,12 +53,7 @@ class ResolvedImage: _SCHEME_RE = re.compile(r"^[A-Za-z][A-Za-z0-9+.\-]*://") -async def resolve_image_source( - src: str, - ctx: ResolveContext, - *, - permitted: tuple = ("image",), -) -> ResolvedImage: +async def resolve_image_source(src: str, ctx: ResolveContext, *, permitted: tuple = ("image",)) -> ResolvedImage: if not isinstance(src, str) or not src.strip(): raise SourceNotFound("image_url is required", src=str(src)) s = src.strip() @@ -96,7 +65,6 @@ async def resolve_image_source( if reason: raise SourceUnsafe(reason, src=s) return _finalize(await _download_to_bytes(s), "", "http", s, permitted) - if _SCHEME_RE.match(s) and not s.lower().startswith("file://"): raise UnsupportedScheme( "Unrecognized image source scheme. Use an http(s) URL, a local " @@ -104,16 +72,14 @@ async def resolve_image_source( src=s, ) - # Everything else is a filesystem path — including bare relative names - # like "pic.png" (a path-shape gate here regressed them once). + # Everything else is a filesystem path — including bare relative names like "pic.png" + # (a path-shape gate here regressed them once). candidate = s[len("file://"):] if s.lower().startswith("file://") else s p = Path(os.path.expanduser(candidate)) host_target = _permitted_host_read_target(p, ctx) if host_target is not None and host_target.is_file(): - # Shared credential-read guard: refuse secret-bearing files (.env, - # auth.json) with a specific error rather than relying on the magic - # sniff to reject them incidentally. Guard import is best-effort; a - # real block always propagates. + # Shared credential-read guard: refuse secret-bearing files (.env, auth.json) with a + # specific error. Guard import is best-effort; a real block always propagates. try: from agent.file_safety import raise_if_read_blocked except Exception: # noqa: BLE001 — guard unavailable: proceed @@ -147,9 +113,8 @@ def _resolve_data_url(s: str) -> tuple[bytes, str]: def _http_block_reason(url: str) -> Optional[str]: - """Block reason, or None when allowed. Refuses policy-blocked URLs BEFORE any - network I/O; ``_download_image`` re-checks per attempt and against the final - redirect target — the second evaluation is intentional, not redundant.""" + """Block reason, or None when allowed. Refuses policy-blocked URLs BEFORE any network I/O; + ``_download_image`` re-checks per attempt and against the final redirect target (intentional).""" from tools.url_safety import is_safe_url from tools.website_policy import check_website_access @@ -163,7 +128,6 @@ def _http_block_reason(url: str) -> Optional[str]: async def _download_to_bytes(url: str) -> bytes: import tempfile - from tools.vision_tools import _download_image with tempfile.NamedTemporaryFile(suffix=".img", delete=False) as tf: @@ -183,30 +147,27 @@ def _is_local_terminal_backend() -> bool: return os.getenv("TERMINAL_ENV", "local").strip().lower() in ("local", "") -def _media_cache_roots() -> list: - """Host-side media caches: the only host paths vision may read under a - non-local backend (gateway inbound media + the tools' own download temp dirs).""" - from hermes_constants import get_hermes_home +# Host-side media caches: the only host paths vision may read under a non-local backend +# (gateway inbound media + the tools' own download temp dirs). +_MEDIA_CACHE_SUBDIRS = ( + "cache", # cache/images, cache/vision, cache/video(s), cache/audio + "images", # desktop/clipboard/PDF uploads (tui_gateway) + "image_cache", "audio_cache", "video_cache", "temp_vision_images", "temp_video_files", +) + +def _media_cache_roots() -> list: + from hermes_constants import get_hermes_home home = get_hermes_home() - return [ - home / "cache", # cache/images, cache/vision, cache/video(s), cache/audio - home / "images", # desktop/clipboard/PDF uploads (tui_gateway) - home / "image_cache", - home / "audio_cache", - home / "video_cache", - home / "temp_vision_images", - home / "temp_video_files", - ] + return [home / sub for sub in _MEDIA_CACHE_SUBDIRS] def _permitted_host_read_target(p: Path, ctx: ResolveContext) -> Optional[Path]: """Host path to read, or ``None`` if a host read is not permitted. - Local backend: any path. Non-local: only paths resolving inside a media - cache root (a container-visible cache path is first translated back to its - host mount); anything else returns ``None`` so the caller exec-reads inside - the sandbox instead of touching the host filesystem. + Local backend: any path. Non-local: only paths resolving inside a media cache root (a + container-visible cache path is first translated back to its host mount); anything else + returns ``None`` so the caller exec-reads inside the sandbox instead. """ if _is_local_terminal_backend(): try: @@ -216,9 +177,8 @@ def _permitted_host_read_target(p: Path, ctx: ResolveContext) -> Optional[Path]: from tools.credential_files import from_agent_visible_cache_path - host_candidate = Path(from_agent_visible_cache_path(str(p))) try: - real = host_candidate.resolve() + real = Path(from_agent_visible_cache_path(str(p))).resolve() except Exception: # noqa: BLE001 — cannot resolve -> not a safe host read return None for root in _media_cache_roots(): @@ -235,25 +195,18 @@ def _get_active_env(task_id: Optional[str]): return None try: from tools.terminal_tool import get_active_env - return get_active_env(task_id) except Exception: return None def _ensure_container_env(task_id: Optional[str]) -> None: - """Lazily bring up the sandbox before an in-sandbox read. - - Vision never triggered environment creation, so a session whose first action - was ``vision_analyze`` on a container-only path found no active env until a - terminal command created one. Best-effort: failure leaves the env absent and - the caller hits the fail-closed error. - """ + """Lazily bring up the sandbox before an in-sandbox read (vision may be a session's first + action). Best-effort: failure leaves the env absent and the caller hits the fail-closed error.""" if not task_id: return try: from tools.terminal_tool import ensure_task_env - ensure_task_env(task_id) except Exception: pass @@ -262,19 +215,17 @@ def _ensure_container_env(task_id: Optional[str]) -> None: async def _resolve_container_fallback( p: Path, ctx: ResolveContext, src: str, permitted: tuple = ("image",) ) -> ResolvedImage: - """Read the bytes inside the sandbox; fail-closed when no env exists (a - non-cache host path under a sandbox must never leak via a host fallback). + """Read the bytes inside the sandbox; fail-closed when no env exists (a non-cache host + path under a sandbox must never leak via a host fallback). - Cold-start retry: under Docker the first exec against a fresh container can - fail (empty pipe) while an identical second call succeeds, so retry once - after a short delay. On final failure the container's own output is folded - into the error so "no such file" / "permission denied" / "never came up" - are distinguishable. + Cold-start retry: under Docker the first exec against a fresh container can fail (empty + pipe) while an identical second call succeeds, so retry once after a short delay. On final + failure the container's own output is folded into the error so "no such file" / + "permission denied" / "never came up" are distinguishable. """ import shlex _ensure_container_env(ctx.task_id) - env = _get_active_env(ctx.task_id) if env is None: raise SourceNotFound( @@ -282,13 +233,11 @@ async def _resolve_container_fallback( f"session is available to read it", src=src, origin="container") - # Bound the read INSIDE the sandbox: head -c caps at ingest-limit+1 so - # /dev/zero can't stream unbounded base64 into host memory (the +1 - # distinguishes "at the cap" from "over"). The input redirect avoids argv, so - # leading-dash paths can't parse as options; base64 -w0 is GNU-only, hence - # tr -d for BusyBox. env.execute blocks — keep it off the event loop. - qp = shlex.quote(str(p)) - cmd = f"head -c {_MAX_INGEST_BYTES + 1} < {qp} | base64 | tr -d '\\n'" + # Bound the read INSIDE the sandbox: head -c caps at ingest-limit+1 so /dev/zero can't + # stream unbounded base64 into host memory (the +1 distinguishes "at the cap" from + # "over"). The input redirect avoids argv, so leading-dash paths can't parse as options; + # base64 -w0 is GNU-only, hence tr -d for BusyBox. env.execute blocks — keep it off the loop. + cmd = f"head -c {_MAX_INGEST_BYTES + 1} < {shlex.quote(str(p))} | base64 | tr -d '\\n'" last_res: dict = {"returncode": 1, "output": ""} for attempt in range(2): @@ -301,9 +250,7 @@ async def _resolve_container_fallback( diag = (last_res.get("output") or "").strip().splitlines() first = next((ln.strip() for ln in diag if ln.strip()), "") suffix = f" ({first[:200]})" if first else "" - raise SourceNotFound( - f"could not read '{p}' inside the sandbox{suffix}", - src=src, origin="container") + raise SourceNotFound(f"could not read '{p}' inside the sandbox{suffix}", src=src, origin="container") try: data = base64.b64decode(last_res.get("output", ""), validate=True) except Exception as exc: @@ -318,9 +265,9 @@ def _finalize( ) -> ResolvedImage: """Chokepoint: 50MB ingest cap + type check. - Images are typed by magic bytes. Video (opt-in) is typed by extension plus - an mp4 container sniff — sufficient because every downstream consumer - re-validates, so a wrong guess is a clean rejection there, not a hole here. + Images are typed by magic bytes. Video (opt-in) is typed by extension plus an mp4 + container sniff — sufficient because every downstream consumer re-validates, so a wrong + guess is a clean rejection there, not a hole here. """ from tools.vision_tools import _detect_image_mime_type_from_bytes @@ -332,25 +279,21 @@ def _finalize( if "image" not in permitted: raise NotAnImage("source is an image, but this argument takes a video", src=src, origin=origin) return ResolvedImage(data=data, mime=sniffed, origin=origin) - if "image" in permitted and b" Optional[str]: - """Video MIME from the extension table, else the ISO base-media ``ftyp`` - magic at offset 4 (covers extensionless data: URLs / query-string URLs).""" + """Video MIME from the extension table, else the ISO base-media ``ftyp`` magic at + offset 4 (covers extensionless data: URLs / query-string URLs).""" from urllib.parse import urlsplit - from tools.vision_tools import _detect_video_mime_type path_part = urlsplit(src).path if _SCHEME_RE.match(src) else src @@ -367,18 +310,15 @@ async def resolve_local_source_to_data_url( ) -> str: """Convert a path-like media source into a ``data:`` URL via the resolver. - Dispatch-layer chokepoint for generation tools: providers historically read - model-supplied local paths off the HOST regardless of backend — broken under - a sandbox (the file lives there) and inconsistent with the confinement model - vision enforces. URL-shaped sources (http/https/data) pass through untouched. + Dispatch-layer chokepoint for generation tools: providers historically read model-supplied + local paths off the HOST regardless of backend — broken under a sandbox and inconsistent + with vision's confinement. URL-shaped sources (http/https/data) pass through untouched. Callers apply this only under a non-local backend (local keeps host reads). """ s = (src or "").strip() if not s or s.lower().startswith(("http://", "https://", "data:")): return src - resolved = await resolve_image_source( - s, ResolveContext(task_id=task_id), permitted=permitted - ) + resolved = await resolve_image_source(s, ResolveContext(task_id=task_id), permitted=permitted) encoded = base64.b64encode(resolved.data).decode("ascii") mime = resolved.mime or "application/octet-stream" return f"data:{mime};base64,{encoded}"