From 6906a6a52ed61f5104aefb9eaa0bde715d08321c Mon Sep 17 00:00:00 2001 From: AceronX Date: Tue, 21 Jul 2026 09:05:45 -0400 Subject: [PATCH] fix: harden agent, benchmark, and UI against security, correctness, and stability issues --- .github/scripts/round_manager.py | 27 ++- agent/agent.py | 34 ++- .../src/imagent_bench/agent_loader.py | 6 +- imagent-bench/src/imagent_bench/models.py | 12 +- imagent-bench/src/imagent_bench/runner.py | 28 ++- imagent-bench/src/imagent_bench/scoring.py | 42 +++- imagent-bench/tests/test_agent_loader.py | 38 +++ imagent-bench/tests/test_models.py | 18 ++ imagent-bench/tests/test_runner.py | 50 +++- imagent-bench/tests/test_scoring.py | 64 ++++++ imagent-ui/app/api/openrouter/verify/route.ts | 16 ++ .../app/api/playground/generate/route.ts | 98 ++++++-- .../playground/runs/[runId]/image/route.ts | 11 +- imagent-ui/app/components/GenerationChat.tsx | 71 ------ imagent-ui/app/styles.css | 81 ------- imagent-ui/deploy/Caddyfile | 15 ++ imagent-ui/lib/playground.ts | 70 +++++- imagent-ui/lib/security.ts | 98 ++++++++ imagent-ui/next.config.mjs | 40 +++- imagent_runtime/cli.py | 26 ++- pyproject.toml | 2 +- tests/test_image_agent.py | 217 ++++++++++++++++++ 22 files changed, 856 insertions(+), 208 deletions(-) create mode 100644 imagent-bench/tests/test_agent_loader.py create mode 100644 imagent-bench/tests/test_models.py create mode 100644 imagent-ui/lib/security.ts diff --git a/.github/scripts/round_manager.py b/.github/scripts/round_manager.py index 1f077b6..db06fd6 100644 --- a/.github/scripts/round_manager.py +++ b/.github/scripts/round_manager.py @@ -351,7 +351,12 @@ def evaluate_pr(pr: dict[str, Any], baseline_score: float, benchmark_config: str baseline_commit = os.environ.get("IMAGENT_BASELINE_COMMIT", "").strip() if baseline_commit: command.extend(["--baseline-commit", baseline_commit]) - completed = subprocess.run([str(part) for part in command], text=True, cwd=candidate_dir) + completed = subprocess.run( + [str(part) for part in command], + text=True, + cwd=candidate_dir, + env=_candidate_subprocess_env(os.environ), + ) finally: subprocess.run(["git", "worktree", "remove", "--force", str(candidate_dir)], check=False) @@ -366,6 +371,20 @@ def evaluate_pr(pr: dict[str, Any], baseline_score: float, benchmark_config: str return CandidateResult(pr=pr, report=report, score=score, delta=delta, head_sha=head_sha) +def _candidate_subprocess_env(base_env: dict) -> dict: + """Return a copy of base_env safe to expose to untrusted candidate code. + + The candidate benchmark subprocess executes contributor-supplied + ``agent/agent.py``, so secrets that only the orchestrator needs must not be + inherited. Strip the GitHub tokens to avoid exfiltration, but keep + ``OPENROUTER_API_KEY`` because the candidate needs it to generate images. + """ + env = dict(base_env) + for name in ("GITHUB_TOKEN", "GH_TOKEN"): + env.pop(name, None) + return env + + def candidate_file_failures(filenames: list[str]) -> list[str]: failures: list[str] = [] if CANDIDATE_AGENT_PATH not in filenames: @@ -402,6 +421,12 @@ def policy_reasons(report: dict[str, Any]) -> list[str]: def is_merge_improvement_reason(reason: str) -> bool: + # NOTE: This is intentionally coupled to imagent-bench's human-readable policy + # reason strings (e.g. "score improvement 0.50 is below required 1.00"). A + # structured status field emitted by the benchmark would be more robust than + # matching English text; until that exists, guard against non-string input. + if not isinstance(reason, str): + return False return reason.startswith("score improvement ") and " is below required " in reason diff --git a/agent/agent.py b/agent/agent.py index 793bf46..42da27a 100644 --- a/agent/agent.py +++ b/agent/agent.py @@ -6,6 +6,7 @@ import os import re import urllib.error +import urllib.parse import urllib.request from pathlib import Path from typing import Any @@ -121,7 +122,7 @@ def _build_context(self, case: dict[str, Any]) -> dict[str, Any]: ) sections = self._sections_from_prompt(prompt) or assets.get("sections", []) required_text = [title, *sections, *assets.get("required_text", [])] - if "reason" in case.get("allowed_tools", []): + if "reason" in (case.get("allowed_tools") or []): display = self._reasoning_display(prompt) if display: required_text.extend([display["answer"], display["display"]]) @@ -161,19 +162,23 @@ def _build_generation_prompt(self, case: dict[str, Any], context: dict[str, Any] return "\n".join(line for line in lines if line) def _read_assets(self, values: Any) -> dict[str, Any]: - for value in values or []: - path = self._case_path(value) - if path.suffix.lower() != ".json": - continue + json_paths = [ + path + for path in (self._case_path(value) for value in values or []) + if path.suffix.lower() == ".json" + ] + merged: dict[str, Any] = {"title": "", "sections": [], "required_text": []} + for path in sorted(json_paths, key=lambda item: item.name): data = json.loads(path.read_text(encoding="utf-8")) if not isinstance(data, dict): continue - return { - "title": str(data.get("title") or ""), - "sections": [str(item) for item in data.get("sections", []) if str(item).strip()], - "required_text": [str(item) for item in data.get("required_text", []) if str(item).strip()], + merged = { + "title": str(data.get("title") or "") or merged["title"], + "sections": [str(item) for item in data.get("sections", []) if str(item).strip()] or merged["sections"], + "required_text": [str(item) for item in data.get("required_text", []) if str(item).strip()] + or merged["required_text"], } - return {"title": "", "sections": [], "required_text": []} + return merged def _read_search_facts(self, values: Any, prompt: str) -> list[str]: prompt_words = set(re.findall(r"[a-z0-9]+", prompt.lower())) @@ -277,6 +282,10 @@ def _image_bytes_from_url(url: str, fallback_media_type: str, backend_config: di except ValueError as exc: raise OpenRouterImageError("OpenRouter image data URL included invalid base64 bytes") from exc + scheme = urllib.parse.urlparse(url).scheme.lower() + if scheme not in {"http", "https"}: + raise OpenRouterImageError(f"OpenRouter image URL scheme is not allowed: {scheme or url!r}") + request = urllib.request.Request(url, method="GET") try: with urllib.request.urlopen(request, timeout=int(backend_config.get("timeout_seconds", 240))) as response: @@ -320,7 +329,10 @@ def _extract_expression(prompt: str) -> str: for match in re.finditer(rf"(? tuple[Any, dict[str, Any]]: raise AgentLoadError("agent manifest entrypoint must include module and attribute") _clear_candidate_modules(module_name) - sys.path.insert(0, str(repository)) + repository_path = str(repository) + # Avoid unbounded sys.path growth across repeated loads; keep the entry so + # candidate agents can still perform lazy imports later. + if repository_path not in sys.path: + sys.path.insert(0, repository_path) try: module = importlib.import_module(module_name) agent_cls = getattr(module, attribute) diff --git a/imagent-bench/src/imagent_bench/models.py b/imagent-bench/src/imagent_bench/models.py index 545a0ff..cbaee5b 100644 --- a/imagent-bench/src/imagent_bench/models.py +++ b/imagent-bench/src/imagent_bench/models.py @@ -4,6 +4,14 @@ from typing import Any +def _safe_int(value: Any) -> int: + # Non-integer ids/seeds should not abort the whole run; fall back to 0. + try: + return int(value) + except (TypeError, ValueError): + return 0 + + @dataclass(frozen=True) class BenchmarkConfig: benchmark_version: str @@ -32,13 +40,13 @@ def from_record(cls, record: dict[str, Any]) -> "BenchmarkCase": raw_id = record.get("id") or record.get("run_id") or record.get("ID") if raw_id is None: raise ValueError("benchmark case is missing id") - numeric_id = int(record.get("ID", record.get("numeric_id", 0)) or 0) + numeric_id = _safe_int(record.get("ID", record.get("numeric_id", 0)) or 0) return cls( id=str(raw_id), numeric_id=numeric_id, prompt=str(record["prompt"]), capability=str(record.get("capability", "plan")), - seed=int(record.get("seed", 0)), + seed=_safe_int(record.get("seed", 0)), allowed_tools=[str(value) for value in record.get("allowed_tools", [])], expected=dict(record.get("expected", {}) or {}), assets=[str(value) for value in record.get("assets", []) or []], diff --git a/imagent-bench/src/imagent_bench/runner.py b/imagent-bench/src/imagent_bench/runner.py index 5618a5d..2325a4c 100644 --- a/imagent-bench/src/imagent_bench/runner.py +++ b/imagent-bench/src/imagent_bench/runner.py @@ -251,7 +251,9 @@ def _ranking( } delta = round(candidate_score - baseline_score, 6) - label = "score-regression" + # Only a negative delta is a regression. A positive or zero delta that is too + # small to clear the smallest improvement threshold is a non-regression. + label = "score-regression" if delta < 0 else "no-significant-change" for name, threshold in sorted(thresholds.items(), key=lambda item: item[1]): if delta >= threshold: label = f"improvement-{name}" @@ -296,7 +298,10 @@ def _percentile(values: list[float], percentile: int) -> float: if len(values) == 1: return round(values[0], 3) ordered = sorted(values) - return round(float(statistics.quantiles(ordered, n=100, method="inclusive")[percentile - 1]), 3) + quantiles = statistics.quantiles(ordered, n=100, method="inclusive") + # quantiles has 99 cut points (indices 0..98); clamp so p<=0 and p=100 stay in range. + index = min(max(percentile - 1, 0), len(quantiles) - 1) + return round(float(quantiles[index]), 3) def _utc_now() -> str: @@ -305,7 +310,8 @@ def _utc_now() -> str: def _append_log(path: Path, message: str) -> None: path.parent.mkdir(parents=True, exist_ok=True) - path.open("a", encoding="utf-8").write(json.dumps({"ts": _utc_now(), "message": message}) + "\n") + with path.open("a", encoding="utf-8") as f: + f.write(json.dumps({"ts": _utc_now(), "message": message}) + "\n") def _validate_local_repository_commit(repository: str | Path, requested_commit: str | None, actual_commit: str) -> None: @@ -359,6 +365,7 @@ def _normalize_repository_identifier(value: str) -> str: text = str(value).strip() if not text: return "" + text = _strip_url_credentials(text) for pattern in ( r"^https?://github\.com/(?P[A-Za-z0-9_.-]+)/(?P[A-Za-z0-9_.-]+?)(?:\.git)?/?$", r"^ssh://git@github\.com/(?P[A-Za-z0-9_.-]+)/(?P[A-Za-z0-9_.-]+?)(?:\.git)?/?$", @@ -371,6 +378,21 @@ def _normalize_repository_identifier(value: str) -> str: return text.removesuffix(".git") +def _strip_url_credentials(text: str) -> str: + # Remove embedded userinfo (e.g. tokens or user:pass) from a URL authority so + # credentials never leak into the normalized identifier or the report. Only + # scheme-based URLs carry userinfo before an "@" in the authority; the SCP + # form (git@github.com:owner/repo.git) has no "//" and is left untouched. + match = re.match(r"^(?P[a-zA-Z][a-zA-Z0-9+.-]*://)(?P.*)$", text) + if not match: + return text + rest = match.group("rest") + authority, sep, tail = rest.partition("/") + if "@" in authority: + authority = authority.rsplit("@", 1)[1] + return f"{match.group('scheme')}{authority}{sep}{tail}" + + def _github_report_metadata() -> tuple[dict[str, Any] | None, dict[str, Any] | None]: payload = _github_event_payload() if not payload: diff --git a/imagent-bench/src/imagent_bench/scoring.py b/imagent-bench/src/imagent_bench/scoring.py index c5a9203..acebd0a 100644 --- a/imagent-bench/src/imagent_bench/scoring.py +++ b/imagent-bench/src/imagent_bench/scoring.py @@ -6,6 +6,7 @@ import mimetypes import os import re +import time import xml.etree.ElementTree as ET from pathlib import Path from typing import Any @@ -37,6 +38,9 @@ def evaluate_case( judge_config = judge_config or {} provider = str(judge_config.get("provider", "mock_text")).strip().lower() if provider in {"openrouter", "openrouter_vision"}: + # Intentional design: when a vision judge is configured, deterministic + # image_contains checks are skipped in favor of the model's holistic score. + # This is not a bug; text-based checks can't reliably validate a raster image. judge_result = evaluate_openrouter_vision(image_path, prompt=prompt, config=judge_config) return [], float(judge_result["overall_score"]), judge_result @@ -112,14 +116,33 @@ def evaluate_openrouter_vision(image_path: Path, *, prompt: str, config: dict[st headers=headers, method="POST", ) - try: - with urllib.request.urlopen(request, timeout=int(config.get("timeout_seconds", 180))) as response: - response_payload = json.loads(response.read().decode("utf-8")) - except urllib.error.HTTPError as exc: - body = exc.read().decode("utf-8", errors="replace") - raise JudgeError(f"OpenRouter judge HTTP {exc.code}: {body}") from exc - except urllib.error.URLError as exc: - raise JudgeError(f"OpenRouter judge request failed: {exc}") from exc + timeout = int(config.get("timeout_seconds", 180)) + max_retries = int(config.get("max_retries", 3)) + delay_seconds = float(config.get("retry_backoff_seconds", 1.0)) + response_payload: dict[str, Any] | None = None + # Retry transient failures (429/5xx and network errors) with exponential + # backoff, mirroring backends/openrouter_backend.py, so a single blip does + # not abort the whole benchmark. + for attempt in range(max_retries + 1): + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + response_payload = json.loads(response.read().decode("utf-8")) + break + except urllib.error.HTTPError as exc: + if exc.code in {429, 500, 502, 503, 504} and attempt < max_retries: + time.sleep(delay_seconds) + delay_seconds *= 2 + continue + body = exc.read().decode("utf-8", errors="replace") + raise JudgeError(f"OpenRouter judge HTTP {exc.code}: {body}") from exc + except urllib.error.URLError as exc: + if attempt < max_retries: + time.sleep(delay_seconds) + delay_seconds *= 2 + continue + raise JudgeError(f"OpenRouter judge request failed: {exc}") from exc + if response_payload is None: # pragma: no cover - defensive, loop always sets or raises + raise JudgeError("OpenRouter judge request failed after retries") content = _openrouter_message_content(response_payload) parsed = _parse_judge_json(content) @@ -146,6 +169,9 @@ def evaluate_openrouter_vision(image_path: Path, *, prompt: str, config: dict[st def _read_visible_text(image_path: Path) -> str: + # Intentional design limitation: only .svg and .txt expose extractable text. + # Raster formats (.png/.jpg/...) return "" because text can't be recovered + # without OCR, which is out of scope here. This is not a bug. if not image_path.exists(): return "" if image_path.suffix.lower() == ".svg": diff --git a/imagent-bench/tests/test_agent_loader.py b/imagent-bench/tests/test_agent_loader.py new file mode 100644 index 0000000..c769b85 --- /dev/null +++ b/imagent-bench/tests/test_agent_loader.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +from imagent_bench.agent_loader import load_agent + + +def _imagent_repository() -> Path: + workspace = Path(__file__).resolve().parents[2] + sibling_repo = workspace / "imagent" + if (sibling_repo / "agent" / "agent.py").exists(): + return sibling_repo + if (workspace / "agent" / "agent.py").exists(): + return workspace + return sibling_repo + + +def test_load_agent_does_not_duplicate_sys_path_entries() -> None: + repository = _imagent_repository() + if not (repository / "agent" / "agent.py").exists(): + pytest.skip("candidate imagent repository not available") + + repository_path = str(repository) + before = sys.path.count(repository_path) + + load_agent(repository) + after_first = sys.path.count(repository_path) + + load_agent(repository) + after_second = sys.path.count(repository_path) + + # A single load adds at most one entry, and a repeated load must never + # add a duplicate regardless of any pre-existing sys.path state. + assert after_first <= before + 1 + assert after_second == after_first diff --git a/imagent-bench/tests/test_models.py b/imagent-bench/tests/test_models.py new file mode 100644 index 0000000..8fdb298 --- /dev/null +++ b/imagent-bench/tests/test_models.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from imagent_bench.models import BenchmarkCase + + +def test_from_record_tolerates_non_integer_id_and_seed() -> None: + case = BenchmarkCase.from_record( + { + "id": "case-1", + "ID": "not-a-number", + "seed": "also-not-a-number", + "prompt": "draw a cat", + } + ) + + assert case.id == "case-1" + assert case.numeric_id == 0 + assert case.seed == 0 diff --git a/imagent-bench/tests/test_runner.py b/imagent-bench/tests/test_runner.py index e9aaae4..33982eb 100644 --- a/imagent-bench/tests/test_runner.py +++ b/imagent-bench/tests/test_runner.py @@ -7,7 +7,14 @@ import pytest -from imagent_bench.runner import BenchmarkRunError, _case_status, _normalize_repository_identifier, run +from imagent_bench.runner import ( + BenchmarkRunError, + _case_status, + _normalize_repository_identifier, + _percentile, + _ranking, + run, +) def _imagent_repository() -> Path: @@ -178,6 +185,47 @@ def test_normalize_repository_identifier_handles_common_github_urls() -> None: assert _normalize_repository_identifier("git@github.com:imagent-ai/imagent-bench.git") == "imagent-ai/imagent-bench" +def test_normalize_repository_identifier_strips_embedded_credentials() -> None: + token_url = "https://ghp_secretToken123@github.com/o/r.git" + userpass_url = "https://u:p@github.com/o/r.git" + + assert _normalize_repository_identifier(token_url) == "o/r" + assert _normalize_repository_identifier(userpass_url) == "o/r" + # No credential substring may survive into the normalized identifier. + assert "ghp_secretToken123" not in _normalize_repository_identifier(token_url) + assert "@" not in _normalize_repository_identifier(token_url) + assert "u:p" not in _normalize_repository_identifier(userpass_url) + # SCP form must remain unaffected by the credential stripping. + assert _normalize_repository_identifier("git@github.com:o/r.git") == "o/r" + + +def test_percentile_handles_p100_and_p95_without_indexerror() -> None: + values = [float(v) for v in range(1, 101)] + + # p=100 previously raised IndexError; it must now clamp to the top cut point. + p100 = _percentile(values, 100) + p95 = _percentile(values, 95) + + assert isinstance(p100, float) + assert p95 <= p100 <= 100.0 + assert p95 == pytest.approx(95.0, abs=1.0) + + +def test_ranking_labels_small_positive_delta_as_non_regression() -> None: + ranking = _ranking({}, candidate_score=95.5, baseline_score=95.0, baseline_commit="abc") + + assert ranking["delta"] == 0.5 + assert ranking["label"] == "no-significant-change" + assert ranking["merge_eligible"] is False + + +def test_ranking_labels_negative_delta_as_regression() -> None: + ranking = _ranking({}, candidate_score=94.0, baseline_score=95.0, baseline_commit="abc") + + assert ranking["delta"] == -1.0 + assert ranking["label"] == "score-regression" + + def test_runner_rejects_local_repository_commit_mismatch(tmp_path: Path) -> None: repository = _imagent_repository() config = Path(__file__).resolve().parents[1] / "configs" / "official.json" diff --git a/imagent-bench/tests/test_scoring.py b/imagent-bench/tests/test_scoring.py index 6c03f38..71bc9fa 100644 --- a/imagent-bench/tests/test_scoring.py +++ b/imagent-bench/tests/test_scoring.py @@ -1,6 +1,8 @@ from __future__ import annotations +import io import json +import urllib.error from pathlib import Path from imagent_bench.scoring import evaluate_openrouter_vision @@ -59,3 +61,65 @@ def read(self) -> bytes: assert result["dimensions"] == {"prompt_alignment": 90.0, "visual_quality": 80.0} assert result["cost_usd"] == 0.002 assert result["judge"]["model"] == "judge/model" + + +def test_openrouter_vision_judge_retries_on_transient_429( + monkeypatch, tmp_path: Path +) -> None: # noqa: ANN001 + image_path = tmp_path / "image.png" + image_path.write_bytes(b"fake-image") + + class FakeResponse: + def __enter__(self): # noqa: ANN001 + return self + + def __exit__(self, *args): # noqa: ANN002 + return None + + def read(self) -> bytes: + return json.dumps( + { + "model": "judge/model", + "choices": [ + { + "message": { + "content": json.dumps( + {"scores": {"prompt_alignment": 90}, "rationale": "ok"} + ) + } + } + ], + "usage": {"cost": 0.001}, + } + ).encode("utf-8") + + calls = {"count": 0} + + def flaky_urlopen(*args, **kwargs): # noqa: ANN002, ANN003 + calls["count"] += 1 + if calls["count"] == 1: + raise urllib.error.HTTPError( + url="https://openrouter.ai", + code=429, + msg="Too Many Requests", + hdrs=None, + fp=io.BytesIO(b"rate limited"), + ) + return FakeResponse() + + monkeypatch.setenv("OPENROUTER_API_KEY", "test-key") + monkeypatch.setattr("urllib.request.urlopen", flaky_urlopen) + monkeypatch.setattr("imagent_bench.scoring.time.sleep", lambda *_: None) + + result = evaluate_openrouter_vision( + image_path, + prompt="Create a clean product image.", + config={ + "provider": "openrouter_vision", + "model": "judge/model", + "dimensions": {"prompt_alignment": 1.0}, + }, + ) + + assert calls["count"] == 2 + assert result["overall_score"] == 90.0 diff --git a/imagent-ui/app/api/openrouter/verify/route.ts b/imagent-ui/app/api/openrouter/verify/route.ts index a8f425b..50260bf 100644 --- a/imagent-ui/app/api/openrouter/verify/route.ts +++ b/imagent-ui/app/api/openrouter/verify/route.ts @@ -4,8 +4,11 @@ import { IMAGENT_GENERATION_MODEL_ID, IMAGENT_GENERATION_MODEL_OPTION } from "@/lib/models"; +import { getClientIp, isRequestAuthorized, rateLimit } from "@/lib/security"; import { resolvePublicSiteUrl } from "@/lib/site"; +const VERIFY_RATE_LIMIT_PER_MINUTE = 20; + type VerifyRequest = { apiKey?: string; }; @@ -47,6 +50,19 @@ const OPENROUTER_KEY_URL = "https://openrouter.ai/api/v1/key"; const OPENROUTER_IMAGE_MODELS_URL = "https://openrouter.ai/api/v1/images/models"; export async function POST(request: Request) { + if (!isRequestAuthorized(request)) { + return NextResponse.json({ error: "Unauthorized." }, { status: 401 }); + } + + const clientIp = getClientIp(request); + const limit = rateLimit(`verify:${clientIp}`, VERIFY_RATE_LIMIT_PER_MINUTE); + if (!limit.allowed) { + return NextResponse.json( + { error: "Too many requests. Please slow down and try again shortly." }, + { status: 429, headers: { "Retry-After": String(limit.retryAfterSeconds) } } + ); + } + const body = await parseJson(request); const publicSiteUrl = resolvePublicSiteUrl(); const requestedApiKey = String(body.apiKey || "").trim(); diff --git a/imagent-ui/app/api/playground/generate/route.ts b/imagent-ui/app/api/playground/generate/route.ts index 267b130..5edaaf1 100644 --- a/imagent-ui/app/api/playground/generate/route.ts +++ b/imagent-ui/app/api/playground/generate/route.ts @@ -4,13 +4,23 @@ import path from "node:path"; import { promisify } from "node:util"; import { NextResponse } from "next/server"; import { + BACKGROUND_OPTIONS, contentTypeForImage, DEFAULT_GENERATION_MODEL, getResolvedPlaygroundRuntime, + pruneRunDirectories, + QUALITY_OPTIONS, resolveRunDirectory, validateRunId, writeRunManifest } from "@/lib/playground"; +import { + getClientIp, + isRequestAuthorized, + rateLimit, + releaseGenerationSlot, + tryAcquireGenerationSlot +} from "@/lib/security"; import { resolvePublicSiteUrl } from "@/lib/site"; type GenerateRequest = { @@ -42,20 +52,66 @@ type AgentResult = { const execFileAsync = promisify(execFile); +// Abuse controls for this expensive, credit-consuming endpoint. +const GENERATION_TIMEOUT_MS = 120_000; +const MAX_CONCURRENT_GENERATIONS = 2; +const RATE_LIMIT_PER_MINUTE = 10; +const MAX_PROMPT_LENGTH = 4000; + +async function parseJson(request: Request): Promise { + try { + return (await request.json()) as T; + } catch { + return null; + } +} + export async function POST(request: Request) { const started = performance.now(); - const body = (await request.json()) as GenerateRequest; + + if (!isRequestAuthorized(request)) { + return NextResponse.json({ error: "Unauthorized." }, { status: 401 }); + } + + const clientIp = getClientIp(request); + const limit = rateLimit(`generate:${clientIp}`, RATE_LIMIT_PER_MINUTE); + if (!limit.allowed) { + return NextResponse.json( + { error: "Too many requests. Please slow down and try again shortly." }, + { status: 429, headers: { "Retry-After": String(limit.retryAfterSeconds) } } + ); + } + + const body = await parseJson(request); + if (!body) { + return NextResponse.json({ error: "Invalid JSON request body." }, { status: 400 }); + } + const publicSiteUrl = resolvePublicSiteUrl(); const prompt = String(body.prompt || "").trim(); const model = DEFAULT_GENERATION_MODEL; const quality = String(body.quality || "auto").trim(); const background = String(body.background || "auto").trim(); - const runtime = await getResolvedPlaygroundRuntime(); - const apiKey = String(body.apiKey || (runtime.hasServerApiKey ? process.env.OPENROUTER_API_KEY : "") || "").trim(); if (!prompt) { return NextResponse.json({ error: "Prompt is required." }, { status: 400 }); } + if (prompt.length > MAX_PROMPT_LENGTH) { + return NextResponse.json( + { error: `Prompt must be ${MAX_PROMPT_LENGTH} characters or fewer.` }, + { status: 400 } + ); + } + if (!(QUALITY_OPTIONS as readonly string[]).includes(quality)) { + return NextResponse.json({ error: "Unsupported quality option." }, { status: 400 }); + } + if (!(BACKGROUND_OPTIONS as readonly string[]).includes(background)) { + return NextResponse.json({ error: "Unsupported background option." }, { status: 400 }); + } + + const runtime = await getResolvedPlaygroundRuntime(); + const apiKey = String(body.apiKey || (runtime.hasServerApiKey ? process.env.OPENROUTER_API_KEY : "") || "").trim(); + if (!apiKey) { return NextResponse.json({ error: "OpenRouter API key is required." }, { status: 400 }); } @@ -68,14 +124,22 @@ export async function POST(request: Request) { ); } - const runId = `ui-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; - if (!validateRunId(runId)) { - return NextResponse.json({ error: "Failed to allocate a valid run ID." }, { status: 500 }); + if (!tryAcquireGenerationSlot(MAX_CONCURRENT_GENERATIONS)) { + return NextResponse.json( + { error: "The generation queue is busy. Please try again in a moment." }, + { status: 429 } + ); } - const outputDir = resolveRunDirectory(runId); - const requestPath = path.join(outputDir, "request.json"); + + const runId = `ui-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; try { + if (!validateRunId(runId)) { + return NextResponse.json({ error: "Failed to allocate a valid run ID." }, { status: 500 }); + } + const outputDir = resolveRunDirectory(runId); + const requestPath = path.join(outputDir, "request.json"); + await fs.mkdir(outputDir, { recursive: true }); await fs.writeFile( requestPath, @@ -102,7 +166,9 @@ export async function POST(request: Request) { ...process.env, OPENROUTER_API_KEY: apiKey }, - maxBuffer: 10 * 1024 * 1024 + maxBuffer: 10 * 1024 * 1024, + timeout: GENERATION_TIMEOUT_MS, + killSignal: "SIGKILL" }); if (stderr.trim()) { console.warn(stderr); @@ -151,11 +217,13 @@ export async function POST(request: Request) { traceUrl: traceFileName ? `/api/playground/runs/${encodeURIComponent(runId)}/trace` : undefined }); } catch (error) { - return NextResponse.json( - { - error: error instanceof Error ? error.message : "Imagent generation failed." - }, - { status: 502 } - ); + // Log the detailed cause (Python stderr, filesystem paths, timeouts) on the + // server only; never leak it to the client. + console.error("Imagent generation failed:", error); + return NextResponse.json({ error: "Imagent generation failed. Please try again." }, { status: 502 }); + } finally { + releaseGenerationSlot(); + // Best-effort cap on data/agent-runs/ growth; never blocks the response path. + void pruneRunDirectories().catch(() => {}); } } diff --git a/imagent-ui/app/api/playground/runs/[runId]/image/route.ts b/imagent-ui/app/api/playground/runs/[runId]/image/route.ts index 8d8b40e..68b5207 100644 --- a/imagent-ui/app/api/playground/runs/[runId]/image/route.ts +++ b/imagent-ui/app/api/playground/runs/[runId]/image/route.ts @@ -1,7 +1,7 @@ import { promises as fs } from "node:fs"; import path from "node:path"; import { NextResponse } from "next/server"; -import { loadRunManifest, resolveRunDirectory } from "@/lib/playground"; +import { loadRunManifest, resolveRunDirectory, safeInlineContentType } from "@/lib/playground"; type RouteContext = { params: Promise<{ runId: string }>; @@ -15,13 +15,18 @@ export async function GET(_request: Request, context: RouteContext) { } const artifactPath = path.join(resolveRunDirectory(runId), manifest.imageFileName); + const isSvg = + manifest.imageMediaType === "image/svg+xml" || manifest.imageFileName.toLowerCase().endsWith(".svg"); try { const bytes = await fs.readFile(artifactPath); return new NextResponse(bytes, { headers: { "Cache-Control": "private, max-age=31536000, immutable", - "Content-Disposition": `inline; filename="${manifest.imageFileName}"`, - "Content-Type": manifest.imageMediaType + // Raster images render inline; svg is downgraded to a neutral type and + // forced to download so it cannot execute script in our origin. + "Content-Disposition": `${isSvg ? "attachment" : "inline"}; filename="${manifest.imageFileName}"`, + "Content-Type": safeInlineContentType(manifest.imageMediaType, manifest.imageFileName), + "X-Content-Type-Options": "nosniff" } }); } catch { diff --git a/imagent-ui/app/components/GenerationChat.tsx b/imagent-ui/app/components/GenerationChat.tsx index 77162df..2c48111 100644 --- a/imagent-ui/app/components/GenerationChat.tsx +++ b/imagent-ui/app/components/GenerationChat.tsx @@ -1,7 +1,6 @@ "use client"; import { - type ChangeEvent, type FormEvent, type KeyboardEvent as ReactKeyboardEvent, type MouseEvent, @@ -18,7 +17,6 @@ import { ChevronDown, Download, FileJson, - ImagePlus, KeyRound, Loader2, MessageCirclePlus, @@ -65,11 +63,6 @@ type ChatSession = { messages: ChatMessage[]; }; -type SelectedInputImage = { - name: string; - previewUrl: string; -}; - type ModalState = | { type: "settings" } | { type: "edit"; sessionId: string } @@ -151,7 +144,6 @@ export function GenerationChat() { const [sessions, setSessions] = useState([]); const [activeSessionId, setActiveSessionId] = useState(""); const [prompt, setPrompt] = useState(""); - const [selectedInputImage, setSelectedInputImage] = useState(null); const [level, setLevel] = useState("auto"); const [apiKey, setApiKey] = useState(""); const [draftApiKey, setDraftApiKey] = useState(""); @@ -173,7 +165,6 @@ export function GenerationChat() { const closeModalRef = useRef<() => void>(() => {}); const preserveModalTriggerRef = useRef(false); const promptTextareaRef = useRef(null); - const imageInputRef = useRef(null); async function loadRuntimeStatus() { try { @@ -236,14 +227,6 @@ export function GenerationChat() { textarea.style.overflowY = textarea.scrollHeight > maxHeight ? "auto" : "hidden"; }, [prompt]); - useEffect(() => { - return () => { - if (selectedInputImage?.previewUrl) { - URL.revokeObjectURL(selectedInputImage.previewUrl); - } - }; - }, [selectedInputImage?.previewUrl]); - useEffect(() => { if (!modal) { return; @@ -391,7 +374,6 @@ export function GenerationChat() { if (reusableSession) { setActiveSessionId(reusableSession.id); setPrompt(""); - clearInputImage(); setSidebarCollapsed(false); return; } @@ -400,7 +382,6 @@ export function GenerationChat() { setSessions((current) => [session, ...current]); setActiveSessionId(session.id); setPrompt(""); - clearInputImage(); setSidebarCollapsed(false); } @@ -420,31 +401,6 @@ export function GenerationChat() { } } - function openImagePicker() { - imageInputRef.current?.click(); - } - - function selectInputImage(event: ChangeEvent) { - const file = event.target.files?.[0]; - event.target.value = ""; - - if (!file || !file.type.startsWith("image/")) { - return; - } - - setSelectedInputImage({ - name: file.name, - previewUrl: URL.createObjectURL(file) - }); - } - - function clearInputImage() { - setSelectedInputImage(null); - if (imageInputRef.current) { - imageInputRef.current.value = ""; - } - } - function selectTemplate(template: string) { void sendPrompt(template); } @@ -471,7 +427,6 @@ export function GenerationChat() { appendSessionMessages(sessionId, [userMessage], titleFromPrompt(content), now); setPrompt(""); - clearInputImage(); if (runtimeStatus && !runtimeStatus.ready) { appendSessionMessages(sessionId, [ @@ -921,33 +876,7 @@ export function GenerationChat() {
- {selectedInputImage ? ( -
- {/* eslint-disable-next-line @next/next/no-img-element */} - - -
- ) : null}
- -