Spaces:
Sleeping
Sleeping
| 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 | |