tests/test_shell_visual_engine.py

import re

from app.agents.shell_visual_engine import ShellGrounding, ShellSetupCode, ShellVisualEngine


class _FakeClient:
    def __init__(self, obj):
        self._obj = obj

    def structured_complete(self, messages, schema, **kw):
        return self._obj


def test_extract_grounding_3d_happy_path_keeps_anchor():
    chunks = [{"source": "paper.pdf", "text": "The hemoglobin molecule forms a tetramer of four globular subunits."}]
    grounding = ShellGrounding(
        renderable=True,
        anchor="forms a tetramer of four globular subunits",
        scene_brief="Render hemoglobin as four globular subunits arranged as a tetramer.",
    )
    engine = ShellVisualEngine(client=_FakeClient(grounding))
    result = engine._extract_grounding("hemoglobin", "3d", chunks, "high_school")

    assert result.renderable is True
    assert result.anchor == "forms a tetramer of four globular subunits"
    assert result.scene_brief == "Render hemoglobin as four globular subunits arranged as a tetramer."


def test_extract_grounding_2d_anim_happy_path_keeps_anchor():
    chunks = [{"source": "paper.pdf", "text": "During depolarization, sodium ions rush into the cell, flipping the membrane potential."}]
    grounding = ShellGrounding(
        renderable=True,
        anchor="sodium ions rush into the cell, flipping the membrane potential",
        scene_brief="Animate sodium ions flowing into the cell as the membrane potential flips positive.",
    )
    engine = ShellVisualEngine(client=_FakeClient(grounding))
    result = engine._extract_grounding("depolarization", "2d_anim", chunks, "high_school")

    assert result.renderable is True
    assert result.anchor == "sodium ions rush into the cell, flipping the membrane potential"
    assert result.scene_brief == "Animate sodium ions flowing into the cell as the membrane potential flips positive."


def test_extract_grounding_decline_path_returned_as_is():
    chunks = [{"source": "paper.pdf", "text": "This section discusses abstract policy tradeoffs with no physical form."}]
    grounding = ShellGrounding(
        renderable=False,
        decline_reason="'policy tradeoffs' has no concrete spatial object in the source — better explored in chat.",
    )
    engine = ShellVisualEngine(client=_FakeClient(grounding))
    result = engine._extract_grounding("policy tradeoffs", "3d", chunks, "high_school")

    assert result.renderable is False
    assert result.anchor == ""
    assert result.decline_reason == "'policy tradeoffs' has no concrete spatial object in the source — better explored in chat."


def test_extract_grounding_discards_hallucinated_anchor():
    chunks = [{"source": "paper.pdf", "text": "The heart has four chambers that pump blood through the body."}]
    grounding = ShellGrounding(
        renderable=True,
        anchor="this text was never in the source at all",
        scene_brief="Render a four-chambered heart pumping blood.",
    )
    engine = ShellVisualEngine(client=_FakeClient(grounding))
    result = engine._extract_grounding("heart anatomy", "3d", chunks, "high_school")

    assert result.renderable is True
    assert result.anchor == ""
    assert result.scene_brief == "Render a four-chambered heart pumping blood."


def _engine():
    return ShellVisualEngine(client=_FakeClient(None))


def test_assemble_shell_replaces_setup_code_token_for_both_kinds():
    engine = _engine()
    for kind in ("3d", "2d_anim"):
        html = engine._assemble_shell(kind, "group.add(new THREE.Mesh());" + "x" * 30)
        assert "{SETUP_CODE_JSON}" not in html


def test_assemble_shell_escapes_hostile_script_close_tag():
    hostile = "group.add(new THREE.Mesh()); /* </script><script>alert(1)</script> */"
    engine = _engine()
    for kind in ("3d", "2d_anim"):
        html = engine._assemble_shell(kind, hostile)
        # The raw closing tag from the hostile string must never survive unescaped.
        assert "<script>alert(1)</script>" not in html
        # It must appear only in its escaped form -- both occurrences from the
        # hostile string end up as the literal 3-char sequence <\/script>.
        assert html.count("<\\/script>") == 2


def test_assemble_shell_globals_lockdown_params_present_for_both_kinds():
    shadowed_names = ["window", "document", "fetch", "eval", "parent", "top", "self", "Function"]
    engine = _engine()
    for kind in ("3d", "2d_anim"):
        html = engine._assemble_shell(kind, "y" * 40)
        for name in shadowed_names:
            assert f'"{name}"' in html


