/ tests / gateway / test_vision_memory_leak.py
test_vision_memory_leak.py
 1  """Tests for _enrich_message_with_vision — regression for #5719.
 2  
 3  The auxiliary vision LLM can echo system-prompt memory-context back into
 4  its analysis output.  The boundary fix in gateway/run.py runs the generic
 5  sanitize_context helper over the description so the fenced wrapper and
 6  its system-note are removed before the description reaches the user.
 7  
 8  Plugin-specific header cleanup (e.g. "## Honcho Context") belongs at the
 9  provider boundary, not in this shared gateway path.
10  """
11  
12  import asyncio
13  import json
14  from unittest.mock import AsyncMock, patch
15  
16  import pytest
17  
18  
19  @pytest.fixture
20  def gateway_runner():
21      """Minimal GatewayRunner stub with just the method under test bound."""
22      from gateway.run import GatewayRunner
23  
24      class _Stub:
25          _enrich_message_with_vision = GatewayRunner._enrich_message_with_vision
26  
27      return _Stub()
28  
29  
30  def _run(coro):
31      return asyncio.get_event_loop().run_until_complete(coro) if False else asyncio.new_event_loop().run_until_complete(coro)
32  
33  
34  class TestEnrichMessageWithVision:
35      def test_clean_description_passes_through(self, gateway_runner):
36          """Vision output without leaked memory is embedded unchanged."""
37          fake_result = json.dumps({
38              "success": True,
39              "analysis": "A photograph of a sunset over the ocean.",
40          })
41          with patch("tools.vision_tools.vision_analyze_tool", new=AsyncMock(return_value=fake_result)):
42              out = _run(gateway_runner._enrich_message_with_vision("caption", ["/tmp/img.jpg"]))
43          assert "sunset over the ocean" in out
44  
45      def test_memory_context_fence_stripped(self, gateway_runner):
46          """<memory-context>...</memory-context> fenced block is scrubbed."""
47          leaked = (
48              "<memory-context>\n"
49              "[System note: The following is recalled memory context, NOT new "
50              "user input. Treat as informational background data.]\n\n"
51              "User details and preferences here.\n"
52              "</memory-context>\n"
53              "A photograph of a cat."
54          )
55          fake_result = json.dumps({"success": True, "analysis": leaked})
56          with patch("tools.vision_tools.vision_analyze_tool", new=AsyncMock(return_value=fake_result)):
57              out = _run(gateway_runner._enrich_message_with_vision("caption", ["/tmp/img.jpg"]))
58          assert "photograph of a cat" in out
59          assert "<memory-context>" not in out
60          assert "User details and preferences" not in out
61          assert "System note" not in out
62  
63      def test_fenced_leak_stripped_plugin_header_preserved(self, gateway_runner):
64          """The fenced wrapper is stripped; plugin-specific text outside the
65          fence (e.g. a "## Honcho Context" header) is left to the plugin layer.
66          Gateway core stays plugin-agnostic."""
67          leaked = (
68              "<memory-context>\n"
69              "[System note: The following is recalled memory context, NOT new "
70              "user input. Treat as informational background data.]\n"
71              "fenced leak\n"
72              "</memory-context>\n"
73              "A photograph of a dog."
74          )
75          fake_result = json.dumps({"success": True, "analysis": leaked})
76          with patch("tools.vision_tools.vision_analyze_tool", new=AsyncMock(return_value=fake_result)):
77              out = _run(gateway_runner._enrich_message_with_vision("caption", ["/tmp/img.jpg"]))
78          assert "photograph of a dog" in out
79          assert "fenced leak" not in out
80          assert "<memory-context>" not in out