113 lines
4.6 KiB
Python
113 lines
4.6 KiB
Python
import json
|
|
|
|
from app.agent import AgentStep
|
|
from app.agent_runs import AgentRunRepository, AgentService, Origin, RunRecorder, step_record, use_origin
|
|
from tests.support import Env
|
|
|
|
SECRET = "super-secret-value"
|
|
|
|
|
|
def step(result="ok", args=None, error=None):
|
|
return AgentStep(1, "get_service_status", args or {"host": "web01", "service": "CPU"}, result, 0.123, error)
|
|
|
|
|
|
def test_step_record_redacts_and_truncates():
|
|
record = step_record(step(result=f"x {SECRET} " + "y" * 2000, args={"host": SECRET}), [SECRET])
|
|
assert SECRET not in json.dumps(record) and len(record["result"]) <= 500
|
|
assert record["duration"] == 0.12 and record["error"] is None
|
|
|
|
|
|
async def test_recorder_writes_steps_run_and_audit_without_secrets():
|
|
env = Env()
|
|
repo = AgentRunRepository(env.db)
|
|
recorder = RunRecorder(env.db, repo, Origin("telegram", "111"), "ask", f"frage {SECRET}", lambda: [SECRET])
|
|
await recorder(step(result=f"treffer {SECRET}", args={"host": "web01"}))
|
|
await recorder(step(error="ToolArgumentError"))
|
|
run_id = await recorder.finish(f"antwort {SECRET}", "step_limit")
|
|
run = repo.get(run_id)
|
|
assert (run["backend"], run["sender"], run["command"], run["abort_reason"]) == ("telegram", "111", "ask", "step_limit")
|
|
assert len(run["steps"]) == 2 and run["steps"][0]["tool"] == "get_service_status"
|
|
assert run["steps"][1]["error"] == "ToolArgumentError"
|
|
dump = json.dumps(run) + json.dumps([dict(r) for r in env.db.query("SELECT * FROM audit")])
|
|
assert SECRET not in dump
|
|
tool_audits = env.db.query("SELECT user, target, new_value FROM audit WHERE action='agent.tool'")
|
|
assert len(tool_audits) == 2 and tool_audits[0]["user"] == "telegram:111"
|
|
assert tool_audits[0]["target"] == "get_service_status" and "web01" in tool_audits[0]["new_value"]
|
|
|
|
|
|
class ScriptedChat:
|
|
def __init__(self):
|
|
self.replies = [
|
|
{"content": "", "tool_calls": [{"function": {"name": "get_problems", "arguments": {}}}]},
|
|
{"content": "Alles gut", "tool_calls": []},
|
|
]
|
|
self.closed = False
|
|
|
|
async def chat(self, model, messages, tools=None, options=None):
|
|
return self.replies.pop(0)
|
|
|
|
async def aclose(self):
|
|
self.closed = True
|
|
|
|
|
|
class FakeReader:
|
|
closed = False
|
|
|
|
async def problems(self):
|
|
return []
|
|
|
|
async def aclose(self):
|
|
self.closed = True
|
|
|
|
|
|
async def test_agent_service_runs_records_and_closes_clients():
|
|
env = Env()
|
|
chat, reader = ScriptedChat(), FakeReader()
|
|
repo = AgentRunRepository(env.db)
|
|
service = AgentService(env.store, env.db, repo, lambda h, s: [], lambda: [], lambda s: chat, lambda s: reader)
|
|
with use_origin("matrix", "@alice:example.org"):
|
|
result = await service.ask("Was ist los?")
|
|
assert result.text == "Alles gut" and chat.closed and reader.closed
|
|
run = repo.recent()[0]
|
|
assert (run["backend"], run["sender"], run["command"], run["question"]) == (
|
|
"matrix", "@alice:example.org", "ask", "Was ist los?")
|
|
assert [s["tool"] for s in run["steps"]] == ["get_problems"]
|
|
assert env.db.query("SELECT COUNT(*) FROM audit WHERE action='agent.tool'")[0][0] == 1
|
|
|
|
|
|
async def test_agent_service_origin_command_override_for_deep_analysis():
|
|
env = Env()
|
|
chat = ScriptedChat()
|
|
chat.replies = [{"content": "ok", "tool_calls": []}]
|
|
repo = AgentRunRepository(env.db)
|
|
service = AgentService(env.store, env.db, repo, lambda h, s: [], lambda: [], lambda s: chat, lambda s: FakeReader())
|
|
with use_origin("telegram", "111", "deep"):
|
|
await service.explain("web01", "CPU load")
|
|
run = repo.recent()[0]
|
|
assert run["command"] == "deep" and run["question"] == "web01 / CPU load"
|
|
|
|
|
|
async def test_agent_service_unconfigured_is_friendly_and_closes():
|
|
env = Env()
|
|
reader = FakeReader()
|
|
service = AgentService(env.store, env.db, AgentRunRepository(env.db), lambda h, s: [], lambda: [],
|
|
lambda s: None, lambda s: reader)
|
|
result = await service.ask("x")
|
|
assert "nicht verfügbar" in result.text and reader.closed
|
|
assert AgentRunRepository(env.db).recent() == []
|
|
|
|
|
|
async def test_agent_service_failure_does_not_leak_exception_text():
|
|
env = Env()
|
|
|
|
class Boom(ScriptedChat):
|
|
async def chat(self, *args, **kwargs):
|
|
raise RuntimeError(SECRET)
|
|
|
|
chat = Boom()
|
|
service = AgentService(env.store, env.db, AgentRunRepository(env.db), lambda h, s: [], lambda: [SECRET],
|
|
lambda s: chat, lambda s: FakeReader())
|
|
result = await service.ask("x")
|
|
assert SECRET not in result.text and chat.closed
|
|
assert SECRET not in json.dumps(AgentRunRepository(env.db).recent())
|