def test_assemble_shell_globals_lockdown_param_and_arg_counts_match_for_both_kinds():
    """Regression guard for the parity the module docstring calls out: the
    `new Function("api", <shadowed names...>, SETUP_CODE_JSON)` parameter
    list and the real invocation `contentFn(api, undefined, undefined, ...)`
    must pass exactly the same number of positional arguments -- one real
    value (`api`) plus exactly one `undefined` per shadowed name. If a future
    edit added or removed a shadowed name in only one of the two spots, it
    would either leave a dangerous global un-shadowed (security gap) or throw
    a runtime arity mismatch -- nothing else in this file catches that today.
    """
    engine = _engine()
    for kind in ("3d", "2d_anim"):
        html = engine._assemble_shell(kind, "v" * 40)

        function_call_match = re.search(r"new Function\(([^)]*)\)", html)
        assert function_call_match, f"could not find a `new Function(...)` call for kind={kind}"
        # Every argument to `new Function` here is a quoted string literal --
        # the shadowed param names, plus the setup-code body (itself a JSON
        # string literal) as the trailing argument. Extracting each quoted
        # token this way (rather than a naive comma split) stays correct even
        # though the setup-code string could itself contain commas.
        function_args = re.findall(r'"(?:[^"\\]|\\.)*"', function_call_match.group(1))
        # The last quoted argument is the function body (the setup code), not
        # a shadowed name -- everything before it is a parameter name (api +
        # the shadowed globals).
        param_name_count = len(function_args) - 1

        invocation_matches = re.findall(r"contentFn\(([^)]*)\)", html)
        assert len(invocation_matches) == 1, f"expected exactly one contentFn(...) call site for kind={kind}"
        invocation_args = [arg.strip() for arg in invocation_matches[0].split(",")]
        arg_count = len(invocation_args)

        assert param_name_count == arg_count, (
            f"kind={kind}: new Function declares {param_name_count} parameter names "
            f"but contentFn(...) passes {arg_count} arguments -- these must match "
            "exactly (api's real value + one undefined per shadowed global)"
        )


def test_assemble_shell_3d_has_pinned_threejs_cdn_2d_does_not():
    engine = _engine()
    html_3d = engine._assemble_shell("3d", "z" * 40)
    html_2d = engine._assemble_shell("2d_anim", "z" * 40)

    assert "https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js" in html_3d
    assert "https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/examples/js/controls/OrbitControls.js" in html_3d
    assert "three.js/r128" not in html_2d


def test_assemble_shell_no_dark_background_uses_cream():
    engine = _engine()
    for kind in ("3d", "2d_anim"):
        html = engine._assemble_shell(kind, "w" * 40)
        assert "#0f0f0f" not in html
        assert "#FAF7F2" in html


def test_assemble_shell_has_well_formed_script_block():
    engine = _engine()
    for kind in ("3d", "2d_anim"):
        html = engine._assemble_shell(kind, "group.add(new THREE.Mesh()); return { update: function(t) {} };")
        blocks = re.findall(r"<script[^>]*>(.*?)</script>", html, re.DOTALL)
        assert len(blocks) >= 1


class _StagedFakeClient:
    """Fakes the two internal `structured_complete` calls `generate()` makes (grounding,
    then setup-code) with a single client, dispatching on the schema class passed as the
    second argument -- and records every call's (schema, messages) for later inspection.
    If `setup_code_response` is left None, a setup-code-stage call is treated as a test
    failure (used to prove the decline path never reaches that stage)."""

    def __init__(self, grounding_response, setup_code_response=None):
        self._grounding_response = grounding_response
        self._setup_code_response = setup_code_response
        self.calls = []  # list of (schema, messages)

    def structured_complete(self, messages, schema, **kw):
        self.calls.append((schema, messages))
        if schema is ShellGrounding:
            return self._grounding_response
        if schema is ShellSetupCode:
            if self._setup_code_response is None:
                raise AssertionError("setup-code stage should never have been reached on the decline path")
            return self._setup_code_response
        raise AssertionError(f"unexpected schema passed to structured_complete: {schema}")


