odysseus/tests/test_photo_editor_family.py

"""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"]))