fix: avoid swallowing prompt manager interrupts

This commit is contained in:
Jacob Molz 2026-05-26 10:37:21 -04:00
parent f8e8c0e9e6
commit 8a8e496acf
2 changed files with 43 additions and 10 deletions

View file

@ -7,7 +7,7 @@ including template rendering, term merging, JSON validation, and error handling.
import pytest import pytest
import json import json
from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import AsyncMock
from trustgraph.template.prompt_manager import PromptManager, PromptConfiguration, Prompt from trustgraph.template.prompt_manager import PromptManager, PromptConfiguration, Prompt
@ -344,6 +344,42 @@ class TestPromptManager:
assert pm.terms == {} # Default empty terms assert pm.terms == {} # Default empty terms
assert len(pm.prompts) == 0 assert len(pm.prompts) == 0
def test_load_config_does_not_swallow_keyboard_interrupt(self, monkeypatch):
"""KeyboardInterrupt should propagate out of config parsing."""
pm = PromptManager()
def interrupt(_value):
raise KeyboardInterrupt
monkeypatch.setattr("trustgraph.template.prompt_manager.json.loads", interrupt)
with pytest.raises(KeyboardInterrupt):
pm.load_config({"system": json.dumps("Test")})
@pytest.mark.asyncio
async def test_json_parse_does_not_swallow_system_exit(self):
"""SystemExit should propagate out of JSON response parsing."""
pm = PromptManager()
config = {
"system": json.dumps("Test"),
"template-index": json.dumps(["json_response"]),
"template.json_response": json.dumps({
"prompt": "Generate JSON",
"response-type": "json"
})
}
pm.load_config(config)
def exit_parse(_text):
raise SystemExit(2)
pm.parse_json = exit_parse
mock_llm = AsyncMock()
mock_llm.return_value = "{}"
with pytest.raises(SystemExit):
await pm.invoke("json_response", {}, mock_llm)
@pytest.mark.unit @pytest.mark.unit
class TestPromptManagerJsonl: class TestPromptManagerJsonl:
@ -585,4 +621,4 @@ not json at all
assert len(result) == 2 assert len(result) == 2
assert result[0] == {"any": "structure"} assert result[0] == {"any": "structure"}
assert result[1] == {"completely": "different"} assert result[1] == {"completely": "different"}

View file

@ -31,12 +31,12 @@ class PromptManager:
try: try:
system = json.loads(config["system"]) system = json.loads(config["system"])
except: except (KeyError, TypeError, json.JSONDecodeError):
system = "Be helpful." system = "Be helpful."
try: try:
ix = json.loads(config["template-index"]) ix = json.loads(config["template-index"])
except: except (KeyError, TypeError, json.JSONDecodeError):
ix = [] ix = []
prompts = {} prompts = {}
@ -68,8 +68,8 @@ class PromptManager:
try: try:
self.system_template = ibis.Template(self.config.system_template) self.system_template = ibis.Template(self.config.system_template)
except: except Exception as e:
raise RuntimeError("Error in system template") raise RuntimeError(f"Error in system template: {e}")
self.templates = {} self.templates = {}
for k, v in self.prompts.items(): for k, v in self.prompts.items():
@ -136,8 +136,6 @@ class PromptManager:
terms = self.terms | self.prompts[id].terms | input terms = self.terms | self.prompts[id].terms | input
resp_type = self.prompts[id].response_type
return self.templates[id].render(terms) return self.templates[id].render(terms)
async def invoke(self, id, input, llm): async def invoke(self, id, input, llm):
@ -161,7 +159,7 @@ class PromptManager:
if resp_type == "json": if resp_type == "json":
try: try:
obj = self.parse_json(resp) obj = self.parse_json(resp)
except: except (json.JSONDecodeError, TypeError):
logger.error(f"JSON parse failed: {resp}") logger.error(f"JSON parse failed: {resp}")
raise RuntimeError("JSON parse fail") raise RuntimeError("JSON parse fail")
@ -195,4 +193,3 @@ class PromptManager:
return objects return objects
raise RuntimeError(f"Response type {resp_type} not known") raise RuntimeError(f"Response type {resp_type} not known")