From 8865c975aaab5d5d755503d2d09e3a75d9eabd88 Mon Sep 17 00:00:00 2001 From: pewdiepie-archdaemon Date: Mon, 26 Oct 2026 16:53:52 +0000 Subject: [PATCH] Keep chat-editor integration and shutdown checks with the runtime review --- tests/test_photo_editor_family.py | 503 ++++++++++++++++++++++++++++++ tests/test_shutdown_signal.py | 51 +++ 2 files changed, 554 insertions(+) create mode 100644 tests/test_photo_editor_family.py create mode 100644 tests/test_shutdown_signal.py diff --git a/tests/test_photo_editor_family.py b/tests/test_photo_editor_family.py new file mode 100644 index 000000000..b29111360 --- /dev/null +++ b/tests/test_photo_editor_family.py @@ -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"])) diff --git a/tests/test_shutdown_signal.py b/tests/test_shutdown_signal.py new file mode 100644 index 000000000..2858e9e01 --- /dev/null +++ b/tests/test_shutdown_signal.py @@ -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