def _setup_code_messages(client):
    calls = [messages for schema, messages in client.calls if schema is ShellSetupCode]
    assert len(calls) == 1, "expected exactly one ShellSetupCode structured_complete call"
    return calls[0]


def test_generate_3d_happy_path_returns_assembled_html():
    grounding = ShellGrounding(renderable=True, scene_brief="Render a tetrahedron with four triangular faces.")
    setup = ShellSetupCode(
        setup_code="const { THREE, scene, camera, group } = api; group.add(new THREE.Mesh());"
    )
    client = _StagedFakeClient(grounding, setup)
    engine = ShellVisualEngine(client=client)
    chunks = [{"source": "paper.pdf", "text": "A tetrahedron has four triangular faces."}]

    result = engine.generate("tetrahedron", "3d", chunks, "high_school")

    assert result.animation_type == "3d"
    assert result.html_code
    assert "<script>" in result.html_code
    assert result.explanation == grounding.scene_brief


def test_generate_2d_anim_happy_path_returns_assembled_html():
    grounding = ShellGrounding(renderable=True, scene_brief="Animate a ball bouncing under gravity.")
    setup = ShellSetupCode(
        setup_code="const { ctx, width, height } = api; return { draw(ctx, width, height, time, dt) {} };"
    )
    client = _StagedFakeClient(grounding, setup)
    engine = ShellVisualEngine(client=client)
    chunks = [{"source": "paper.pdf", "text": "A ball bounces under gravity, losing height each bounce."}]

    result = engine.generate("bouncing ball", "2d_anim", chunks, "high_school")

    assert result.animation_type == "2d_anim"
    assert result.html_code
    assert "<script>" in result.html_code
    assert result.explanation == grounding.scene_brief


def test_generate_decline_path_short_circuits_before_setup_code():
    grounding = ShellGrounding(
        renderable=False,
        decline_reason="'policy tradeoffs' has no concrete spatial object in the source — better explored in chat.",
    )
    # No setup_code_response given -- any structured_complete(..., ShellSetupCode) call raises.
    client = _StagedFakeClient(grounding, setup_code_response=None)
    engine = ShellVisualEngine(client=client)
    chunks = [{"source": "paper.pdf", "text": "This section discusses abstract policy tradeoffs."}]

    result = engine.generate("policy tradeoffs", "3d", chunks, "high_school")

    assert result.animation_type == "2d_text"
    assert "policy tradeoffs" in result.html_code
    assert result.explanation == grounding.decline_reason
    # Exactly one structured_complete call happened, and it was the grounding stage --
    # the setup-code stage was genuinely never reached.
    assert len(client.calls) == 1
    assert client.calls[0][0] is ShellGrounding


def test_generate_setup_code_prompt_instructs_group_not_scene_for_3d():
    grounding = ShellGrounding(renderable=True, scene_brief="Render a tetrahedron with four triangular faces.")
    setup = ShellSetupCode(setup_code="const { THREE, scene, camera, group } = api; group.add(new THREE.Mesh());")
    client = _StagedFakeClient(grounding, setup)
    engine = ShellVisualEngine(client=client)
    chunks = [{"source": "paper.pdf", "text": "A tetrahedron has four triangular faces."}]

    engine.generate("tetrahedron", "3d", chunks, "high_school")

    prompt_text = " ".join(m["content"] for m in _setup_code_messages(client))
    assert "Do NOT call `scene.add(...)` directly" in prompt_text
    assert "add every mesh/line/points/group you create to `group`" in prompt_text


def test_generate_setup_code_prompt_never_sees_raw_source_text():
    marker = "UNIQUE_SOURCE_MARKER_TEXT_12345"
    grounding = ShellGrounding(renderable=True, scene_brief="Render a tetrahedron with four triangular faces.")
    setup = ShellSetupCode(setup_code="const { THREE, scene, camera, group } = api; group.add(new THREE.Mesh());")
    client = _StagedFakeClient(grounding, setup)
    engine = ShellVisualEngine(client=client)
    chunks = [{"source": "paper.pdf", "text": f"A tetrahedron has four triangular faces. {marker}"}]

    engine.generate("tetrahedron", "3d", chunks, "high_school")

    prompt_text = " ".join(m["content"] for m in _setup_code_messages(client))
    assert marker not in prompt_text
    assert grounding.scene_brief in prompt_text