6278 lines
290 KiB
Python
6278 lines
290 KiB
Python
# retoor <retoor@molodetz.nl>
|
||
import contextlib
|
||
import io
|
||
import json
|
||
import os
|
||
import re
|
||
import shutil
|
||
import sqlite3
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
import urllib.error
|
||
from datetime import datetime, timedelta, timezone
|
||
from unittest import mock
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
import tai
|
||
|
||
|
||
VALID_SKILL = """---
|
||
name: pdf-forms
|
||
description: >
|
||
Fill PDF forms and extract field data.
|
||
Use when the user mentions PDF documents.
|
||
---
|
||
# PDF forms
|
||
Do the thing.
|
||
"""
|
||
|
||
|
||
class SkillTests(unittest.TestCase):
|
||
def write_skill(self, root, entry, body):
|
||
folder = os.path.join(root, entry)
|
||
os.makedirs(folder, exist_ok=True)
|
||
path = os.path.join(folder, "SKILL.md")
|
||
with open(path, "w", encoding="utf-8") as handle:
|
||
handle.write(body)
|
||
return path
|
||
|
||
def test_parse_valid(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
path = self.write_skill(tmp, "x", VALID_SKILL)
|
||
skill = tai.parse_skill_file(path)
|
||
self.assertEqual(skill["name"], "pdf-forms")
|
||
self.assertIn("Fill PDF forms", skill["description"])
|
||
self.assertIn("Do the thing.", skill["body"])
|
||
|
||
def test_parse_rejects_bad(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
missing = self.write_skill(tmp, "a", "---\nname: x\n---\nbody\n")
|
||
self.assertIsNone(tai.parse_skill_file(missing))
|
||
bad_name = self.write_skill(tmp, "b", "---\nname: Bad_Name!\ndescription: d\n---\nbody\n")
|
||
self.assertIsNone(tai.parse_skill_file(bad_name))
|
||
no_front = self.write_skill(tmp, "c", "just markdown\n")
|
||
self.assertIsNone(tai.parse_skill_file(no_front))
|
||
|
||
def test_discover_project_wins(self):
|
||
with tempfile.TemporaryDirectory() as home, tempfile.TemporaryDirectory() as project:
|
||
self.write_skill(os.path.join(home, "skills"), "dup", VALID_SKILL)
|
||
other = VALID_SKILL.replace("Do the thing.", "Project variant.")
|
||
self.write_skill(os.path.join(project, ".tai", "skills"), "dup", other)
|
||
found = tai.discover_skills(home, project)
|
||
self.assertEqual(list(found), ["pdf-forms"])
|
||
self.assertIn("Project variant.", found["pdf-forms"]["body"])
|
||
|
||
def test_catalog(self):
|
||
skills = {"b-skill": {"description": "second"}, "a-skill": {"description": "first"}}
|
||
catalog = tai.skill_catalog(skills)
|
||
self.assertLess(catalog.index("a-skill"), catalog.index("b-skill"))
|
||
self.assertIn("load_skill", catalog)
|
||
empty = tai.skill_catalog({})
|
||
self.assertIn("Buildable skill blueprints", empty)
|
||
self.assertIn("bot-creator", empty)
|
||
|
||
|
||
class FakeApp:
|
||
def __init__(self):
|
||
self.skills = {}
|
||
self.env = "home"
|
||
self.store = None
|
||
|
||
|
||
class ToolTests(unittest.TestCase):
|
||
def test_load_skill_unknown(self):
|
||
tools = tai.Tools(FakeApp())
|
||
result = tools.dispatch("load_skill", json.dumps({"name": "nope"}))
|
||
self.assertIn("unknown skill", result)
|
||
|
||
def test_terminal_without_tmux(self):
|
||
tools = tai.Tools(FakeApp())
|
||
with tempfile.TemporaryDirectory() as empty:
|
||
with mock.patch.dict(os.environ, {"PATH": empty}):
|
||
self.assertEqual(tools.dispatch("get_current_terminal_content", "{}"), "tmux not available")
|
||
|
||
def test_box_helpers_present(self):
|
||
self.assertIn("faster-whisper", tai.BOX_CONTAINERFILE)
|
||
self.assertIn("edge-tts", tai.BOX_CONTAINERFILE)
|
||
self.assertIn("sleep", tai.BOX_CONTAINERFILE)
|
||
self.assertIn("WhisperModel", tai.BOX_STT)
|
||
self.assertIn("Communicate", tai.BOX_TTS)
|
||
|
||
|
||
class TerminalContextTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.hist = os.path.join(self.tmp.name, "hist")
|
||
with open(self.hist, "w") as handle:
|
||
handle.write("git status\ngit push\n")
|
||
self.paths = mock.patch.object(tai, "shell_history_paths", return_value=[self.hist])
|
||
self.paths.start()
|
||
self.which = mock.patch.object(tai.shutil, "which", return_value="/usr/bin/tmux")
|
||
self.which.start()
|
||
self.calls = []
|
||
|
||
def tearDown(self):
|
||
self.which.stop()
|
||
self.paths.stop()
|
||
self.tmp.cleanup()
|
||
|
||
def tmux(self, panes=(), captures=None, list_error=None):
|
||
captures = captures or {}
|
||
|
||
def side(argv, **kwargs):
|
||
self.calls.append(list(argv))
|
||
sub = argv[1] if len(argv) > 1 else ""
|
||
if sub == "display-message":
|
||
return mock.Mock(returncode=0, stdout="main:0.0 bash", stderr="")
|
||
if sub == "list-panes":
|
||
if list_error is not None:
|
||
return mock.Mock(returncode=1, stdout="", stderr=list_error)
|
||
return mock.Mock(returncode=0, stdout="".join(line + "\n" for line in panes), stderr="")
|
||
if sub == "capture-pane":
|
||
key = argv[argv.index("-t") + 1] if "-t" in argv else ""
|
||
if key in captures:
|
||
out = captures[key]
|
||
if isinstance(out, tuple):
|
||
return mock.Mock(returncode=out[0], stdout=out[1], stderr=out[2])
|
||
return mock.Mock(returncode=0, stdout=out, stderr="")
|
||
return mock.Mock(returncode=1, stdout="", stderr="can't find pane: %s" % key)
|
||
raise AssertionError("unexpected tmux call %r" % (argv,))
|
||
|
||
return mock.patch.object(tai.subprocess, "run", side_effect=side)
|
||
|
||
def pane(self, session, window, index, pid, active, command, path):
|
||
return "\t".join([session, str(window), str(index), pid, "1" if active else "0", command, path])
|
||
|
||
def test_history_section_bash_order(self):
|
||
section = tai.shell_history_section()
|
||
self.assertIn("## Recent shell history", section)
|
||
self.assertLess(section.index("git status"), section.index("git push"))
|
||
|
||
def test_history_section_missing_files(self):
|
||
with mock.patch.object(tai, "shell_history_paths", return_value=[os.path.join(self.tmp.name, "nope")]):
|
||
self.assertEqual(tai.shell_history_section(), "")
|
||
|
||
def test_history_zsh_extended_stripped(self):
|
||
with open(self.hist, "w") as handle:
|
||
handle.write(": 1690000000:0;vim foo\n: 1690000001:5;make\n")
|
||
section = tai.shell_history_section()
|
||
self.assertIn("vim foo", section)
|
||
self.assertNotIn("1690000000", section)
|
||
|
||
def test_history_fish_cmds(self):
|
||
fish = os.path.join(self.tmp.name, "fish_history")
|
||
with open(fish, "w") as handle:
|
||
handle.write("- cmd: ls -la\n when: 1690000000\n- cmd: pwd\n when: 1690000001\n")
|
||
with mock.patch.object(tai, "shell_history_paths", return_value=[fish]):
|
||
section = tai.shell_history_section()
|
||
self.assertIn("ls -la", section)
|
||
self.assertIn("pwd", section)
|
||
self.assertNotIn("when:", section)
|
||
|
||
def test_history_caps_keep_newest(self):
|
||
with open(self.hist, "w") as handle:
|
||
for num in range(200):
|
||
handle.write("cmd-%03d %s\n" % (num, "x" * 100))
|
||
section = tai.shell_history_section()
|
||
self.assertLessEqual(len(section), tai.SHELL_HISTORY_CHARS + 200)
|
||
self.assertIn("cmd-199", section)
|
||
self.assertNotIn("cmd-000", section)
|
||
|
||
def test_content_appends_history(self):
|
||
tools = tai.Tools(FakeApp())
|
||
with self.tmux(captures={"": "visible pane output"}):
|
||
result = tools.dispatch("get_current_terminal_content", "{}")
|
||
self.assertIn("visible pane output", result)
|
||
self.assertIn("## Recent shell history", result)
|
||
self.assertIn("git push", result)
|
||
|
||
def test_content_target_other_pane(self):
|
||
tools = tai.Tools(FakeApp())
|
||
with self.tmux(captures={"%5": "other pane output"}):
|
||
result = tools.dispatch("get_current_terminal_content", json.dumps({"target": "%5"}))
|
||
self.assertIn("other pane output", result)
|
||
self.assertIn("pane %5", result)
|
||
self.assertTrue(any("-t" in call and "%5" in call for call in self.calls))
|
||
|
||
def test_content_bad_target_keeps_history(self):
|
||
tools = tai.Tools(FakeApp())
|
||
with self.tmux(captures={}):
|
||
result = tools.dispatch("get_current_terminal_content", json.dumps({"target": "%9"}))
|
||
self.assertTrue(result.startswith("terminal capture failed"))
|
||
self.assertIn("git push", result)
|
||
|
||
def test_list_terminals_counts_and_marks_active(self):
|
||
tools = tai.Tools(FakeApp())
|
||
panes = (self.pane("main", 0, 0, "%1", True, "bash", "/home/u"),
|
||
self.pane("main", 0, 1, "%2", False, "vim", "/home/u"),
|
||
self.pane("other", 1, 0, "%3", False, "python", "/tmp"))
|
||
with self.tmux(panes=panes):
|
||
result = tools.dispatch("list_terminals", "{}")
|
||
self.assertIn("2 sessions", result)
|
||
self.assertIn("2 windows", result)
|
||
self.assertIn("3 panes", result)
|
||
self.assertIn("main:0.0 %1 bash /home/u *", result)
|
||
self.assertIn("other:1.0 %3 python /tmp", result)
|
||
|
||
def test_list_terminals_no_server(self):
|
||
tools = tai.Tools(FakeApp())
|
||
with self.tmux(list_error="no server running on /tmp/tmux-1000/default"):
|
||
self.assertEqual(tools.dispatch("list_terminals", "{}"), "no open terminals (tmux server not running)")
|
||
|
||
def test_tail_terminals_labels_each_pane(self):
|
||
tools = tai.Tools(FakeApp())
|
||
panes = (self.pane("main", 0, 0, "%1", True, "bash", "/home/u"),
|
||
self.pane("main", 0, 1, "%2", False, "vim", "/home/u"))
|
||
with self.tmux(panes=panes, captures={"%1": "one\ntwo\nthree", "%2": "alpha\nbeta"}):
|
||
result = tools.dispatch("tail_terminals", "{}")
|
||
self.assertIn("=== main:0.0 %1 (bash) ===", result)
|
||
self.assertIn("three", result)
|
||
self.assertIn("=== main:0.1 %2 (vim) ===", result)
|
||
self.assertIn("beta", result)
|
||
self.assertTrue(any("-3" in call for call in self.calls))
|
||
|
||
def test_tail_terminals_clamps_lines(self):
|
||
tools = tai.Tools(FakeApp())
|
||
panes = (self.pane("main", 0, 0, "%1", True, "bash", "/home/u"),)
|
||
with self.tmux(panes=panes, captures={"%1": "x"}):
|
||
tools.dispatch("tail_terminals", json.dumps({"lines": 99}))
|
||
self.assertTrue(any("-10" in call for call in self.calls))
|
||
|
||
def test_tail_terminals_dead_pane_skipped_loudly(self):
|
||
tools = tai.Tools(FakeApp())
|
||
panes = (self.pane("main", 0, 0, "%1", True, "bash", "/home/u"),
|
||
self.pane("main", 0, 1, "%2", False, "vim", "/home/u"))
|
||
with self.tmux(panes=panes, captures={"%1": "alive", "%2": (1, "", "can't find pane")}):
|
||
result = tools.dispatch("tail_terminals", "{}")
|
||
self.assertIn("alive", result)
|
||
self.assertIn("(empty)", result)
|
||
|
||
def test_tools_lazy_load_on_keywords(self):
|
||
names = {schema["function"]["name"] for schema in tai.select_tools("list all open terminal sessions")}
|
||
self.assertIn("list_terminals", names)
|
||
names = {schema["function"]["name"] for schema in tai.select_tools("tail every terminal pane")}
|
||
self.assertIn("tail_terminals", names)
|
||
names = {schema["function"]["name"] for schema in tai.select_tools("what is in my bash history")}
|
||
self.assertIn("get_current_terminal_content", names)
|
||
|
||
def test_schemas_registered(self):
|
||
schemas = {schema["function"]["name"]: schema["function"] for schema in tai.TOOL_SCHEMAS}
|
||
self.assertIn("target", schemas["get_current_terminal_content"]["parameters"]["properties"])
|
||
self.assertIn("list_terminals", schemas)
|
||
self.assertIn("lines", schemas["tail_terminals"]["parameters"]["properties"])
|
||
|
||
|
||
class AskUserTests(unittest.TestCase):
|
||
def tools(self, **attrs):
|
||
app = FakeApp()
|
||
for key, value in attrs.items():
|
||
setattr(app, key, value)
|
||
return tai.Tools(app)
|
||
|
||
def tty(self):
|
||
fake_stdin = io.StringIO("")
|
||
fake_stdin.isatty = lambda: True
|
||
return mock.patch.object(sys, "stdin", fake_stdin)
|
||
|
||
def test_schema_registered_and_core(self):
|
||
schemas = {schema["function"]["name"]: schema["function"] for schema in tai.TOOL_SCHEMAS}
|
||
self.assertIn("ask_user", schemas)
|
||
self.assertIn("question", schemas["ask_user"]["parameters"]["required"])
|
||
self.assertIn("ask_user", tai.CORE_TOOLS)
|
||
self.assertIn("question", tai.TOOL_TAGS["ask_user"])
|
||
self.assertIn("ask_user", {schema["function"]["name"] for schema in tai.select_tools("hello")})
|
||
|
||
def test_system_prompt_points_at_tool(self):
|
||
self.assertIn("ask_user", tai.DEFAULT_SYSTEM)
|
||
self.assertIn("ask_user", tai.AUTONOMOUS_NOTE)
|
||
|
||
def test_validation(self):
|
||
tools = self.tools()
|
||
self.assertIn("non-empty string", tools.dispatch("ask_user", json.dumps({})))
|
||
self.assertIn("non-empty string", tools.dispatch("ask_user", json.dumps({"question": " "})))
|
||
self.assertIn("must be a list", tools.dispatch("ask_user", json.dumps({"question": "q?", "options": "nope"})))
|
||
|
||
def test_free_input(self):
|
||
tools = self.tools()
|
||
with self.tty(), mock.patch("builtins.input", return_value="do the simple thing") as asked, mock.patch("builtins.print") as shown:
|
||
result = tools.dispatch("ask_user", json.dumps({"question": "Which approach?"}))
|
||
self.assertEqual(result, "user answered: do the simple thing")
|
||
self.assertTrue(tools.app.user_stepped_in)
|
||
asked.assert_called_once()
|
||
self.assertIn("free input", asked.call_args.args[0])
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list)
|
||
self.assertIn("Which approach?", printed)
|
||
|
||
def test_option_by_number(self):
|
||
tools = self.tools()
|
||
with self.tty(), mock.patch("builtins.input", return_value="2"), mock.patch("builtins.print") as shown:
|
||
result = tools.dispatch("ask_user", json.dumps({"question": "Pick one", "options": ["Ship it", "Staging first", "Abort"]}))
|
||
self.assertEqual(result, "user chose option 2: Staging first")
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list)
|
||
self.assertIn("1) Ship it", printed)
|
||
self.assertIn("3) Abort", printed)
|
||
|
||
def test_option_by_text_case_insensitive(self):
|
||
tools = self.tools()
|
||
with self.tty(), mock.patch("builtins.input", return_value="staging FIRST"), mock.patch("builtins.print"):
|
||
result = tools.dispatch("ask_user", json.dumps({"question": "Pick one", "options": ["Ship it", "Staging first"]}))
|
||
self.assertEqual(result, "user chose option 2: Staging first")
|
||
|
||
def test_out_of_range_number_is_free_text(self):
|
||
tools = self.tools()
|
||
with self.tty(), mock.patch("builtins.input", return_value="9"), mock.patch("builtins.print"):
|
||
result = tools.dispatch("ask_user", json.dumps({"question": "Pick one", "options": ["A", "B"]}))
|
||
self.assertEqual(result, "user answered: 9")
|
||
|
||
def test_empty_answer_keeps_turn_alive(self):
|
||
tools = self.tools()
|
||
with self.tty(), mock.patch("builtins.input", return_value=" "), mock.patch("builtins.print"):
|
||
result = tools.dispatch("ask_user", json.dumps({"question": "q?"}))
|
||
self.assertIn("no answer (empty)", result)
|
||
self.assertIn("proceed", result)
|
||
|
||
def test_eof_keeps_turn_alive(self):
|
||
tools = self.tools()
|
||
with self.tty(), mock.patch("builtins.input", side_effect=EOFError), mock.patch("builtins.print"):
|
||
result = tools.dispatch("ask_user", json.dumps({"question": "q?"}))
|
||
self.assertIn("no answer (input closed)", result)
|
||
|
||
def test_non_tty_refuses_without_blocking(self):
|
||
tools = self.tools()
|
||
fake_stdin = io.StringIO("")
|
||
fake_stdin.isatty = lambda: False
|
||
with mock.patch.object(sys, "stdin", fake_stdin), mock.patch("builtins.input") as asked:
|
||
result = tools.dispatch("ask_user", json.dumps({"question": "q?"}))
|
||
self.assertIn("no interactive user", result)
|
||
asked.assert_not_called()
|
||
|
||
def test_auto_mode_refuses(self):
|
||
tools = self.tools(auto=True)
|
||
with mock.patch("builtins.input") as asked:
|
||
result = tools.dispatch("ask_user", json.dumps({"question": "q?"}))
|
||
self.assertIn("autonomous mode", result)
|
||
asked.assert_not_called()
|
||
|
||
def test_background_worker_refuses(self):
|
||
tools = self.tools(persist=False)
|
||
with mock.patch("builtins.input") as asked:
|
||
result = tools.dispatch("ask_user", json.dumps({"question": "q?"}))
|
||
self.assertIn("background workers", result)
|
||
asked.assert_not_called()
|
||
deep = self.tools(persist=True, depth=1)
|
||
with mock.patch("builtins.input") as asked:
|
||
result = deep.dispatch("ask_user", json.dumps({"question": "q?"}))
|
||
self.assertIn("background workers", result)
|
||
asked.assert_not_called()
|
||
|
||
def test_options_capped(self):
|
||
tools = self.tools()
|
||
many = ["opt%d" % num for num in range(20)]
|
||
with self.tty(), mock.patch("builtins.input", return_value="1"), mock.patch("builtins.print") as shown:
|
||
result = tools.dispatch("ask_user", json.dumps({"question": "q?", "options": many}))
|
||
self.assertEqual(result, "user chose option 1: opt0")
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list)
|
||
self.assertIn("12) opt11", printed)
|
||
self.assertNotIn("13) opt12", printed)
|
||
|
||
|
||
class SysinfoTests(unittest.TestCase):
|
||
def test_collect_all_checks_visible(self):
|
||
report = tai.collect_sysinfo()
|
||
lines = report.splitlines()
|
||
self.assertEqual(len(lines), len(tai.SYSINFO_CHECKS))
|
||
for (name, _func), line in zip(tai.SYSINFO_CHECKS, lines):
|
||
self.assertTrue(line.startswith(name + ": "), line)
|
||
self.assertRegex(line, r"\(\d+ms\)$")
|
||
|
||
def test_subset(self):
|
||
report = tai.collect_sysinfo(["os", "root"])
|
||
lines = report.splitlines()
|
||
self.assertEqual(len(lines), 2)
|
||
self.assertTrue(lines[0].startswith("os: "))
|
||
self.assertTrue(lines[1].startswith("root: "))
|
||
|
||
def test_checks_run_in_parallel(self):
|
||
def slow():
|
||
time.sleep(0.3)
|
||
return "ok"
|
||
|
||
stubs = tuple(("slow%d" % pos, slow) for pos in range(4))
|
||
with mock.patch("tai.SYSINFO_CHECKS", stubs):
|
||
started = time.time()
|
||
report = tai.collect_sysinfo()
|
||
self.assertLess(time.time() - started, 1.0)
|
||
self.assertEqual(len(report.splitlines()), 4)
|
||
|
||
def test_timeout_surfaces(self):
|
||
def stuck():
|
||
time.sleep(5)
|
||
return "never"
|
||
|
||
with mock.patch("tai.SYSINFO_CHECKS", (("stuck", stuck),)):
|
||
with mock.patch.object(tai.Consts, "SYSINFO_TIMEOUT", 1):
|
||
report = tai.collect_sysinfo()
|
||
self.assertIn("stuck: timed out after 1s", report)
|
||
|
||
def test_check_error_surfaces(self):
|
||
def broken():
|
||
raise RuntimeError("boom")
|
||
|
||
with mock.patch("tai.SYSINFO_CHECKS", (("broken", broken),)):
|
||
report = tai.collect_sysinfo()
|
||
self.assertIn("broken: error: boom", report)
|
||
|
||
def test_venv_detection(self):
|
||
with mock.patch.dict(os.environ, {"VIRTUAL_ENV": "/tmp/fake-venv"}):
|
||
self.assertIn("/tmp/fake-venv", tai.sysinfo_check_venv())
|
||
|
||
def test_root_shape(self):
|
||
result = tai.sysinfo_check_root()
|
||
self.assertTrue(result.startswith(("yes", "no", "unknown")), result)
|
||
|
||
def test_tool_dispatch(self):
|
||
tools = tai.Tools(FakeApp())
|
||
self.assertIn("os: ", tools.dispatch("sysinfo", "{}"))
|
||
subset = tools.dispatch("sysinfo", json.dumps({"checks": ["python"]}))
|
||
self.assertEqual(len(subset.splitlines()), 1)
|
||
self.assertIn("unknown checks: nope", tools.dispatch("sysinfo", json.dumps({"checks": ["nope"]})))
|
||
self.assertIn("non-empty list", tools.dispatch("sysinfo", json.dumps({"checks": []})))
|
||
|
||
def test_sysinfo_command(self):
|
||
self.assertIn("sysinfo", tai.COMMANDS)
|
||
self.assertIn("/sysinfo", tai.HELP_TEXT)
|
||
with mock.patch.object(tai.SysInfo, "collect_sysinfo", return_value="fake-report") as collector:
|
||
with mock.patch("builtins.print") as printer:
|
||
self.assertTrue(tai.handle_command(mock.Mock(), "/sysinfo"))
|
||
collector.assert_called_once_with()
|
||
printer.assert_called_once_with("fake-report")
|
||
|
||
|
||
class FakeSkillApp:
|
||
def __init__(self, config):
|
||
self.depth = 0
|
||
self.profile = "t"
|
||
self.config = config
|
||
self.store = mock.Mock()
|
||
self.store.redact = lambda text, profile=None: text
|
||
self.skills = {}
|
||
self.applied = 0
|
||
self.runner_override = None
|
||
|
||
def apply_system(self):
|
||
self.applied += 1
|
||
|
||
|
||
class CreateSkillTests(unittest.TestCase):
|
||
def config_in(self, tmp):
|
||
old_home = os.environ.get("TAI_HOME")
|
||
os.environ["TAI_HOME"] = tmp
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
try:
|
||
return tai.Config(FakeArgs())
|
||
finally:
|
||
if old_home is None:
|
||
os.environ.pop("TAI_HOME", None)
|
||
else:
|
||
os.environ["TAI_HOME"] = old_home
|
||
|
||
def test_prompt_is_self_contained(self):
|
||
prompt = tai.create_skill_prompt("pdf-forms", "fill pdf forms", "/tmp/x/pdf-forms")
|
||
self.assertIn("pdf-forms", prompt)
|
||
self.assertIn("fill pdf forms", prompt)
|
||
self.assertIn("/tmp/x/pdf-forms", prompt)
|
||
self.assertIn("sysinfo", prompt)
|
||
self.assertIn("SKILL.md", prompt)
|
||
self.assertIn("name: pdf-forms", prompt)
|
||
self.assertIn("two independent sources", prompt)
|
||
|
||
def test_validation(self):
|
||
tools = tai.Tools(FakeSkillApp(None))
|
||
self.assertIn("invalid skill name", tools.dispatch("create_skill", json.dumps({"name": "Bad_Name!", "brief": "b"})))
|
||
self.assertIn("invalid skill name", tools.dispatch("create_skill", json.dumps({"brief": "b"})))
|
||
self.assertIn("empty brief", tools.dispatch("create_skill", json.dumps({"name": "ok-name"})))
|
||
self.assertIn("scope must be", tools.dispatch("create_skill", json.dumps({"name": "ok-name", "brief": "b", "scope": "moon"})))
|
||
deep = FakeSkillApp(None)
|
||
deep.depth = 2
|
||
self.assertIn("depth limit", tai.Tools(deep).dispatch("create_skill", json.dumps({"name": "ok-name", "brief": "b"})))
|
||
|
||
def test_creates_and_refreshes(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
app = FakeSkillApp(self.config_in(tmp))
|
||
seen = {}
|
||
|
||
def runner(task, profile, timeout):
|
||
seen["task"] = task
|
||
seen["profile"] = profile
|
||
seen["timeout"] = timeout
|
||
skill_dir = os.path.join(tmp, ".tai", "skills", "pdf-forms")
|
||
os.makedirs(skill_dir, exist_ok=True)
|
||
with open(os.path.join(skill_dir, "SKILL.md"), "w", encoding="utf-8") as handle:
|
||
handle.write(VALID_SKILL)
|
||
return "researched and wrote the skill"
|
||
|
||
app.runner_override = runner
|
||
with mock.patch("os.getcwd", return_value=tmp):
|
||
result = tai.Tools(app).dispatch("create_skill", json.dumps({"name": "pdf-forms", "brief": "fill pdf forms"}))
|
||
self.assertIn("created at", result)
|
||
self.assertIn("pdf-forms", app.skills)
|
||
self.assertEqual(app.applied, 1)
|
||
self.assertEqual(seen["profile"], "t")
|
||
self.assertIn("fill pdf forms", seen["task"])
|
||
self.assertIn("SKILL.md", seen["task"])
|
||
|
||
def test_replace_and_missing(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
config = self.config_in(tmp)
|
||
skill_dir = os.path.join(tmp, ".tai", "skills", "pdf-forms")
|
||
os.makedirs(skill_dir, exist_ok=True)
|
||
with open(os.path.join(skill_dir, "SKILL.md"), "w", encoding="utf-8") as handle:
|
||
handle.write(VALID_SKILL)
|
||
app = FakeSkillApp(config)
|
||
app.runner_override = lambda task, profile, timeout: "rewrote it"
|
||
with mock.patch("os.getcwd", return_value=tmp):
|
||
result = tai.Tools(app).dispatch("create_skill", json.dumps({"name": "pdf-forms", "brief": "b"}))
|
||
self.assertIn("replaced at", result)
|
||
app2 = FakeSkillApp(config)
|
||
app2.runner_override = lambda task, profile, timeout: "gave up"
|
||
with mock.patch("os.getcwd", return_value=tmp):
|
||
result2 = tai.Tools(app2).dispatch("create_skill", json.dumps({"name": "other-skill", "brief": "b"}))
|
||
self.assertIn("was not created (done)", result2)
|
||
self.assertNotIn("other-skill", app2.skills)
|
||
self.assertEqual(app2.applied, 0)
|
||
|
||
def test_home_scope(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
app = FakeSkillApp(self.config_in(tmp))
|
||
|
||
def runner(task, profile, timeout):
|
||
skill_dir = os.path.join(tmp, "skills", "home-skill")
|
||
os.makedirs(skill_dir, exist_ok=True)
|
||
with open(os.path.join(skill_dir, "SKILL.md"), "w", encoding="utf-8") as handle:
|
||
handle.write(VALID_SKILL.replace("pdf-forms", "home-skill"))
|
||
return "done"
|
||
|
||
app.runner_override = runner
|
||
with mock.patch("os.getcwd", return_value="/nonexistent-dir"):
|
||
result = tai.Tools(app).dispatch("create_skill", json.dumps({"name": "home-skill", "brief": "b", "scope": "home"}))
|
||
self.assertIn("created at", result)
|
||
self.assertIn("home-skill", app.skills)
|
||
|
||
|
||
class FakeSecretApp:
|
||
def __init__(self, store, approve=True):
|
||
self.store = store
|
||
self.env = "home"
|
||
self.config = mock.Mock(auto_approve=True)
|
||
self.approvals = []
|
||
self.approve = approve
|
||
|
||
def ask_approval(self, command):
|
||
self.approvals.append(command)
|
||
return self.approve
|
||
|
||
def show_diff(self):
|
||
pass
|
||
|
||
|
||
class SecretsTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(config, tai.Seal(config.home, "vault-test-1"))
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def test_store_list_delete_roundtrip(self):
|
||
app = FakeSecretApp(self.store)
|
||
tools = tai.Tools(app)
|
||
self.assertIn("stored secret 'api'", tools.dispatch("store_secret", json.dumps({"name": "api", "value": "token-abc-123"})))
|
||
self.assertEqual(self.store.load_secret("api"), "token-abc-123")
|
||
raw = self.store.db.execute("SELECT value FROM secrets WHERE name = 'api'").fetchone()[0]
|
||
self.assertTrue(raw.startswith("tai1$"))
|
||
self.assertNotIn("token-abc", raw)
|
||
listed = tools.dispatch("list_secrets", "{}")
|
||
self.assertIn("api", listed)
|
||
self.assertNotIn("token-abc", listed)
|
||
self.assertIn("deleted secret 'api'", tools.dispatch("delete_secret", json.dumps({"name": "api"})))
|
||
self.assertIsNone(self.store.load_secret("api"))
|
||
self.assertEqual(len(app.approvals), 1)
|
||
|
||
def test_validation(self):
|
||
tools = tai.Tools(FakeSecretApp(self.store))
|
||
self.assertIn("invalid secret name", tools.dispatch("store_secret", json.dumps({"name": "Bad Name!", "value": "x"})))
|
||
self.assertIn("empty value", tools.dispatch("store_secret", json.dumps({"name": "ok", "value": ""})))
|
||
self.assertIn("unknown secret", tools.dispatch("delete_secret", json.dumps({"name": "nope"})))
|
||
self.assertIn("unknown tool", tools.dispatch("get_secret", json.dumps({"name": "x"})))
|
||
|
||
def test_shell_blind_injection(self):
|
||
self.store.save_secret("demo", "injected-value-42")
|
||
app = FakeSecretApp(self.store)
|
||
tools = tai.Tools(app)
|
||
result = tools.dispatch("shell", json.dumps({"command": "echo $TAI_SECRET_DEMO", "secrets": ["demo"]}))
|
||
self.assertIn("[redacted:demo]", result)
|
||
self.assertNotIn("injected-value-42", result)
|
||
tools.dispatch("shell", json.dumps({"command": "echo $TAI_SECRET_DEMO", "secrets": ["demo"]}))
|
||
self.assertEqual(len(app.approvals), 1)
|
||
self.assertIn("unknown secret 'nope'", tools.dispatch("shell", json.dumps({"command": "echo hi", "secrets": ["nope"]})))
|
||
self.assertIn("must be a list", tools.dispatch("shell", json.dumps({"command": "echo hi", "secrets": "demo"})))
|
||
|
||
def test_shell_secret_denied(self):
|
||
self.store.save_secret("demo", "injected-value-42")
|
||
app = FakeSecretApp(self.store, approve=False)
|
||
result = tai.Tools(app).dispatch("shell", json.dumps({"command": "echo hi", "secrets": ["demo"]}))
|
||
self.assertEqual(result, "denied by user")
|
||
|
||
def test_web_fetch_auth(self):
|
||
self.store.save_secret("api", "fetch-token-7")
|
||
seen = {}
|
||
|
||
class FakeResponse:
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, *exc):
|
||
return False
|
||
|
||
def read(self, limit=0):
|
||
return b"<html><body>ok fetch-token-7 here</body></html>"
|
||
|
||
def fake_urlopen(request, timeout=30):
|
||
seen[request.full_url] = {key.lower(): value for key, value in request.headers.items()}
|
||
return FakeResponse()
|
||
|
||
app = FakeSecretApp(self.store)
|
||
tools = tai.Tools(app)
|
||
with mock.patch("urllib.request.urlopen", side_effect=fake_urlopen):
|
||
result = tools.dispatch("web_fetch", json.dumps({"url": "https://api.example.test/v1", "auth_secret": "api"}))
|
||
keyed = tools.dispatch("web_fetch", json.dumps({"url": "https://api.example.test/v2", "auth_secret": "api", "auth_header": "X-Api-Key", "auth_prefix": ""}))
|
||
first = seen["https://api.example.test/v1"]
|
||
second = seen["https://api.example.test/v2"]
|
||
self.assertEqual(first["authorization"], "Bearer fetch-token-7")
|
||
self.assertIn("tai/", first["user-agent"])
|
||
self.assertEqual(second["x-api-key"], "fetch-token-7")
|
||
self.assertIn("[redacted:api]", result)
|
||
self.assertNotIn("fetch-token-7", result)
|
||
self.assertIn("[redacted:api]", keyed)
|
||
self.assertEqual(len(app.approvals), 1)
|
||
self.assertIn("unknown secret", tools.dispatch("web_fetch", json.dumps({"url": "https://api.example.test/v1", "auth_secret": "nope"})))
|
||
self.assertIn("invalid auth header", tools.dispatch("web_fetch", json.dumps({"url": "https://api.example.test/v1", "auth_secret": "api", "auth_header": "Bad\nHeader"})))
|
||
|
||
def test_redact(self):
|
||
self.store.save_secret("long", "abcdefghij")
|
||
self.store.save_secret("short", "abc")
|
||
self.store.save_secret("sub", "cdef")
|
||
text = self.store.redact("see abcdefghij and abc here")
|
||
self.assertEqual(text, "see [redacted:long] and abc here")
|
||
|
||
def test_run_turn_scrubs_and_redacts(self):
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
config = tai.Config(FakeArgs())
|
||
agent = tai.Agent(config, self.store, persist=False, quiet=True)
|
||
call_reply = {"role": "assistant", "content": "storing now hunter2-leak", "reasoning": "", "tool_calls": [{"id": "c9", "name": "store_secret", "arguments": json.dumps({"name": "leak", "value": "hunter2-leak"})}], "backend": "x"}
|
||
final_reply = {"role": "assistant", "content": "done", "reasoning": "", "tool_calls": [], "backend": "x"}
|
||
with mock.patch.object(agent.chat, "complete", side_effect=[call_reply, final_reply]):
|
||
result = agent.run_turn("remember it", capture=True)
|
||
self.assertEqual(result, "done")
|
||
blob = json.dumps(agent.messages)
|
||
self.assertNotIn("hunter2-leak", blob)
|
||
self.assertIn("[redacted:leak]", blob)
|
||
self.assertEqual(self.store.load_secret("leak"), "hunter2-leak")
|
||
self.assertEqual(self.store.search_events("t", "hunter2"), [])
|
||
|
||
def test_secret_repl(self):
|
||
agent = mock.Mock()
|
||
agent.store = self.store
|
||
agent.profile = "t"
|
||
agent.ask_approval = lambda command, guidance=True: True
|
||
agent.tools = tai.Tools(agent)
|
||
with mock.patch("getpass.getpass", return_value="repl-value-1"):
|
||
with mock.patch("builtins.print") as printer:
|
||
self.assertTrue(tai.handle_command(agent, "/secret set wifi"))
|
||
self.assertEqual(self.store.load_secret("wifi"), "repl-value-1")
|
||
printed = "\n".join(str(call.args[0]) for call in printer.call_args_list if call.args)
|
||
self.assertIn("stored secret 'wifi'", printed)
|
||
self.assertNotIn("repl-value-1", printed)
|
||
with mock.patch("builtins.print") as printer:
|
||
tai.handle_command(agent, "/secret list")
|
||
tai.handle_command(agent, "/secret")
|
||
listed = "\n".join(str(call.args[0]) for call in printer.call_args_list if call.args)
|
||
self.assertIn("- wifi", listed)
|
||
self.assertNotIn("repl-value-1", listed)
|
||
with mock.patch("builtins.print") as printer:
|
||
tai.handle_command(agent, "/secret delete wifi")
|
||
tai.handle_command(agent, "/secret delete nope")
|
||
tai.handle_command(agent, "/secret frobnicate")
|
||
tai.handle_command(agent, "/secret set Bad Name")
|
||
removed = "\n".join(str(call.args[0]) for call in printer.call_args_list if call.args)
|
||
self.assertIn("deleted secret 'wifi'", removed)
|
||
self.assertIn("unknown secret", removed)
|
||
self.assertIn("use /secret set", removed)
|
||
self.assertIn("invalid secret name", removed)
|
||
self.assertIsNone(self.store.load_secret("wifi"))
|
||
|
||
def test_legacy_memory_migrates_to_vault(self):
|
||
self.assertIn("store_secret", tai.DEFAULT_SYSTEM)
|
||
self.assertNotIn("collect and keep", tai.DEFAULT_SYSTEM)
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
config = tai.Config(FakeArgs())
|
||
old_system = "You are tai. Memory: call remember whenever you learn durable facts, " + tai.LEGACY_PASSWORD_NOTE + "; call it with a forget instruction."
|
||
self.store.save_system("old", old_system)
|
||
agent = tai.Agent(config, self.store, persist=False, quiet=True)
|
||
agent.switch_profile("old", silent=True)
|
||
self.assertNotIn(tai.LEGACY_PASSWORD_NOTE, agent.system_message)
|
||
self.assertIn("store_secret", agent.system_message)
|
||
self.assertIn("store_secret", self.store.load_system("old"))
|
||
|
||
|
||
class FakeSchedulerApp:
|
||
def __init__(self, store, approve=True):
|
||
self.store = store
|
||
self.profile = "t"
|
||
self.approvals = []
|
||
self.approve = approve
|
||
|
||
def ask_approval(self, command):
|
||
self.approvals.append(command)
|
||
return self.approve
|
||
|
||
|
||
class SchedulerTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.seal = tai.Seal(self.config.home, "sched-test-1")
|
||
self.store = tai.Store(self.config, self.seal)
|
||
with tai.AGENTS_LOCK:
|
||
tai.AGENTS.clear()
|
||
tai.AGENTS_NEXT[0] = 1
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
with tai.AGENTS_LOCK:
|
||
tai.AGENTS.clear()
|
||
tai.AGENTS_NEXT[0] = 1
|
||
|
||
def test_parse_at_and_duration(self):
|
||
parsed = tai.parse_schedule_at("2026-10-08T09:00")
|
||
self.assertEqual(datetime.fromisoformat(parsed).utcoffset().total_seconds(), 0)
|
||
self.assertEqual(tai.local_display(parsed), "2026-10-08 09:00")
|
||
with self.assertRaises(ValueError):
|
||
tai.parse_schedule_at("not a date")
|
||
self.assertEqual(tai.parse_duration("10m"), 600)
|
||
self.assertEqual(tai.parse_duration("2h"), 7200)
|
||
self.assertEqual(tai.parse_duration("1d"), 86400)
|
||
self.assertEqual(tai.parse_duration("45"), 45)
|
||
self.assertIsNone(tai.parse_duration("nope"))
|
||
self.assertEqual(tai.format_delay(4000), "1h")
|
||
self.assertEqual(tai.format_delay(-5), "overdue")
|
||
|
||
def test_schedule_validation(self):
|
||
tools = tai.Tools(FakeSchedulerApp(self.store))
|
||
self.assertIn("empty prompt", tools.dispatch("schedule", json.dumps({"at": "2026-10-08T09:00"})))
|
||
self.assertIn("exactly one", tools.dispatch("schedule", json.dumps({"prompt": "x"})))
|
||
self.assertIn("exactly one", tools.dispatch("schedule", json.dumps({"prompt": "x", "at": "2026-10-08T09:00", "every": 60})))
|
||
self.assertIn("invalid datetime", tools.dispatch("schedule", json.dumps({"prompt": "x", "at": "soon"})))
|
||
self.assertIn("at least 60", tools.dispatch("schedule", json.dumps({"prompt": "x", "every": 5})))
|
||
self.assertIn("invalid profile", tools.dispatch("schedule", json.dumps({"prompt": "x", "every": 60, "profile": "bad name!"})))
|
||
denied = tai.Tools(FakeSchedulerApp(self.store, approve=False))
|
||
self.assertEqual(denied.dispatch("schedule", json.dumps({"prompt": "x", "every": 60})), "denied by user")
|
||
self.assertEqual(self.store.list_schedules(), [])
|
||
|
||
def test_schedule_roundtrip(self):
|
||
app = FakeSchedulerApp(self.store)
|
||
tools = tai.Tools(app)
|
||
once = tools.dispatch("schedule", json.dumps({"prompt": "water plants", "name": "plants", "at": "2026-10-08T09:00"}))
|
||
self.assertIn("scheduled #1", once)
|
||
rep = tools.dispatch("schedule", json.dumps({"prompt": "check mail", "every": 3600}))
|
||
self.assertIn("scheduled #2", rep)
|
||
raw = self.store.db.execute("SELECT prompt FROM schedules WHERE id = 1").fetchone()[0]
|
||
self.assertTrue(raw.startswith("tai1$"))
|
||
listed = tools.dispatch("schedules", "{}")
|
||
self.assertIn("#1", listed)
|
||
self.assertIn("plants", listed)
|
||
self.assertIn("every 1h", listed)
|
||
self.assertIn("water plants", listed)
|
||
self.assertIn("deleted schedule #1", tools.dispatch("unschedule", json.dumps({"id": 1})))
|
||
self.assertIn("no schedule #1", tools.dispatch("unschedule", json.dumps({"id": 1})))
|
||
self.assertEqual(len(app.approvals), 3)
|
||
|
||
def test_tick_fires_once_async(self):
|
||
self.store.add_schedule("job", "do the thing", "t", 0, "2020-01-01T00:00:00+00:00", 60)
|
||
calls = []
|
||
|
||
def stub(task, profile, timeout):
|
||
calls.append((task, profile, timeout))
|
||
return "stub-result"
|
||
|
||
self.assertEqual(tai.scheduler_tick(self.config, self.seal, runner=stub), 1)
|
||
self.assertEqual(calls[0][0], "do the thing")
|
||
self.assertEqual(self.store.db.execute("SELECT status FROM schedules WHERE id = 1").fetchone()[0], "done")
|
||
self.assertEqual(tai.scheduler_tick(self.config, self.seal, runner=stub), 0)
|
||
self.assertEqual(len(calls), 1)
|
||
deadline = time.time() + 5
|
||
row = ("", "")
|
||
while time.time() < deadline:
|
||
row = self.store.db.execute("SELECT last_status, last_result FROM schedules WHERE id = 1").fetchone()
|
||
if row[0]:
|
||
break
|
||
time.sleep(0.05)
|
||
self.assertEqual(row[0], "done")
|
||
self.assertIn("stub-result", row[1])
|
||
|
||
def test_tick_repeat_advances_without_backfill(self):
|
||
past = (datetime.now(timezone.utc) - timedelta(days=2)).isoformat()
|
||
self.store.add_schedule("hourly", "ping", "t", 3600, past, 60)
|
||
calls = []
|
||
self.assertEqual(tai.scheduler_tick(self.config, self.seal, runner=lambda task, profile, timeout: calls.append(task) or "ok"), 1)
|
||
self.assertEqual(len(calls), 1)
|
||
row = self.store.db.execute("SELECT status, next_run FROM schedules WHERE id = 1").fetchone()
|
||
self.assertEqual(row[0], "pending")
|
||
self.assertGreater(datetime.fromisoformat(row[1]), datetime.now(timezone.utc))
|
||
|
||
def test_tick_skips_claimed(self):
|
||
self.store.add_schedule("job", "do it", "t", 0, "2020-01-01T00:00:00+00:00", 60)
|
||
self.store.db.execute("UPDATE schedules SET status = 'done' WHERE id = 1")
|
||
self.store.db.commit()
|
||
calls = []
|
||
self.assertEqual(tai.scheduler_tick(self.config, self.seal, runner=lambda task, profile, timeout: calls.append(task) or "ok"), 0)
|
||
self.assertEqual(calls, [])
|
||
|
||
def test_schedule_repl(self):
|
||
agent = mock.Mock()
|
||
agent.store = self.store
|
||
agent.profile = "t"
|
||
with mock.patch("builtins.print") as printer:
|
||
tai.handle_command(agent, "/schedule every 1h check mail")
|
||
tai.handle_command(agent, "/schedule at 2026-10-08T09:00 water plants")
|
||
created = "\n".join(str(call.args[0]) for call in printer.call_args_list if call.args)
|
||
self.assertIn("scheduled #1", created)
|
||
self.assertIn("scheduled #2", created)
|
||
with mock.patch("builtins.print") as printer:
|
||
tai.handle_command(agent, "/schedules")
|
||
listed = "\n".join(str(call.args[0]) for call in printer.call_args_list if call.args)
|
||
self.assertIn("#1", listed)
|
||
self.assertIn("check mail", listed)
|
||
with mock.patch("builtins.print") as printer:
|
||
tai.handle_command(agent, "/schedule soon x")
|
||
tai.handle_command(agent, "/schedule every 5s x")
|
||
tai.handle_command(agent, "/schedule at nope x")
|
||
tai.handle_command(agent, "/unschedule 1")
|
||
tai.handle_command(agent, "/unschedule 1")
|
||
tai.handle_command(agent, "/unschedule x")
|
||
errors = "\n".join(str(call.args[0]) for call in printer.call_args_list if call.args)
|
||
self.assertIn("use /schedule at", errors)
|
||
self.assertIn("at least 60", errors)
|
||
self.assertIn("invalid datetime", errors)
|
||
self.assertIn("deleted schedule #1", errors)
|
||
self.assertIn("no schedule #1", errors)
|
||
self.assertIn("use /unschedule", errors)
|
||
|
||
def test_delegation_wording(self):
|
||
self.assertIn("Delegation:", tai.DEFAULT_SYSTEM)
|
||
self.assertIn("fork background subagents", tai.DEFAULT_SYSTEM)
|
||
self.assertIn("schedule future work", tai.DEFAULT_SYSTEM)
|
||
|
||
|
||
class RecordGraphTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "record-test-1"))
|
||
self.app = FakeSecretApp(self.store)
|
||
self.tools = tai.Tools(self.app)
|
||
with tai.AGENTS_LOCK:
|
||
tai.AGENTS.clear()
|
||
tai.AGENTS_NEXT[0] = 1
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
with tai.AGENTS_LOCK:
|
||
tai.AGENTS.clear()
|
||
tai.AGENTS_NEXT[0] = 1
|
||
|
||
def mem_id(self, text):
|
||
match = re.search(r"mem:[0-9a-f]{16}", text)
|
||
self.assertIsNotNone(match, "no mem id in %r" % text)
|
||
return match.group(0)
|
||
|
||
def test_save_read_search_delete_roundtrip(self):
|
||
saved = self.tools.dispatch("record_save", json.dumps({"title": "site notes", "content": "alpha beta gamma", "kind": "note", "tags": ["Site Visit"]}))
|
||
record_id = self.mem_id(saved)
|
||
self.assertIn("16 chars", saved)
|
||
self.assertIn("site-visit", saved)
|
||
page = self.tools.dispatch("record_read", json.dumps({"id": record_id}))
|
||
self.assertIn("chars 0-16 of 16", page)
|
||
self.assertIn("alpha beta gamma", page)
|
||
found = self.tools.dispatch("record_search", json.dumps({"query": "beta"}))
|
||
self.assertIn(record_id, found)
|
||
self.assertIn("site notes", found)
|
||
deleted = self.tools.dispatch("record_delete", json.dumps({"id": record_id}))
|
||
self.assertIn("deleted record", deleted)
|
||
self.assertIn("unknown record", self.tools.dispatch("record_read", json.dumps({"id": record_id})))
|
||
self.assertEqual(self.app.approvals, ["delete record " + record_id])
|
||
|
||
def test_read_pages_offsets(self):
|
||
saved = self.tools.dispatch("record_save", json.dumps({"title": "paged", "content": "".join("%04d" % num for num in range(250))}))
|
||
record_id = self.mem_id(saved)
|
||
first = self.tools.dispatch("record_read", json.dumps({"id": record_id, "limit": 100}))
|
||
self.assertIn("chars 0-100 of 1000", first)
|
||
second = self.tools.dispatch("record_read", json.dumps({"id": record_id, "offset": 100, "limit": 100}))
|
||
self.assertIn("chars 100-200 of 1000", second)
|
||
self.assertNotEqual(first.splitlines()[-1], second.splitlines()[-1])
|
||
record = self.store.get_record(record_id)
|
||
self.assertGreaterEqual(record["reads"], 3)
|
||
|
||
def test_record_validation_errors(self):
|
||
self.assertIn("content is empty", self.tools.dispatch("record_save", json.dumps({"title": "x"})))
|
||
self.assertIn("unknown kind", self.tools.dispatch("record_save", json.dumps({"content": "x", "kind": "song"})))
|
||
self.assertIn("tags must be a list", self.tools.dispatch("record_save", json.dumps({"content": "x", "tags": "nope"})))
|
||
self.assertIn("unknown record", self.tools.dispatch("record_read", json.dumps({"id": "mem:0123456789abcdef"})))
|
||
self.assertIn("no records match", self.tools.dispatch("record_search", json.dumps({"query": "nothing-here-zzz"})))
|
||
|
||
def test_delete_denied_keeps_record(self):
|
||
saved = self.tools.dispatch("record_save", json.dumps({"content": "keep me"}))
|
||
record_id = self.mem_id(saved)
|
||
self.app.approve = False
|
||
self.assertIn("denied by user", self.tools.dispatch("record_delete", json.dumps({"id": record_id})))
|
||
self.assertIsNotNone(self.store.get_record(record_id))
|
||
|
||
def test_search_filters_kind_and_tags(self):
|
||
self.tools.dispatch("record_save", json.dumps({"title": "a", "content": "shared word", "kind": "note", "tags": ["team"]}))
|
||
self.tools.dispatch("record_save", json.dumps({"title": "b", "content": "shared word", "kind": "output", "tags": ["team"]}))
|
||
by_kind = self.tools.dispatch("record_search", json.dumps({"query": "shared", "kind": "output"}))
|
||
self.assertIn("[output]", by_kind)
|
||
self.assertNotIn("[note]", by_kind)
|
||
by_tag = self.tools.dispatch("record_search", json.dumps({"tags": ["team"]}))
|
||
self.assertIn("[output]", by_tag)
|
||
self.assertIn("[note]", by_tag)
|
||
missing = self.tools.dispatch("record_search", json.dumps({"tags": ["team", "other"]}))
|
||
self.assertIn("no records match", missing)
|
||
|
||
def test_graph_link_and_query(self):
|
||
first = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "plan", "content": "the plan"})))
|
||
second = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "bravo-log-title", "content": "the log"})))
|
||
self.tools.dispatch("store_secret", json.dumps({"name": "deploy", "value": "token-xyz-9"}))
|
||
linked = self.tools.dispatch("graph_link", json.dumps({"src": first, "dst": second, "relation": "Follows Up"}))
|
||
self.assertIn("follows-up", linked)
|
||
self.tools.dispatch("graph_link", json.dumps({"src": second, "dst": "secret:deploy", "relation": "uses"}))
|
||
view = self.tools.dispatch("graph_query", json.dumps({"node": first}))
|
||
self.assertIn(first + " :: plan", view)
|
||
self.assertIn(second, view)
|
||
self.assertIn("secret:deploy", view)
|
||
self.assertIn("bravo-log-title", view)
|
||
self.assertIn("unknown node", self.tools.dispatch("graph_query", json.dumps({"node": "mem:0123456789abcdef"})))
|
||
self.assertIn("unknown node", self.tools.dispatch("graph_link", json.dumps({"src": first, "dst": "mem:0123456789abcdef"})))
|
||
|
||
def test_traverse_caps_depth(self):
|
||
ids = [self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "n%d" % num, "content": "x"}))) for num in range(6)]
|
||
for pos in range(5):
|
||
self.store.add_edge(ids[pos], ids[pos + 1], "next")
|
||
shallow = self.store.traverse(ids[0], depth=1, limit=50)
|
||
self.assertEqual(len(shallow), 2)
|
||
deep = self.store.traverse(ids[0], depth=99, limit=50)
|
||
self.assertEqual(len(deep), 5)
|
||
self.assertEqual(max(item["depth"] for item in deep), 4)
|
||
|
||
def test_shell_output_spills_to_record(self):
|
||
result = self.tools.dispatch("shell", json.dumps({"command": "python3 -c \"print('0123456789' * 700)\""}))
|
||
self.assertIn("exit 0", result)
|
||
record_id = self.mem_id(result)
|
||
self.assertIn("record_read pages the rest", result)
|
||
page = self.tools.dispatch("record_read", json.dumps({"id": record_id, "limit": 20}))
|
||
self.assertIn("of 7000", page)
|
||
found = self.tools.dispatch("record_search", json.dumps({"tags": ["shell"]}))
|
||
self.assertIn(record_id, found)
|
||
|
||
def test_poll_spills_once_and_reuses(self):
|
||
big = "z" * 7000
|
||
agent_id = tai.spawn_agent("big task", "t", 60, None, None, 0, lambda task, profile, timeout: big)
|
||
first = self.tools.dispatch("poll", json.dumps({"id": agent_id, "wait": 5}))
|
||
record_id = self.mem_id(first)
|
||
second = self.tools.dispatch("poll", json.dumps({"id": agent_id}))
|
||
self.assertIn(record_id, second)
|
||
rows = self.store.db.execute("SELECT COUNT(*) FROM records").fetchone()[0]
|
||
self.assertEqual(rows, 1)
|
||
|
||
def test_record_repl_commands(self):
|
||
record_id = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "repl note", "content": "visible words"})))
|
||
self.app.tools = self.tools
|
||
with mock.patch("builtins.print") as shown:
|
||
tai.handle_command(self.app, "/records visible")
|
||
tai.handle_command(self.app, "/record " + record_id)
|
||
tai.handle_command(self.app, "/graph " + record_id)
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list)
|
||
self.assertIn(record_id, printed)
|
||
self.assertIn("repl note", printed)
|
||
|
||
def test_wal_mode_enabled(self):
|
||
mode = self.store.db.execute("PRAGMA journal_mode").fetchone()[0]
|
||
self.assertEqual(mode.lower(), "wal")
|
||
|
||
|
||
class SecretMetaTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(config, tai.Seal(config.home, "meta-test-1"))
|
||
self.tools = tai.Tools(FakeSecretApp(self.store))
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def test_metadata_roundtrip(self):
|
||
result = self.tools.dispatch("store_secret", json.dumps({"name": "db", "value": "pw-12345", "username": "ops", "host": "db.internal", "port": 5432, "notes": "primary", "tags": ["Prod DB"]}))
|
||
self.assertIn("stored secret 'db'", result)
|
||
listed = self.tools.dispatch("list_secrets", "{}")
|
||
self.assertIn("ops@db.internal:5432", listed)
|
||
self.assertIn("prod-db", listed)
|
||
self.assertNotIn("pw-12345", listed)
|
||
infos = self.store.list_secret_infos()
|
||
self.assertEqual(infos[0]["meta"].get("username"), "ops")
|
||
self.assertEqual(infos[0]["meta"].get("port"), 5432)
|
||
|
||
def test_minimum_info_suffices(self):
|
||
result = self.tools.dispatch("store_secret", json.dumps({"name": "plain", "value": "v-abcdef"}))
|
||
self.assertIn("stored secret 'plain'", result)
|
||
self.assertIn("- plain [secret]", self.tools.dispatch("list_secrets", "{}"))
|
||
|
||
|
||
class ScheduleTagTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "schetag-1"))
|
||
self.tools = tai.Tools(FakeSchedulerApp(self.store))
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def test_schedule_tags_shown(self):
|
||
self.tools.dispatch("schedule", json.dumps({"prompt": "water plants", "every": 3600, "tags": ["Home Chores"]}))
|
||
listed = self.tools.dispatch("schedules", "{}")
|
||
self.assertIn("home-chores", listed)
|
||
items = self.store.list_schedules()
|
||
self.assertIn("schedule", items[0]["tags"])
|
||
|
||
|
||
class ModeTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = False
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "mode-test-1"))
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def test_interactive_always_enables_yolo(self):
|
||
agent = tai.Agent(self.config, self.store, persist=True, quiet=True)
|
||
with mock.patch("sys.stdin") as fake_stdin:
|
||
fake_stdin.isatty.return_value = True
|
||
with mock.patch("builtins.input", return_value="Y"):
|
||
with mock.patch("builtins.print"):
|
||
self.assertTrue(agent.ask_approval("do it"))
|
||
self.assertTrue(agent.yolo)
|
||
with mock.patch("builtins.input", side_effect=AssertionError("must not ask")):
|
||
self.assertTrue(agent.ask_approval("do it again"))
|
||
|
||
def test_yolo_never_asks(self):
|
||
agent = tai.Agent(self.config, self.store, persist=True, quiet=True, yolo=True)
|
||
with mock.patch("builtins.input", side_effect=AssertionError("must not ask")):
|
||
self.assertTrue(agent.ask_approval("anything"))
|
||
|
||
def test_auto_gets_autonomous_note(self):
|
||
agent = tai.Agent(self.config, self.store, persist=True, quiet=True, auto=True)
|
||
agent.apply_system()
|
||
self.assertIn("Autonomous mode", agent.messages[0]["content"])
|
||
|
||
|
||
class AuditTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "audit-test-1"))
|
||
self.app = FakeSecretApp(self.store)
|
||
self.app.skills = {}
|
||
self.tools = tai.Tools(self.app)
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def workfile(self, name="note.txt", content="v1"):
|
||
path = os.path.join(self.tmp.name, name)
|
||
with open(path, "w", encoding="utf-8") as handle:
|
||
handle.write(content)
|
||
return path
|
||
|
||
def test_write_audits_and_records(self):
|
||
path = os.path.join(self.tmp.name, "fresh.txt")
|
||
self.assertIn("wrote", self.tools.dispatch("write_file", json.dumps({"path": path, "content": "hello"})))
|
||
rows = self.store.audit_history(path)
|
||
self.assertEqual(len(rows), 1)
|
||
self.assertEqual(rows[0]["action"], "write")
|
||
self.assertIsNone(rows[0]["old_size"])
|
||
self.assertEqual(rows[0]["new_size"], 5)
|
||
record_id = self.store.file_record_id(path)
|
||
self.assertIsNotNone(record_id)
|
||
found = self.tools.dispatch("record_search", json.dumps({"tags": ["txt"]}))
|
||
self.assertIn(record_id, found)
|
||
self.tools.dispatch("read_file", json.dumps({"path": path}))
|
||
self.tools.dispatch("write_file", json.dumps({"path": path, "content": "hello again"}))
|
||
rows = self.store.audit_history(path)
|
||
self.assertEqual(len(rows), 2)
|
||
self.assertEqual(rows[0]["old_size"], 5)
|
||
self.assertEqual(rows[0]["new_size"], 11)
|
||
|
||
def test_edit_captures_old_and_new(self):
|
||
path = self.workfile()
|
||
self.tools.dispatch("read_file", json.dumps({"path": path}))
|
||
self.assertIn("edited", self.tools.dispatch("edit_file", json.dumps({"path": path, "find": "v1", "replace": "v2"})))
|
||
row = self.store.audit_get(self.store.audit_history(path)[0]["id"])
|
||
self.assertEqual(row["action"], "edit")
|
||
self.assertEqual(row["old"], "v1")
|
||
self.assertEqual(row["new"], "v2")
|
||
|
||
def test_big_file_truncates_with_marker(self):
|
||
path = os.path.join(self.tmp.name, "big.txt")
|
||
content = "x" * 60000
|
||
self.tools.dispatch("write_file", json.dumps({"path": path, "content": content}))
|
||
row = self.store.audit_get(self.store.audit_history(path)[0]["id"])
|
||
self.assertEqual(row["new_size"], 60000)
|
||
self.assertIn("[truncated, full size 60000 bytes]", row["new"])
|
||
record = self.store.get_record(self.store.file_record_id(path))
|
||
self.assertIn("[truncated, full size 60000 bytes]", record["content"])
|
||
|
||
def test_shell_rm_snapshots_before_delete(self):
|
||
first = self.workfile("a.txt", "alpha")
|
||
second = self.workfile("b.txt", "beta")
|
||
result = self.tools.dispatch("shell", json.dumps({"command": "rm a.txt b.txt", "workdir": self.tmp.name}))
|
||
self.assertIn("exit 0", result)
|
||
self.assertFalse(os.path.exists(first))
|
||
snaps = self.store.audit_history(tag="shell")
|
||
self.assertEqual(len(snaps), 2)
|
||
by_path = {row["path"]: row for row in snaps}
|
||
full = self.store.audit_get(by_path[first]["id"])
|
||
self.assertEqual(full["old"], "alpha")
|
||
self.assertIsNone(full["new"])
|
||
self.assertIn("rm a.txt b.txt", full["message"])
|
||
|
||
def test_shell_redirect_snapshots_target(self):
|
||
path = self.workfile("out.txt", "old words")
|
||
self.tools.dispatch("shell", json.dumps({"command": "echo new > out.txt", "workdir": self.tmp.name}))
|
||
snaps = self.store.audit_history(tag="shell")
|
||
self.assertEqual(len(snaps), 1)
|
||
full = self.store.audit_get(snaps[0]["id"])
|
||
self.assertEqual(full["old"], "old words")
|
||
|
||
def test_shell_parser_units(self):
|
||
self.assertEqual(tai.segment_targets(["rm", "-rf", "a", "b"]), ["a", "b"])
|
||
self.assertEqual(tai.segment_targets(["sudo", "rm", "x"]), ["x"])
|
||
self.assertEqual(tai.segment_targets(["VAR=1", "cmd", ">", "out"]), ["out"])
|
||
self.assertEqual(tai.segment_targets(["tee", "t1", "t2"]), ["t1", "t2"])
|
||
self.assertEqual(tai.segment_targets(["mv", "a", "b"]), ["b"])
|
||
self.assertEqual(tai.segment_targets(["cmd", "2>", "/dev/null"]), [])
|
||
self.assertEqual(tai.segment_targets(["truncate", "-s", "0", "f"]), ["f"])
|
||
self.assertEqual(tai.segment_targets(["dd", "if=a", "of=b"]), ["b"])
|
||
sub = os.path.join(self.tmp.name, "sub")
|
||
os.makedirs(sub)
|
||
self.workfile("sub/one.txt", "1")
|
||
self.workfile("sub/two.txt", "2")
|
||
hits = tai.shell_target_paths("rm -rf sub", self.tmp.name)
|
||
self.assertEqual(len(hits), 2)
|
||
first = self.workfile("t1.txt", "1")
|
||
piped = tai.shell_target_paths("echo x | tee t1.txt missing.txt", self.tmp.name)
|
||
self.assertEqual(piped, [first])
|
||
|
||
def test_delete_roundtrip_and_deny(self):
|
||
gone = self.workfile("gone.txt", "bye")
|
||
kept = self.workfile("kept.txt", "hi")
|
||
self.tools.dispatch("read_file", json.dumps({"path": gone}))
|
||
self.tools.dispatch("read_file", json.dumps({"path": kept}))
|
||
self.assertIn("deleted", self.tools.dispatch("delete_file", json.dumps({"path": gone})))
|
||
self.assertFalse(os.path.exists(gone))
|
||
self.assertIsNone(self.store.file_record_id(gone))
|
||
row = self.store.audit_get(self.store.audit_history(gone)[0]["id"])
|
||
self.assertEqual(row["action"], "delete")
|
||
self.assertEqual(row["old"], "bye")
|
||
self.app.approve = False
|
||
self.assertIn("denied by user", self.tools.dispatch("delete_file", json.dumps({"path": kept})))
|
||
self.assertTrue(os.path.exists(kept))
|
||
|
||
def test_audit_tool_filters(self):
|
||
first = self.workfile("f1.txt", "one")
|
||
self.workfile("f2.txt", "two")
|
||
self.tools.dispatch("read_file", json.dumps({"path": first}))
|
||
self.tools.dispatch("edit_file", json.dumps({"path": first, "find": "one", "replace": "uno"}))
|
||
by_path = self.tools.dispatch("audit", json.dumps({"path": first}))
|
||
self.assertIn("edit", by_path)
|
||
self.assertNotIn("f2.txt", by_path)
|
||
self.assertIn("invalid limit", self.tools.dispatch("audit", json.dumps({"limit": "zzz"})))
|
||
self.assertIn("no audit rows match", self.tools.dispatch("audit", json.dumps({"path": "/nope/nothing"})))
|
||
|
||
def test_restore_post_pre_and_undelete(self):
|
||
path = self.workfile("time.txt", "v1")
|
||
self.tools.dispatch("read_file", json.dumps({"path": path}))
|
||
self.tools.dispatch("edit_file", json.dumps({"path": path, "find": "v1", "replace": "v2"}))
|
||
rows = {row["action"]: row for row in self.store.audit_history(path)}
|
||
self.assertIn("restored", self.tools.dispatch("restore", json.dumps({"id": rows["edit"]["id"]})))
|
||
with open(path, encoding="utf-8") as handle:
|
||
self.assertEqual(handle.read(), "v2")
|
||
self.tools.dispatch("shell", json.dumps({"command": "rm time.txt", "workdir": self.tmp.name}))
|
||
snap = self.store.audit_history(path, tag="shell")[0]
|
||
self.assertIn("restored", self.tools.dispatch("restore", json.dumps({"id": snap["id"]})))
|
||
with open(path, encoding="utf-8") as handle:
|
||
self.assertEqual(handle.read(), "v2")
|
||
history = self.store.audit_history(path)
|
||
self.assertEqual(history[0]["action"], "restore")
|
||
self.assertIn("restored from audit", history[0]["message"])
|
||
|
||
def test_restore_refuses_truncated(self):
|
||
path = os.path.join(self.tmp.name, "huge.txt")
|
||
self.tools.dispatch("write_file", json.dumps({"path": path, "content": "y" * 60000}))
|
||
row_id = self.store.audit_history(path)[0]["id"]
|
||
result = self.tools.dispatch("restore", json.dumps({"id": row_id}))
|
||
self.assertIn("truncated", result)
|
||
self.assertIn("cannot restore safely", result)
|
||
|
||
def test_restore_repl_command(self):
|
||
path = self.workfile("repl.txt", "keep")
|
||
self.tools.dispatch("read_file", json.dumps({"path": path}))
|
||
self.tools.dispatch("edit_file", json.dumps({"path": path, "find": "keep", "replace": "changed"}))
|
||
row_id = self.store.audit_history(path)[0]["id"]
|
||
self.app.tools = self.tools
|
||
with mock.patch("builtins.print") as shown:
|
||
tai.handle_command(self.app, "/audit " + path)
|
||
tai.handle_command(self.app, "/restore %d" % row_id)
|
||
tai.handle_command(self.app, "/restore nope")
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list)
|
||
self.assertIn("edit", printed)
|
||
self.assertIn("restored", printed)
|
||
self.assertIn("use /restore <audit id>", printed)
|
||
|
||
def test_lazy_blueprint_builds_on_load(self):
|
||
with mock.patch.object(self.tools, "run_create_skill", return_value="created") as creator:
|
||
result = self.tools.dispatch("load_skill", json.dumps({"name": "bot-creator"}))
|
||
self.assertIn("building skill 'bot-creator' from blueprint", result)
|
||
creator.assert_called_once()
|
||
sent = creator.call_args.args[0]
|
||
self.assertEqual(sent["name"], "bot-creator")
|
||
self.assertIn("Deep-research", sent["brief"])
|
||
self.assertIn("unknown skill", self.tools.dispatch("load_skill", json.dumps({"name": "nope"})))
|
||
|
||
|
||
class ReleaseTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "release-test-1"))
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def script(self, version="1.2.3"):
|
||
path = os.path.join(self.tmp.name, "tai-copy.py")
|
||
with open(path, "w", encoding="utf-8") as handle:
|
||
handle.write("#!/usr/bin/env python3\nVERSION = \"%s\"\nprint('hi')\n" % version)
|
||
return path
|
||
|
||
def test_bump_parts(self):
|
||
self.assertEqual(tai.release_bump(self.script(), "patch"), ("1.2.3", "1.2.4"))
|
||
self.assertEqual(tai.release_bump(self.script(), "minor"), ("1.2.3", "1.3.0"))
|
||
self.assertEqual(tai.release_bump(self.script(), "major"), ("1.2.3", "2.0.0"))
|
||
with open(self.script(), encoding="utf-8") as handle:
|
||
self.assertIn('VERSION = "1.2.3"', handle.read())
|
||
with self.assertRaises(ValueError):
|
||
tai.release_bump(self.script(), "banana")
|
||
flat = os.path.join(self.tmp.name, "flat.py")
|
||
with open(flat, "w", encoding="utf-8") as handle:
|
||
handle.write("no version here\n")
|
||
with self.assertRaises(ValueError):
|
||
tai.release_bump(flat, "patch")
|
||
|
||
def test_do_release_end_to_end(self):
|
||
path = self.script()
|
||
folder = tai.backups_dir(self.tmp.name)
|
||
old, new, dest, row_id = tai.do_release(self.store, path, folder, "minor", "add audit trail")
|
||
self.assertEqual((old, new), ("1.2.3", "1.3.0"))
|
||
self.assertTrue(dest.startswith(folder))
|
||
self.assertIn("tai-1.3.0-", os.path.basename(dest))
|
||
with open(path, encoding="utf-8") as handle:
|
||
text = handle.read()
|
||
self.assertIn('VERSION = "1.3.0"', text)
|
||
self.assertIn("print('hi')", text)
|
||
row = self.store.audit_get(row_id)
|
||
self.assertEqual(row["action"], "release")
|
||
self.assertEqual(row["message"], "add audit trail")
|
||
self.assertEqual(row["old"], "1.2.3")
|
||
self.assertEqual(row["new"], "1.3.0")
|
||
self.assertIn("release", row["tags"])
|
||
with self.assertRaises(ValueError):
|
||
tai.do_release(self.store, path, folder, "patch", " ")
|
||
|
||
def test_backup_prune_and_skip(self):
|
||
folder = tai.backups_dir(self.tmp.name)
|
||
os.makedirs(folder)
|
||
now = time.time()
|
||
for pos in range(12):
|
||
name = os.path.join(folder, "tai-1.0.0-20200101T00000%dZ-abc%d.py" % (pos, pos))
|
||
with open(name, "w", encoding="utf-8") as handle:
|
||
handle.write("old")
|
||
stamp = now - (12 - pos) * 60
|
||
os.utime(name, (stamp, stamp))
|
||
path = self.script()
|
||
dest, _digest = tai.snapshot_self(path, folder, "1.2.3")
|
||
kept = [entry for entry in os.listdir(folder) if entry.endswith(".py")]
|
||
self.assertEqual(len(kept), tai.BACKUP_KEEP)
|
||
self.assertIn(os.path.basename(dest), kept)
|
||
with open(path, "w", encoding="utf-8") as handle:
|
||
handle.write("#!/usr/bin/env python3\nVERSION = \"9.9.9\"\nprint('changed')\n")
|
||
first = tai.ensure_self_backup(self.config, self.store, script_file=path)
|
||
self.assertIsNotNone(first)
|
||
again = tai.ensure_self_backup(self.config, self.store, script_file=path)
|
||
self.assertIsNone(again)
|
||
snaps = self.store.audit_history(tag="snapshot")
|
||
self.assertEqual(len(snaps), 1)
|
||
self.assertIn("backed up as", snaps[0]["message"])
|
||
|
||
def test_release_tool_validates_safely(self):
|
||
app = FakeSecretApp(self.store, approve=False)
|
||
tools = tai.Tools(app)
|
||
self.assertIn("part must be", tools.dispatch("release", json.dumps({"part": "banana", "message": "x"})))
|
||
self.assertIn("message is required", tools.dispatch("release", json.dumps({"part": "patch", "message": ""})))
|
||
denied = tools.dispatch("release", json.dumps({"part": "patch", "message": "try bump"}))
|
||
self.assertIn("denied by user", denied)
|
||
self.assertEqual(len(app.approvals), 1)
|
||
self.assertIn("release %s (patch)" % tai.next_version(tai.VERSION, "patch"), app.approvals[0])
|
||
|
||
def test_boot_creates_self_backup(self):
|
||
import argparse as ap
|
||
|
||
args = ap.Namespace(profile="t", yes=True, yolo=False, auto=False)
|
||
with mock.patch.dict(os.environ, {"TAI_PASSPHRASE": "release-test-1"}):
|
||
with mock.patch("builtins.print"):
|
||
agent = tai.boot(args)
|
||
try:
|
||
names = os.listdir(tai.backups_dir(self.tmp.name))
|
||
self.assertEqual(len(names), 1)
|
||
self.assertTrue(names[0].startswith("tai-%s-" % tai.VERSION))
|
||
snaps = agent.store.audit_history(tag="snapshot")
|
||
self.assertEqual(len(snaps), 1)
|
||
finally:
|
||
agent.store.close()
|
||
|
||
|
||
class TaggingTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "tag-test-1"))
|
||
self.app = FakeSecretApp(self.store)
|
||
self.tools = tai.Tools(self.app)
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def mem_id(self, text):
|
||
return re.search(r"mem:[0-9a-f]{16}", text).group(0)
|
||
|
||
def test_singular_units(self):
|
||
self.assertEqual(tai.singular_noun("servers"), "server")
|
||
self.assertEqual(tai.singular_noun("cities"), "city")
|
||
self.assertEqual(tai.singular_noun("boxes"), "box")
|
||
self.assertEqual(tai.singular_noun("branches"), "branch")
|
||
self.assertEqual(tai.singular_noun("glass"), "glass")
|
||
self.assertEqual(tai.singular_noun("status"), "status")
|
||
self.assertEqual(tai.singular_noun("news"), "news")
|
||
self.assertEqual(tai.singular_noun("physics"), "physics")
|
||
self.assertEqual(tai.singular_noun("api"), "api")
|
||
|
||
def test_plural_merges_into_known_singular(self):
|
||
first = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "one", "content": "about a server", "tags": ["server"]})))
|
||
self.assertIn("server", self.store.item_tags(first))
|
||
second = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "two", "content": "more", "tags": ["Servers"]})))
|
||
tags = self.store.item_tags(second)
|
||
self.assertIn("server", tags)
|
||
self.assertNotIn("servers", tags)
|
||
self.assertNotIn("servers", self.store.known_tags())
|
||
|
||
def test_irregular_words_never_corrupt(self):
|
||
record_id = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "words", "content": "plain", "tags": ["news", "glass", "status", "physics"]})))
|
||
tags = self.store.item_tags(record_id)
|
||
for word in ("news", "glass", "status", "physics"):
|
||
self.assertIn(word, tags)
|
||
|
||
def test_known_words_attach_automatically(self):
|
||
self.tools.dispatch("record_save", json.dumps({"title": "seed", "content": "nothing yet", "tags": ["deploy", "server"]}))
|
||
record_id = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "friday", "content": "we deploy the thing friday", "tags": []})))
|
||
self.assertIn("deploy", self.store.item_tags(record_id))
|
||
self.assertNotIn("friday", self.store.item_tags(record_id))
|
||
plural = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "fleet", "content": "all servers rebooted", "tags": []})))
|
||
self.assertIn("server", self.store.item_tags(plural))
|
||
|
||
def test_new_records_link_to_same_tag_neighbors(self):
|
||
first = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "a", "content": "deploy alpha", "tags": ["deploy"]})))
|
||
saved = self.tools.dispatch("record_save", json.dumps({"title": "b", "content": "deploy beta", "tags": ["deploy"]}))
|
||
second = self.mem_id(saved)
|
||
self.assertIn(first, saved)
|
||
edges = self.store.edges_for(second)
|
||
self.assertEqual(len(edges), 1)
|
||
self.assertEqual(edges[0]["relation"], "shares-deploy")
|
||
self.assertEqual(edges[0]["other"], first)
|
||
|
||
def test_baseline_only_records_do_not_link(self):
|
||
self.tools.dispatch("record_save", json.dumps({"title": "a", "content": "lorem ipsum", "tags": []}))
|
||
second = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "b", "content": "dolor sit", "tags": []})))
|
||
self.assertEqual(self.store.edges_for(second), [])
|
||
|
||
def test_auto_link_caps_at_three(self):
|
||
for pos in range(5):
|
||
self.tools.dispatch("record_save", json.dumps({"title": "n%d" % pos, "content": "shared topic here", "tags": ["topic"]}))
|
||
sixth = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "n5", "content": "shared topic again", "tags": ["topic"]})))
|
||
edges = [edge for edge in self.store.edges_for(sixth) if edge["direction"] == "out"]
|
||
self.assertEqual(len(edges), 3)
|
||
|
||
def test_search_expands_singular_plural(self):
|
||
legacy = self.store.add_record("note", "old", "zeta-content", ["servers"])
|
||
self.assertIn("servers", self.store.item_tags(legacy))
|
||
modern = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "new", "content": "zeta-content", "tags": ["server"]})))
|
||
found = self.tools.dispatch("record_search", json.dumps({"tags": ["server"]}))
|
||
self.assertIn(legacy, found)
|
||
self.assertIn(modern, found)
|
||
|
||
def test_tags_tool_counts_and_prefix(self):
|
||
self.tools.dispatch("record_save", json.dumps({"title": "a", "content": "x", "tags": ["deploy", "friday"]}))
|
||
self.tools.dispatch("record_save", json.dumps({"title": "b", "content": "y", "tags": ["deploy"]}))
|
||
listed = self.tools.dispatch("tags", "{}")
|
||
self.assertIn("deploy (2)", listed)
|
||
self.assertIn("friday (1)", listed)
|
||
self.assertLess(listed.index("deploy (2)"), listed.index("friday (1)"))
|
||
prefixed = self.tools.dispatch("tags", json.dumps({"prefix": "fri"}))
|
||
self.assertIn("friday (1)", prefixed)
|
||
self.assertNotIn("deploy", prefixed)
|
||
self.app.tools = self.tools
|
||
with mock.patch("builtins.print") as shown:
|
||
tai.handle_command(self.app, "/tags")
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list)
|
||
self.assertIn("deploy (2)", printed)
|
||
|
||
|
||
class ProfileTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "profile-test-1"))
|
||
self.app = FakeSecretApp(self.store)
|
||
self.tools = tai.Tools(self.app)
|
||
with tai.AGENTS_LOCK:
|
||
tai.AGENTS.clear()
|
||
tai.AGENTS_NEXT[0] = 1
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
with tai.AGENTS_LOCK:
|
||
tai.AGENTS.clear()
|
||
tai.AGENTS_NEXT[0] = 1
|
||
|
||
def mem_id(self, text):
|
||
return re.search(r"mem:[0-9a-f]{16}", text).group(0)
|
||
|
||
def test_records_isolated(self):
|
||
record_id = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "mine", "content": "t-only words", "tags": ["t-tag"]})))
|
||
self.assertIsNone(self.store.get_record(record_id, "other"))
|
||
self.assertEqual(self.store.search_records("t-only", profile="other"), [])
|
||
self.assertEqual(len(self.store.search_records("t-only", profile="t")), 1)
|
||
self.app.profile = "other"
|
||
self.assertIn("no records match", self.tools.dispatch("record_search", json.dumps({"query": "t-only"})))
|
||
self.assertIn("unknown record", self.tools.dispatch("record_read", json.dumps({"id": record_id})))
|
||
|
||
def test_secrets_isolated(self):
|
||
self.tools.dispatch("store_secret", json.dumps({"name": "api", "value": "t-value-123"}))
|
||
self.assertIsNone(self.store.load_secret("api", "other"))
|
||
self.assertEqual(self.store.list_secret_infos("other"), [])
|
||
self.store.save_secret("api", "other-value-456", None, None, "other")
|
||
self.assertEqual(self.store.load_secret("api", "t"), "t-value-123")
|
||
self.assertEqual(self.store.load_secret("api", "other"), "other-value-456")
|
||
self.assertIn("[redacted:api]", self.store.redact("leak t-value-123 here", "t"))
|
||
self.assertIn("t-value-123", self.store.redact("leak t-value-123 here", "other"))
|
||
self.app.profile = "other"
|
||
listed = self.tools.dispatch("list_secrets", "{}")
|
||
self.assertIn("api", listed)
|
||
self.assertNotIn("t-value-123", listed)
|
||
|
||
def test_tags_vocab_isolated(self):
|
||
self.tools.dispatch("record_save", json.dumps({"title": "seed", "content": "nothing", "tags": ["deploy"]}))
|
||
self.app.profile = "other"
|
||
record_id = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "fresh", "content": "we deploy friday", "tags": []})))
|
||
self.assertNotIn("deploy", self.store.item_tags(record_id, "other"))
|
||
self.assertEqual(self.store.tag_counts(profile="other"), [("record", 1)])
|
||
self.assertIn("deploy", self.store.known_tags("t"))
|
||
|
||
def test_edges_isolated(self):
|
||
first = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "a", "content": "x", "tags": ["linkable"]})))
|
||
second = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "b", "content": "y", "tags": ["unlinked"]})))
|
||
self.tools.dispatch("graph_link", json.dumps({"src": first, "dst": second, "relation": "uses"}))
|
||
self.app.profile = "other"
|
||
self.assertIn("unknown node", self.tools.dispatch("graph_query", json.dumps({"node": first})))
|
||
self.assertEqual(self.store.edges_for(first, "other"), [])
|
||
self.app.profile = "t"
|
||
view = self.tools.dispatch("graph_query", json.dumps({"node": first}))
|
||
self.assertIn("uses", view)
|
||
|
||
def test_audit_isolated(self):
|
||
path = os.path.join(self.tmp.name, "aud.txt")
|
||
self.tools.dispatch("write_file", json.dumps({"path": path, "content": "t-data"}))
|
||
self.assertEqual(len(self.store.audit_history(profile="t")), 1)
|
||
self.assertEqual(self.store.audit_history(profile="other"), [])
|
||
row_id = self.store.audit_history(profile="t")[0]["id"]
|
||
self.assertIsNone(self.store.audit_get(row_id, "other"))
|
||
self.app.profile = "other"
|
||
self.assertIn("no audit rows match", self.tools.dispatch("audit", "{}"))
|
||
self.assertIn("unknown audit", self.tools.dispatch("restore", json.dumps({"id": row_id})))
|
||
|
||
def test_schedules_isolated(self):
|
||
app = FakeSchedulerApp(self.store)
|
||
tools = tai.Tools(app)
|
||
tools.dispatch("schedule", json.dumps({"prompt": "t-job", "every": 3600}))
|
||
self.assertEqual(len(self.store.list_schedules("t")), 1)
|
||
self.assertEqual(self.store.list_schedules("other"), [])
|
||
row_id = self.store.list_schedules("t")[0]["id"]
|
||
app.profile = "other"
|
||
self.assertIn("no schedule", tools.dispatch("unschedule", json.dumps({"id": row_id})))
|
||
self.assertFalse(self.store.remove_schedule(row_id, "other"))
|
||
self.assertTrue(self.store.remove_schedule(row_id, "t"))
|
||
self.assertIn("another profile", tools.dispatch("schedule", json.dumps({"prompt": "x", "every": 60, "profile": "t"})))
|
||
|
||
def test_fork_and_poll_scoped(self):
|
||
first = tai.spawn_agent("task-t", "t", 60, None, None, 0, lambda task, profile, timeout: "done-t")
|
||
second = tai.spawn_agent("task-other", "other", 60, None, None, 0, lambda task, profile, timeout: "done-other")
|
||
self.assertEqual(tai.poll_agent(first, wait=5, profile="t")[0], "done")
|
||
self.assertEqual(tai.poll_agent(second, profile="t")[0], "missing")
|
||
self.assertEqual([item["id"] for item in tai.list_agents("t")], [first])
|
||
app = FakePrincipal()
|
||
app.store = self.store
|
||
tools = tai.Tools(app)
|
||
self.assertIn("done-t", tools.dispatch("poll", json.dumps({"id": first})))
|
||
self.assertIn("missing", tools.dispatch("poll", json.dumps({"id": second})))
|
||
self.assertIn("another profile", tools.dispatch("fork", json.dumps({"task": "x", "profile": "other"})))
|
||
|
||
def test_switch_resets_identity(self):
|
||
agent = tai.Agent(self.config, self.store, persist=True, quiet=True)
|
||
agent.tools.read_files.add(("home", "/tmp/x"))
|
||
agent.tools.secret_grants.add(("shell", "", ("api",)))
|
||
with mock.patch("builtins.print"):
|
||
agent.switch_profile("newbie")
|
||
self.assertEqual(agent.profile, "newbie")
|
||
self.assertEqual(self.store.profile, "newbie")
|
||
self.assertEqual(agent.tools.read_files, set())
|
||
self.assertEqual(agent.tools.secret_grants, set())
|
||
self.assertEqual([item["role"] for item in agent.messages], ["system"])
|
||
self.assertIn("newbie", self.store.list_profiles())
|
||
self.assertIn("t", self.store.list_profiles())
|
||
|
||
def test_recall_scoped(self):
|
||
self.store.log_event("t", "user", "message", "t recall marker")
|
||
self.store.log_event("other", "user", "message", "other recall marker")
|
||
found = self.tools.dispatch("recall", json.dumps({"query": "recall marker"}))
|
||
self.assertIn("t recall marker", found)
|
||
self.assertNotIn("other recall marker", found)
|
||
|
||
def test_profile_list_is_global(self):
|
||
self.store.save_system("alpha", "system a")
|
||
self.store.save_system("beta", "system b")
|
||
self.assertIn("alpha", self.store.list_profiles())
|
||
self.assertIn("beta", self.store.list_profiles())
|
||
|
||
|
||
class MigrationTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "default"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def test_legacy_db_migrates_to_default(self):
|
||
db = sqlite3.connect(self.config.db_path)
|
||
db.execute("CREATE TABLE secrets (name TEXT PRIMARY KEY, value TEXT, meta TEXT, updated TEXT)")
|
||
db.execute("INSERT INTO secrets VALUES ('k', 'v-legacy', '{}', '2026-01-01')")
|
||
db.execute("CREATE TABLE records (id TEXT PRIMARY KEY, kind TEXT, title TEXT, content TEXT, size INTEGER, reads INTEGER, created TEXT, updated TEXT)")
|
||
db.execute("INSERT INTO records VALUES ('mem:0123456789abcdef', 'note', 'old', 'old words', 9, 0, '', '')")
|
||
db.execute("CREATE TABLE tags (item TEXT, tag TEXT, PRIMARY KEY (item, tag))")
|
||
db.execute("INSERT INTO tags VALUES ('mem:0123456789abcdef', 'legacy')")
|
||
db.execute("CREATE TABLE edges (src TEXT, dst TEXT, relation TEXT, created TEXT, PRIMARY KEY (src, dst, relation))")
|
||
db.execute("CREATE TABLE audit (id INTEGER PRIMARY KEY, ts TEXT, actor TEXT, action TEXT, path TEXT, message TEXT, old_size INTEGER, new_size INTEGER, old TEXT, new TEXT, tags TEXT)")
|
||
db.execute("INSERT INTO audit VALUES (1, 't', 'main', 'write', '/x', '', 0, 9, NULL, 'old words', '')")
|
||
db.execute("CREATE TABLE events (id INTEGER PRIMARY KEY, profile TEXT, ts TEXT, role TEXT, kind TEXT, text TEXT, tags TEXT)")
|
||
db.execute("CREATE TABLE schedules (id INTEGER PRIMARY KEY, name TEXT, prompt TEXT, profile TEXT, every_sec INTEGER, next_run TEXT, timeout INTEGER, status TEXT, last_status TEXT, last_result TEXT, created TEXT, updated TEXT)")
|
||
db.commit()
|
||
db.close()
|
||
store = tai.Store(self.config, tai.Seal(self.config.home, ""))
|
||
try:
|
||
self.assertEqual(store.load_secret("k"), "v-legacy")
|
||
store.save_secret("k", "v-work", None, None, "work")
|
||
self.assertEqual(store.load_secret("k", "work"), "v-work")
|
||
self.assertEqual(store.load_secret("k"), "v-legacy")
|
||
self.assertEqual(store.get_record("mem:0123456789abcdef")["content"], "old words")
|
||
self.assertIn("legacy", store.item_tags("mem:0123456789abcdef"))
|
||
self.assertEqual(len(store.audit_history()), 1)
|
||
self.assertEqual(tai.table_pk_columns(store.db, "secrets"), ["profile", "name"])
|
||
self.assertEqual(tai.table_pk_columns(store.db, "tags"), ["item", "tag", "profile"])
|
||
finally:
|
||
store.close()
|
||
|
||
def test_sparse_legacy_db_auto_heals(self):
|
||
db = sqlite3.connect(self.config.db_path)
|
||
db.execute("CREATE TABLE secrets (name TEXT PRIMARY KEY, value TEXT, updated TEXT)")
|
||
db.execute("INSERT INTO secrets VALUES ('srv', 'pw', 't')")
|
||
db.execute("CREATE TABLE events (id INTEGER PRIMARY KEY, profile TEXT, ts TEXT, role TEXT, kind TEXT, text TEXT)")
|
||
db.execute("INSERT INTO events VALUES (1, 'default', 't', 'user', 'message', 'hello world')")
|
||
db.execute("CREATE TABLE records (id TEXT PRIMARY KEY, kind TEXT, title TEXT, content TEXT, size TEXT, reads INTEGER, created TEXT, updated TEXT)")
|
||
db.execute("INSERT INTO records VALUES ('mem:aaaaaaaaaaaaaaaa', 'note', 't', 'words', '5', 0, '', '')")
|
||
db.commit()
|
||
db.close()
|
||
store = tai.Store(self.config, tai.Seal(self.config.home, ""))
|
||
try:
|
||
self.assertEqual(store.load_secret("srv"), "pw")
|
||
self.assertEqual(store.get_record("mem:aaaaaaaaaaaaaaaa")["content"], "words")
|
||
self.assertEqual(store.get_record("mem:aaaaaaaaaaaaaaaa")["size"], 5)
|
||
self.assertEqual(len(store.search_events("default", "hello")), 1)
|
||
for table in ("secrets", "events", "records", "tags", "edges", "audit", "schedules"):
|
||
self.assertTrue(store.db.execute("SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?", (table,)).fetchone(), table)
|
||
secret_cols = [row[1] for row in store.db.execute("PRAGMA table_info(secrets)").fetchall()]
|
||
event_cols = [row[1] for row in store.db.execute("PRAGMA table_info(events)").fetchall()]
|
||
record_cols = [row[1] for row in store.db.execute("PRAGMA table_info(records)").fetchall()]
|
||
self.assertIn("meta", secret_cols)
|
||
self.assertIn("tags", event_cols)
|
||
self.assertIn("profile", record_cols)
|
||
self.assertIn("added secrets.meta", store.schema_notes)
|
||
self.assertIn("added events.tags", store.schema_notes)
|
||
self.assertIn("added records.profile", store.schema_notes)
|
||
finally:
|
||
store.close()
|
||
|
||
def test_legacy_db_without_meta_migrates(self):
|
||
db = sqlite3.connect(self.config.db_path)
|
||
db.execute("CREATE TABLE secrets (name TEXT PRIMARY KEY, value TEXT, updated TEXT)")
|
||
db.execute("INSERT INTO secrets VALUES ('srv', 'pw-legacy', '2026-01-01')")
|
||
db.commit()
|
||
db.close()
|
||
store = tai.Store(self.config, tai.Seal(self.config.home, ""))
|
||
try:
|
||
self.assertEqual(store.load_secret("srv"), "pw-legacy")
|
||
self.assertEqual(store.secret_meta("srv"), {})
|
||
self.assertEqual(tai.table_pk_columns(store.db, "secrets"), ["profile", "name"])
|
||
names = [row[1] for row in store.db.execute("PRAGMA table_info(secrets)").fetchall()]
|
||
self.assertIn("meta", names)
|
||
store.save_secret("srv", "pw-new", {"host": "example.com"}, None, "default")
|
||
self.assertEqual(store.secret_meta("srv"), {"host": "example.com"})
|
||
finally:
|
||
store.close()
|
||
|
||
|
||
class LazyTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "lazy-test-1"))
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def names(self, text):
|
||
return {schema["function"]["name"] for schema in tai.select_tools(text)}
|
||
|
||
def test_core_minimal(self):
|
||
self.assertEqual(self.names("hello"), set(tai.CORE_TOOLS))
|
||
self.assertEqual(len(tai.CORE_TOOLS), 9)
|
||
|
||
def test_synonym_triggers(self):
|
||
pairs = [
|
||
("remove that file", "delete_file"),
|
||
("set a cron reminder", "schedule"),
|
||
("my password", "store_secret"),
|
||
("undo that change", "restore"),
|
||
("show version history", "audit"),
|
||
("what tags exist", "tags"),
|
||
("publish a release", "release"),
|
||
("search the web", "web_search"),
|
||
("transcribe this", "listen"),
|
||
("check the host specs", "sysinfo"),
|
||
("upcoming agenda", "schedules"),
|
||
("connect these nodes", "graph_link"),
|
||
]
|
||
for text, tool in pairs:
|
||
self.assertIn(tool, self.names(text), "missing %s for %r" % (tool, text))
|
||
|
||
def test_name_mention_loads_family(self):
|
||
found = self.names("use record_save for this")
|
||
for tool in ("record_save", "record_read", "record_search", "record_delete"):
|
||
self.assertIn(tool, found)
|
||
|
||
def test_no_overtrigger(self):
|
||
self.assertEqual(self.names("hello, how are you today"), set(tai.CORE_TOOLS))
|
||
self.assertNotIn("edit_file", self.names("tell me about this"))
|
||
|
||
def test_catalog_lists_lazy(self):
|
||
catalog = tai.tool_catalog()
|
||
self.assertIn("## Tool catalog", catalog)
|
||
for schema in tai.TOOL_SCHEMAS:
|
||
name = schema["function"]["name"]
|
||
if name in tai.CORE_TOOLS:
|
||
continue
|
||
if not tai.tool_env_ok(name):
|
||
continue
|
||
self.assertIn("- %s:" % name, catalog)
|
||
|
||
def test_tool_results_feed_selection(self):
|
||
messages = [
|
||
{"role": "user", "content": "run it"},
|
||
{"role": "assistant", "content": None, "tool_calls": []},
|
||
{"role": "tool", "content": "exit 0\n...[7000 chars spilled to mem:0123456789abcdef, record_read pages the rest]..."},
|
||
]
|
||
text = tai.conversation_text(messages)
|
||
self.assertIn("record_read", text)
|
||
self.assertIn("record_read", self.names(text))
|
||
|
||
def test_run_turn_uses_filtered_payload(self):
|
||
agent = tai.Agent(self.config, self.store, persist=False, quiet=True)
|
||
seen = {}
|
||
reply = {"role": "assistant", "content": "hi", "reasoning": "", "tool_calls": [], "backend": "x"}
|
||
|
||
def fake_complete(messages, tools, stream_sink=None, deadline=None):
|
||
seen["tools"] = {schema["function"]["name"] for schema in tools}
|
||
return dict(reply)
|
||
|
||
with mock.patch.object(agent.chat, "complete", side_effect=fake_complete):
|
||
agent.run_turn("hi", capture=True)
|
||
self.assertEqual(seen["tools"], set(tai.CORE_TOOLS))
|
||
|
||
def test_tools_repl(self):
|
||
agent = mock.Mock()
|
||
with mock.patch("builtins.print") as shown:
|
||
tai.handle_command(agent, "/tools")
|
||
tai.handle_command(agent, "/tools schedule")
|
||
tai.handle_command(agent, "/tools nope")
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list)
|
||
self.assertIn("core, always loaded", printed)
|
||
self.assertIn("schedule", printed)
|
||
self.assertIn("cron", printed)
|
||
self.assertIn("unknown tool", printed)
|
||
|
||
|
||
class InstallOpsTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.home = os.path.join(self.tmp.name, "home")
|
||
os.makedirs(self.home)
|
||
self.env = mock.patch.dict(os.environ, {"HOME": self.home, "TAI_HOME": os.path.join(self.home, ".tai")})
|
||
self.env.start()
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "install-test-1"))
|
||
self.app = FakeSecretApp(self.store)
|
||
self.app.config = self.config
|
||
self.tools = tai.Tools(self.app)
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.env.stop()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
tai._BOX_PYTHON_CACHE.clear()
|
||
|
||
def test_status_reports_all_targets(self):
|
||
with mock.patch.object(tai.Sandbox, "container_engine", return_value=None):
|
||
report = tai.install_report("status", list(tai.INSTALL_TARGETS), self.config.home, self.store)
|
||
for label in (
|
||
"binary:", "bash-hook:", "zsh-hook:", "fish-hook:", "plexer-hook:", "zsh-plexer-hook:",
|
||
"fish-plexer-hook:", "venv:", "scheduler-service:", "telegram-service:", "container:", "vault:",
|
||
):
|
||
self.assertIn(label, report)
|
||
self.assertIn("not installed", report)
|
||
self.assertIn("never touch data", report)
|
||
|
||
def test_binary_install_upgrade_uninstall(self):
|
||
target = os.path.join(self.home, ".local", "bin", "tai.py")
|
||
self.assertIn("installed", tai.op_binary("install", self.config.home))
|
||
self.assertTrue(os.access(target, os.X_OK))
|
||
with open(os.path.abspath(tai.__file__), "rb") as handle:
|
||
self.assertEqual(open(target, "rb").read(), handle.read())
|
||
self.assertIn("already installed", tai.op_binary("install", self.config.home))
|
||
self.assertIn("refreshed", tai.op_binary("upgrade", self.config.home))
|
||
self.assertIn("removed", tai.op_binary("uninstall", self.config.home))
|
||
self.assertFalse(os.path.exists(target))
|
||
self.assertIn("not present", tai.op_binary("uninstall", self.config.home))
|
||
|
||
def test_hook_install_remove_idempotent(self):
|
||
path = os.path.join(self.home, ".bashrc")
|
||
self.assertIn("source ~/.bashrc", tai.op_hook("install", self.config.home))
|
||
with open(path, encoding="utf-8") as handle:
|
||
self.assertIn(tai.BASHRC_MARK_BEGIN, handle.read())
|
||
self.assertIn("already present", tai.op_hook("install", self.config.home))
|
||
self.assertIn("removed", tai.op_hook("uninstall", self.config.home))
|
||
with open(path, encoding="utf-8") as handle:
|
||
self.assertNotIn(tai.BASHRC_MARK_BEGIN, handle.read())
|
||
self.assertTrue(os.path.isfile(path + ".bak-tai"))
|
||
self.assertIn("not present", tai.op_hook("uninstall", self.config.home))
|
||
|
||
def test_zsh_hook_install_remove_idempotent(self):
|
||
path = os.path.join(self.home, ".zshrc")
|
||
with mock.patch.object(tai.Install, "shell_present", return_value=True):
|
||
self.assertIn("source ~/.zshrc", tai.op_zsh_hook("install", self.config.home))
|
||
with open(path, encoding="utf-8") as handle:
|
||
content = handle.read()
|
||
self.assertIn(tai.BASHRC_MARK_BEGIN, content)
|
||
self.assertIn("command_not_found_handler", content)
|
||
self.assertNotIn("command_not_found_handle()", content)
|
||
self.assertIn("already present", tai.op_zsh_hook("install", self.config.home))
|
||
self.assertIn("removed", tai.op_zsh_hook("uninstall", self.config.home))
|
||
with open(path, encoding="utf-8") as handle:
|
||
self.assertNotIn(tai.BASHRC_MARK_BEGIN, handle.read())
|
||
self.assertTrue(os.path.isfile(path + ".bak-tai"))
|
||
self.assertIn("not present", tai.op_zsh_hook("uninstall", self.config.home))
|
||
|
||
def test_fish_hook_install_remove_idempotent(self):
|
||
path = os.path.join(self.home, ".config", "fish", "config.fish")
|
||
with mock.patch.object(tai.Install, "shell_present", return_value=True):
|
||
result = tai.op_fish_hook("install", self.config.home)
|
||
self.assertIn("installed", result)
|
||
with open(path, encoding="utf-8") as handle:
|
||
content = handle.read()
|
||
self.assertIn(tai.BASHRC_MARK_BEGIN, content)
|
||
self.assertIn("function fish_command_not_found", content)
|
||
self.assertIn("already present", tai.op_fish_hook("install", self.config.home))
|
||
self.assertIn("removed", tai.op_fish_hook("uninstall", self.config.home))
|
||
with open(path, encoding="utf-8") as handle:
|
||
self.assertNotIn(tai.BASHRC_MARK_BEGIN, handle.read())
|
||
self.assertTrue(os.path.isfile(path + ".bak-tai"))
|
||
|
||
def test_shell_hook_skips_loudly_when_shell_missing(self):
|
||
with mock.patch.object(tai.Install, "shell_present", return_value=False):
|
||
self.assertEqual(tai.op_zsh_hook("install", self.config.home), "zsh-hook: zsh not found, skipped")
|
||
self.assertEqual(tai.op_fish_hook("upgrade", self.config.home), "fish-hook: fish not found, skipped")
|
||
self.assertEqual(
|
||
tai.op_zsh_plexer_hook("install", self.config.home), "zsh-plexer-hook: zsh not found, skipped",
|
||
)
|
||
self.assertEqual(
|
||
tai.op_fish_plexer_hook("install", self.config.home), "fish-plexer-hook: fish not found, skipped",
|
||
)
|
||
self.assertFalse(os.path.exists(os.path.join(self.home, ".zshrc")))
|
||
self.assertFalse(os.path.exists(os.path.join(self.home, ".config", "fish", "config.fish")))
|
||
|
||
def test_shell_hook_uninstall_cleans_up_even_if_shell_missing_now(self):
|
||
with mock.patch.object(tai.Install, "shell_present", return_value=True):
|
||
tai.op_zsh_hook("install", self.config.home)
|
||
with mock.patch.object(tai.Install, "shell_present", return_value=False):
|
||
self.assertIn("removed", tai.op_zsh_hook("uninstall", self.config.home))
|
||
|
||
def test_zsh_plexer_hook_reuses_bash_syntax_different_file(self):
|
||
bash_path = os.path.join(self.home, ".bashrc")
|
||
zsh_path = os.path.join(self.home, ".zshrc")
|
||
with mock.patch.object(tai.Install, "shell_present", return_value=True):
|
||
tai.op_plexer_hook("install", self.config.home)
|
||
tai.op_zsh_plexer_hook("install", self.config.home)
|
||
with open(bash_path, encoding="utf-8") as handle:
|
||
bash_content = handle.read()
|
||
with open(zsh_path, encoding="utf-8") as handle:
|
||
zsh_content = handle.read()
|
||
self.assertEqual(bash_content, zsh_content)
|
||
self.assertIn('case "$-" in', zsh_content)
|
||
|
||
def test_fish_plexer_hook_uses_fish_syntax(self):
|
||
path = os.path.join(self.home, ".config", "fish", "config.fish")
|
||
with mock.patch.object(tai.Install, "shell_present", return_value=True):
|
||
self.assertIn("installed", tai.op_fish_plexer_hook("install", self.config.home))
|
||
with open(path, encoding="utf-8") as handle:
|
||
content = handle.read()
|
||
self.assertIn("status is-interactive", content)
|
||
self.assertIn("isatty stdin", content)
|
||
self.assertIn("isatty stdout", content)
|
||
self.assertIn("exec \"$HOME/.local/bin/tai.py\" --plexer-wrap", content)
|
||
self.assertNotIn("case \"$-\"", content)
|
||
self.assertNotIn("[ -z \"$TAI_PLEXER_ACTIVE\" ]", content)
|
||
|
||
def test_install_self_detects_current_shell(self):
|
||
with mock.patch.object(tai.Install, "shell_present", return_value=True):
|
||
with mock.patch.dict(os.environ, {"SHELL": "/usr/bin/zsh"}):
|
||
with mock.patch("builtins.print") as shown:
|
||
tai.install_self()
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list)
|
||
self.assertIn("zsh-hook:", printed)
|
||
self.assertIn("source ~/.zshrc", printed)
|
||
|
||
def test_venv_kept_and_removed(self):
|
||
folder = tai.venv_dir(self.config.home)
|
||
python = os.path.join(folder, "bin", "python")
|
||
os.makedirs(os.path.dirname(python))
|
||
with open(python, "w", encoding="utf-8") as handle:
|
||
handle.write("#!/bin/sh\n")
|
||
os.chmod(python, 0o755)
|
||
self.assertIn("kept", tai.op_venv("install", self.config.home))
|
||
self.assertIn(python, tai.service_exec(self.config.home, "/x/tai.py"))
|
||
self.assertIn("removed", tai.op_venv("uninstall", self.config.home))
|
||
self.assertFalse(os.path.isdir(folder))
|
||
self.assertEqual(tai.service_exec(self.config.home, "/x/tai.py"), "/x/tai.py")
|
||
|
||
def test_venv_create_and_unavailable(self):
|
||
with mock.patch.object(tai.Install, "venv_available", return_value=False):
|
||
self.assertIn("unavailable", tai.op_venv("install", self.config.home))
|
||
|
||
def fake_run(argv, stdout=None, stderr=None, universal_newlines=True, timeout=300):
|
||
folder = argv[-1]
|
||
python = os.path.join(folder, "bin", "python")
|
||
os.makedirs(os.path.dirname(python))
|
||
with open(python, "w", encoding="utf-8") as handle:
|
||
handle.write("#!/bin/sh\n")
|
||
os.chmod(python, 0o755)
|
||
outcome = mock.Mock()
|
||
outcome.returncode = 0
|
||
outcome.stderr = ""
|
||
return outcome
|
||
|
||
with mock.patch.object(tai.Install, "venv_available", return_value=True):
|
||
with mock.patch("subprocess.run", side_effect=fake_run):
|
||
self.assertIn("created", tai.op_venv("install", self.config.home))
|
||
|
||
def test_scheduler_unit_write_and_remove(self):
|
||
with mock.patch("shutil.which", return_value=None):
|
||
self.assertIn("unit written", tai.op_scheduler_service("install", self.config.home))
|
||
unit = os.path.join(self.home, ".config", "systemd", "user", "tai-scheduler.service")
|
||
with open(unit, encoding="utf-8") as handle:
|
||
body = handle.read()
|
||
self.assertIn("--scheduler", body)
|
||
self.assertIn("already installed", tai.op_scheduler_service("install", self.config.home))
|
||
self.assertIn("removed", tai.op_scheduler_service("uninstall", self.config.home))
|
||
self.assertFalse(os.path.exists(unit))
|
||
|
||
def test_telegram_needs_token(self):
|
||
with mock.patch.dict(os.environ, {}, clear=False):
|
||
os.environ.pop("TELEGRAM_BOT_TOKEN", None)
|
||
self.assertIn("TELEGRAM_BOT_TOKEN", tai.op_telegram_service("install", self.config.home))
|
||
self.assertIn("not present", tai.op_telegram_service("uninstall", self.config.home))
|
||
|
||
def test_reinstall_keeps_vault(self):
|
||
marker = b"vault-data-marker-7"
|
||
with open(os.path.join(self.config.home, "memory.db"), "wb") as handle:
|
||
handle.write(marker)
|
||
with mock.patch.object(tai.Install, "ensure_venv", return_value=(True, "created fake")):
|
||
report = tai.install_report("reinstall", ["binary", "bash-hook", "venv"], self.config.home, self.store)
|
||
with open(os.path.join(self.config.home, "memory.db"), "rb") as handle:
|
||
self.assertIn(marker, handle.read())
|
||
self.assertIn("never touch data", report)
|
||
|
||
def test_install_tool_validation_and_approval(self):
|
||
self.assertIn("action must be", self.tools.dispatch("install", json.dumps({"action": "explode"})))
|
||
self.assertIn("unknown target", self.tools.dispatch("install", json.dumps({"action": "status", "targets": ["nope"]})))
|
||
with mock.patch.object(tai.Sandbox, "container_engine", return_value=None):
|
||
status = self.tools.dispatch("install", json.dumps({"action": "status"}))
|
||
self.assertIn("scheduler-service:", status)
|
||
self.app.approve = False
|
||
self.assertIn("denied by user", self.tools.dispatch("install", json.dumps({"action": "upgrade", "targets": ["venv"]})))
|
||
self.assertIn("vault always kept", self.app.approvals[0])
|
||
self.app.approve = True
|
||
done = self.tools.dispatch("install", json.dumps({"action": "install", "targets": ["bash-hook"]}))
|
||
self.assertIn("bash-hook:", done)
|
||
|
||
def test_install_mentions_load_tool(self):
|
||
found = {schema["function"]["name"] for schema in tai.select_tools("how do I install tai as a service")}
|
||
self.assertIn("install", found)
|
||
found = {schema["function"]["name"] for schema in tai.select_tools("set up a venv hook")}
|
||
self.assertIn("install", found)
|
||
|
||
def test_box_python_cache_and_containerfile(self):
|
||
self.assertIn("/box/venv", tai.BOX_CONTAINERFILE)
|
||
self.assertIn("python3-venv", tai.BOX_CONTAINERFILE)
|
||
probe = mock.Mock()
|
||
probe.returncode = 0
|
||
with mock.patch.object(tai.Sandbox, "box_exec", return_value=probe) as runner:
|
||
self.assertEqual(tai.box_python("podman"), tai.BOX_PYTHON)
|
||
self.assertEqual(tai.box_python("podman"), tai.BOX_PYTHON)
|
||
self.assertEqual(runner.call_count, 1)
|
||
tai._BOX_PYTHON_CACHE.clear()
|
||
probe.returncode = 1
|
||
with mock.patch.object(tai.Sandbox, "box_exec", return_value=probe):
|
||
self.assertEqual(tai.box_python("podman"), "python3")
|
||
|
||
def test_unit_templates_format(self):
|
||
body = tai.TELEGRAM_UNIT % ("/x/tai.py", "/x/telegram.env")
|
||
self.assertIn("WorkingDirectory=%h", body)
|
||
self.assertIn("/x/tai.py --telegram", body)
|
||
body = tai.SCHEDULER_UNIT % "/x/tai.py"
|
||
self.assertIn("/x/tai.py --scheduler", body)
|
||
|
||
def test_install_repl(self):
|
||
self.app.tools = self.tools
|
||
with mock.patch.object(tai.Sandbox, "container_engine", return_value=None):
|
||
with mock.patch("builtins.print") as shown:
|
||
tai.handle_command(self.app, "/install status")
|
||
tai.handle_command(self.app, "/install frobnicate")
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list)
|
||
self.assertIn("bash-hook:", printed)
|
||
self.assertIn("use /install", printed)
|
||
|
||
|
||
class BotsTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "bots-test-1"))
|
||
self.agent = tai.Agent(self.config, self.store, persist=True, quiet=True)
|
||
self.tools = self.agent.tools
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def test_create_bot_roundtrip(self):
|
||
result = self.tools.dispatch("create_bot", json.dumps({"name": "Helper-Bot", "description": "helps a lot", "rules": "be kind", "behavior": "concise", "nicknames": ["help"]}))
|
||
self.assertIn("created bot 'helper-bot'", result)
|
||
self.assertIn("help", result)
|
||
self.assertIn("helper", result)
|
||
bots = {item["name"]: item for item in self.store.list_bots("t")}
|
||
self.assertIn("helper-bot", bots)
|
||
self.assertIn("help", bots["helper-bot"]["nicknames"])
|
||
self.assertIn("helper", bots["helper-bot"]["nicknames"])
|
||
system = self.store.load_bot_system("t", "helper-bot")
|
||
self.assertIn("helps a lot", system)
|
||
self.assertIn("Rules:", system)
|
||
self.assertIn("Behavior:", system)
|
||
self.assertEqual(self.store.resolve_bot("t", "helper-bot"), "helper-bot")
|
||
self.assertEqual(self.store.resolve_bot("t", "HELP"), "helper-bot")
|
||
self.assertEqual(self.store.resolve_bot("t", "main"), "main")
|
||
self.assertIsNone(self.store.resolve_bot("t", "nope"))
|
||
|
||
def test_create_bot_validation(self):
|
||
self.assertIn("invalid bot name", self.tools.dispatch("create_bot", json.dumps({"name": "Bad Name!"})))
|
||
self.assertIn("default bot", self.tools.dispatch("create_bot", json.dumps({"name": "main", "description": "x"})))
|
||
self.assertIn("at least one", self.tools.dispatch("create_bot", json.dumps({"name": "empty"})))
|
||
self.assertIn("nicknames must be", self.tools.dispatch("create_bot", json.dumps({"name": "x", "description": "y", "nicknames": "nope"})))
|
||
self.tools.dispatch("create_bot", json.dumps({"name": "coder", "description": "writes code", "nicknames": ["cd"]}))
|
||
self.assertIn("taken", self.tools.dispatch("create_bot", json.dumps({"name": "coder", "description": "again"})))
|
||
self.assertIn("taken", self.tools.dispatch("create_bot", json.dumps({"name": "other", "description": "y", "nicknames": ["cd"]})))
|
||
self.assertIn("invalid nickname", self.tools.dispatch("create_bot", json.dumps({"name": "ok", "description": "y", "nicknames": ["bad nick!"]})))
|
||
|
||
def test_switch_bot_resumes(self):
|
||
self.tools.dispatch("create_bot", json.dumps({"name": "coder", "description": "writes code", "nicknames": ["cd"]}))
|
||
self.agent.messages.append({"role": "user", "content": "main hello"})
|
||
self.agent.messages.append({"role": "assistant", "content": "main hi"})
|
||
with mock.patch("builtins.print"):
|
||
self.assertIn("switched", self.agent.switch_bot("coder"))
|
||
self.assertEqual(self.agent.bot, "coder")
|
||
self.assertEqual([item["role"] for item in self.agent.messages], ["system"])
|
||
self.assertIn("writes code", self.agent.system_message)
|
||
self.agent.messages.append({"role": "user", "content": "coder hello"})
|
||
with mock.patch("builtins.print"):
|
||
self.agent.switch_bot("main")
|
||
contents = [item.get("content", "") for item in self.agent.messages]
|
||
self.assertIn("main hello", contents)
|
||
self.assertNotIn("coder hello", contents)
|
||
with mock.patch("builtins.print"):
|
||
self.assertIn("switched", self.agent.switch_bot("cd"))
|
||
self.assertIn("coder hello", [item.get("content", "") for item in self.agent.messages])
|
||
|
||
def test_mention_routes_and_crossposts(self):
|
||
self.tools.dispatch("create_bot", json.dumps({"name": "coder", "description": "coder-brain-9"}))
|
||
seen = []
|
||
reply = {"role": "assistant", "content": "coded-it", "reasoning": "", "tool_calls": [], "backend": "x"}
|
||
|
||
def fake_complete(messages, tools, stream_sink=None, deadline=None):
|
||
seen.append(messages[0]["content"])
|
||
return dict(reply)
|
||
|
||
with mock.patch.object(self.agent.chat, "complete", side_effect=fake_complete):
|
||
answer = self.agent.run_turn("@coder write frob", capture=True)
|
||
self.assertEqual(answer, "coded-it")
|
||
self.assertIn("coder-brain-9", seen[0])
|
||
self.assertEqual(self.agent.bot, "main")
|
||
tail = self.agent.messages[-2:]
|
||
self.assertEqual(tail[0]["content"], "@coder write frob")
|
||
self.assertEqual(tail[1]["content"], "coded-it")
|
||
restored = self.store.load_session("t", "coder")
|
||
texts = [item.get("content", "") for item in restored]
|
||
self.assertIn("write frob", texts)
|
||
self.assertIn("coded-it", texts)
|
||
|
||
def test_mention_unknown_and_self(self):
|
||
reply = {"role": "assistant", "content": "hi", "reasoning": "", "tool_calls": [], "backend": "x"}
|
||
with mock.patch.object(self.agent.chat, "complete", return_value=dict(reply)):
|
||
self.assertIn("unknown bot", self.agent.run_turn("@ghost hi", capture=True))
|
||
self.assertEqual(self.agent.run_turn("@main hi", capture=True), "hi")
|
||
self.assertEqual(self.agent.messages[-2]["content"], "hi")
|
||
|
||
def test_remember_goes_to_current_bot(self):
|
||
self.tools.dispatch("create_bot", json.dumps({"name": "coder", "description": "v1"}))
|
||
with mock.patch("builtins.print"):
|
||
self.agent.switch_bot("coder")
|
||
self.agent.update_system("coder rules v2")
|
||
self.assertIn("v2", self.store.load_bot_system("t", "coder"))
|
||
self.assertNotIn("v2", self.store.load_system("t"))
|
||
|
||
def test_profile_switch_resets_bot(self):
|
||
self.tools.dispatch("create_bot", json.dumps({"name": "coder", "description": "v1"}))
|
||
with mock.patch("builtins.print"):
|
||
self.agent.switch_bot("coder")
|
||
self.agent.switch_profile("other")
|
||
self.assertEqual(self.agent.bot, "main")
|
||
|
||
def test_bots_repl(self):
|
||
self.tools.dispatch("create_bot", json.dumps({"name": "coder", "description": "v1", "nicknames": ["cd"]}))
|
||
with mock.patch("builtins.print") as shown:
|
||
tai.handle_command(self.agent, "/bots")
|
||
tai.handle_command(self.agent, "/bot coder")
|
||
tai.handle_command(self.agent, "/bot")
|
||
tai.handle_command(self.agent, "/bot nope")
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list)
|
||
self.assertIn("* main", printed)
|
||
self.assertIn("coder aka cd", printed)
|
||
self.assertIn("switched to bot", printed)
|
||
self.assertIn("use /bot <name>", printed)
|
||
self.assertIn("unknown bot", printed)
|
||
self.assertEqual(self.agent.bot, "coder")
|
||
|
||
|
||
class MentionTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
self.workdir = tempfile.TemporaryDirectory()
|
||
self.saved_cwd = os.getcwd()
|
||
os.chdir(self.workdir.name)
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "mention-test-1"))
|
||
self.agent = tai.Agent(self.config, self.store, persist=True, quiet=True)
|
||
|
||
def tearDown(self):
|
||
os.chdir(self.saved_cwd)
|
||
self.store.close()
|
||
self.workdir.cleanup()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def fake_chat(self, seen):
|
||
reply = {"role": "assistant", "content": "done-1", "reasoning": "", "tool_calls": [], "backend": "x"}
|
||
|
||
def fake_complete(messages, tools, stream_sink=None, deadline=None):
|
||
seen.append([item.get("content", "") for item in messages if item.get("role") == "user"])
|
||
seen.append(self.agent.chat.provider_pin)
|
||
return dict(reply)
|
||
|
||
return mock.patch.object(self.agent.chat, "complete", side_effect=fake_complete)
|
||
|
||
def tools_dispatch_create_bot(self):
|
||
self.agent.tools.dispatch("create_bot", json.dumps({"name": "coder", "description": "writes code"}))
|
||
|
||
def test_split_leading_tags(self):
|
||
self.assertEqual(tai.split_leading_tags("@a @b rest"), (["a", "b"], "rest"))
|
||
self.assertEqual(tai.split_leading_tags("@a"), (["a"], ""))
|
||
self.assertEqual(tai.split_leading_tags("hi @a"), ([], "hi @a"))
|
||
self.assertEqual(tai.split_leading_tags("@"), ([], "@"))
|
||
self.assertEqual(tai.split_leading_tags(""), ([], ""))
|
||
|
||
def test_resolve_precedence(self):
|
||
self.tools_dispatch_create_bot()
|
||
with open("coder", "w") as handle:
|
||
handle.write("x")
|
||
with open("note.txt", "w") as handle:
|
||
handle.write("x")
|
||
self.assertEqual(tai.resolve_tag("coder", self.agent, os.getcwd()), ("bot", "coder"))
|
||
self.assertEqual(tai.resolve_tag("opencode", self.agent, os.getcwd()), ("provider", "opencode"))
|
||
self.assertEqual(tai.resolve_tag("mimo-v2.6-flash-free", self.agent, os.getcwd()), ("model", "mimo-v2.6-flash-free"))
|
||
self.assertEqual(tai.resolve_tag("auto", self.agent, os.getcwd()), ("auto", "auto"))
|
||
kind, value = tai.resolve_tag("note.txt", self.agent, os.getcwd())
|
||
self.assertEqual(kind, "file")
|
||
self.assertTrue(value.endswith("note.txt"))
|
||
self.assertEqual(tai.resolve_tag("nope-zzz", self.agent, os.getcwd()), ("unknown", "nope-zzz"))
|
||
|
||
def test_combined_tags_route_and_pin(self):
|
||
self.tools_dispatch_create_bot()
|
||
seen = []
|
||
with self.fake_chat(seen):
|
||
answer = self.agent.run_turn("@opencode @coder write frob", capture=True)
|
||
self.assertEqual(answer, "done-1")
|
||
self.assertEqual(self.agent.chat.provider_pin, "opencode")
|
||
self.assertEqual(self.agent.bot, "main")
|
||
self.assertEqual(seen[1], "opencode")
|
||
saved = self.store.load_dir_defaults(tai.dir_key())
|
||
self.assertEqual(saved, {"provider": "opencode", "bot": "coder"})
|
||
|
||
def test_last_tag_wins(self):
|
||
answer = self.agent.run_turn("@ollama @opencode", capture=True)
|
||
self.assertIn("opencode", answer)
|
||
self.assertEqual(self.agent.chat.provider_pin, "opencode")
|
||
|
||
def test_unknown_tag_errors(self):
|
||
answer = self.agent.run_turn("@ghost-zzz hi", capture=True)
|
||
self.assertIn("unknown bot", answer)
|
||
self.assertIn("@ghost-zzz", answer)
|
||
self.assertIsNone(self.agent.chat.provider_pin)
|
||
self.assertEqual(self.store.load_dir_defaults(tai.dir_key()), {})
|
||
|
||
def test_bare_selectors_confirm_and_save(self):
|
||
seen = []
|
||
with self.fake_chat(seen):
|
||
answer = self.agent.run_turn("@opencode", capture=True)
|
||
self.assertIn("provider pinned to opencode", answer)
|
||
self.assertIn("saved as default", answer)
|
||
self.assertEqual(seen, [])
|
||
self.assertEqual(self.store.load_dir_defaults(tai.dir_key()), {"provider": "opencode"})
|
||
|
||
def test_bare_bot_switches_sticky(self):
|
||
self.tools_dispatch_create_bot()
|
||
answer = self.agent.run_turn("@coder", capture=True)
|
||
self.assertIn("switched to bot", answer)
|
||
self.assertEqual(self.agent.bot, "coder")
|
||
self.assertEqual(self.store.load_dir_defaults(tai.dir_key()), {"bot": "coder"})
|
||
|
||
def test_auto_clears(self):
|
||
self.agent.chat.provider_pin = "opencode"
|
||
self.agent.chat.model_pin = "mimo-v2.6-flash-free"
|
||
self.store.save_dir_defaults(tai.dir_key(), provider="opencode", model="mimo-v2.6-flash-free", bot="main")
|
||
answer = self.agent.run_turn("@auto", capture=True)
|
||
self.assertIn("pins cleared", answer)
|
||
self.assertIsNone(self.agent.chat.provider_pin)
|
||
self.assertIsNone(self.agent.chat.model_pin)
|
||
self.assertEqual(self.store.load_dir_defaults(tai.dir_key()), {"bot": "main"})
|
||
|
||
def test_model_tag_pins(self):
|
||
seen = []
|
||
with self.fake_chat(seen):
|
||
self.agent.run_turn("@mimo-v2.6-flash-free hi", capture=True)
|
||
self.assertEqual(self.agent.chat.model_pin, "mimo-v2.6-flash-free")
|
||
self.assertEqual(self.store.load_dir_defaults(tai.dir_key()), {"model": "mimo-v2.6-flash-free"})
|
||
|
||
def test_midtext_file_expands(self):
|
||
with open("doc.txt", "w") as handle:
|
||
handle.write("x")
|
||
seen = []
|
||
with self.fake_chat(seen):
|
||
self.agent.run_turn("explain @doc.txt briefly", capture=True)
|
||
user_text = seen[0][-1]
|
||
self.assertIn(os.path.realpath("doc.txt"), user_text)
|
||
self.assertNotIn("@doc.txt", user_text)
|
||
|
||
def test_midtext_prose_untouched(self):
|
||
seen = []
|
||
with self.fake_chat(seen):
|
||
self.agent.run_turn("ping me @home base", capture=True)
|
||
self.assertEqual(seen[0][-1], "ping me @home base")
|
||
|
||
def test_missing_midtext_path_stays_literal(self):
|
||
seen = []
|
||
with self.fake_chat(seen):
|
||
self.agent.run_turn("read @nope-missing-1.txt now", capture=True)
|
||
self.assertIn("@nope-missing-1.txt", seen[0][-1])
|
||
|
||
def test_leading_file_tag(self):
|
||
with open("doc.txt", "w") as handle:
|
||
handle.write("x")
|
||
seen = []
|
||
with self.fake_chat(seen):
|
||
self.agent.run_turn("@doc.txt summarize", capture=True)
|
||
self.assertTrue(seen[0][-1].startswith(os.path.realpath("doc.txt")))
|
||
|
||
def test_pins_filter_candidates(self):
|
||
cands = [
|
||
{"provider": "opencode", "chat_url": "opencode://run", "models_url": "x", "key": "", "model": "a", "profile": ""},
|
||
{"provider": "ollama", "chat_url": "http://x", "models_url": "y", "key": "", "model": "b", "profile": ""},
|
||
]
|
||
with mock.patch.object(self.agent.chat, "candidates", return_value=cands):
|
||
self.agent.chat.provider_pin = "opencode"
|
||
got = self.agent.chat.ordered_candidates()
|
||
self.assertEqual([cand["provider"] for cand in got], ["opencode"])
|
||
self.agent.chat.provider_pin = None
|
||
self.agent.chat.model_pin = "b"
|
||
got = self.agent.chat.all_candidates()
|
||
self.assertEqual([cand["model"] for cand in got], ["b"])
|
||
|
||
def test_empty_candidates_raise(self):
|
||
cands = [{"provider": "opencode", "chat_url": "opencode://run", "models_url": "x", "key": "", "model": "a", "profile": ""}]
|
||
with mock.patch.object(self.agent.chat, "candidates", return_value=cands):
|
||
self.agent.chat.provider_pin = "no-such-provider"
|
||
with self.assertRaises(tai.BackendError) as raised:
|
||
self.agent.chat.complete([{"role": "user", "content": "hi"}])
|
||
self.assertIn("no-such-provider", str(raised.exception))
|
||
|
||
def test_dir_defaults_roundtrip(self):
|
||
self.store.save_dir_defaults("d1", provider="opencode")
|
||
self.store.save_dir_defaults("d1", bot="coder")
|
||
self.assertEqual(self.store.load_dir_defaults("d1"), {"provider": "opencode", "bot": "coder"})
|
||
self.store.save_dir_defaults("d1", clear=("provider",))
|
||
self.assertEqual(self.store.load_dir_defaults("d1"), {"bot": "coder"})
|
||
self.store.save_dir_defaults("d1", clear=("bot",))
|
||
self.assertEqual(self.store.load_dir_defaults("d1"), {})
|
||
|
||
def test_apply_dir_defaults(self):
|
||
self.tools_dispatch_create_bot()
|
||
self.store.save_dir_defaults(tai.dir_key(), provider="opencode", model="mimo-v2.6-flash-free", bot="coder")
|
||
note = tai.apply_dir_defaults(self.agent)
|
||
self.assertIn("dir defaults", note)
|
||
self.assertEqual(self.agent.chat.provider_pin, "opencode")
|
||
self.assertEqual(self.agent.chat.model_pin, "mimo-v2.6-flash-free")
|
||
self.assertEqual(self.agent.bot, "coder")
|
||
|
||
def test_apply_dir_defaults_skips_unknown_bot(self):
|
||
self.store.save_dir_defaults(tai.dir_key(), bot="ghost-zzz")
|
||
note = tai.apply_dir_defaults(self.agent)
|
||
self.assertEqual(note, "")
|
||
self.assertEqual(self.agent.bot, "main")
|
||
|
||
def test_apply_dir_defaults_env_model_wins(self):
|
||
self.config.explicit_model = "env-model-1"
|
||
self.store.save_dir_defaults(tai.dir_key(), model="mimo-v2.6-flash-free")
|
||
tai.apply_dir_defaults(self.agent)
|
||
self.assertIsNone(self.agent.chat.model_pin)
|
||
|
||
def test_complete_mention(self):
|
||
self.tools_dispatch_create_bot()
|
||
with open("doc.txt", "w") as handle:
|
||
handle.write("x")
|
||
self.assertEqual(tai.complete_mention("@op", 0, self.agent), "@opencode")
|
||
self.assertEqual(tai.complete_mention("@cod", 0, self.agent), "@coder")
|
||
self.assertEqual(tai.complete_mention("@doc", 0, self.agent), "@doc.txt")
|
||
self.assertIsNone(tai.complete_mention("@doc", 5, self.agent))
|
||
self.assertIsNone(tai.complete_mention("xx", 0, self.agent))
|
||
self.assertEqual(tai.complete_input("/he", 0, None), "/help")
|
||
self.assertEqual(tai.complete_input("@op", 0, self.agent), "@opencode")
|
||
|
||
def test_loud_paths_print(self):
|
||
self.agent.quiet = False
|
||
with mock.patch("builtins.print") as shown:
|
||
self.agent.run_turn("@ghost-zzz hi", capture=False)
|
||
self.agent.run_turn("@opencode", capture=False)
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list if call.args)
|
||
self.assertIn("unknown bot @ghost-zzz", printed)
|
||
self.assertIn("provider pinned to opencode", printed)
|
||
|
||
def test_complete_files(self):
|
||
os.makedirs("src/nest", exist_ok=True)
|
||
with open("src/main.py", "w") as handle:
|
||
handle.write("x")
|
||
with open(".hidden", "w") as handle:
|
||
handle.write("x")
|
||
self.assertEqual(tai.complete_files("src/ma", os.getcwd()), ["src/main.py"])
|
||
self.assertEqual(tai.complete_files("src/", os.getcwd()), ["src/main.py", "src/nest/"])
|
||
self.assertNotIn(".hidden", tai.complete_files("", os.getcwd()))
|
||
self.assertIn(".hidden", tai.complete_files(".", os.getcwd()))
|
||
|
||
|
||
class OrientationTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
self.workdir = tempfile.TemporaryDirectory()
|
||
self.saved_cwd = os.getcwd()
|
||
os.chdir(self.workdir.name)
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "orient-test-1"))
|
||
self.agent = tai.Agent(self.config, self.store, persist=True, quiet=True)
|
||
|
||
def tearDown(self):
|
||
os.chdir(self.saved_cwd)
|
||
self.store.close()
|
||
self.workdir.cleanup()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def test_block_states_cwd_and_scope(self):
|
||
block = tai.orientation_block()
|
||
self.assertIn(os.path.realpath(os.getcwd()), block)
|
||
self.assertIn("Stay inside it unless the user names another path", block)
|
||
self.assertIn("scratch or temp dir", block)
|
||
|
||
def test_block_sandbox_variant(self):
|
||
block = tai.orientation_block(sandbox=True)
|
||
self.assertIn("Sandboxed container", block)
|
||
self.assertNotIn(os.path.realpath(os.getcwd()), block)
|
||
|
||
def test_apply_system_injects_block(self):
|
||
system = self.agent.messages[0]["content"]
|
||
self.assertIn("## Working directory", system)
|
||
self.assertIn(os.path.realpath(os.getcwd()), system)
|
||
|
||
def test_reset_history_keeps_block(self):
|
||
self.agent.messages.append({"role": "user", "content": "hi"})
|
||
self.agent.reset_history()
|
||
self.assertEqual(len(self.agent.messages), 1)
|
||
self.assertIn("## Working directory", self.agent.messages[0]["content"])
|
||
|
||
def test_muse_prompt_carries_block(self):
|
||
prompt = tai.muse_prompt(self.agent.messages, [])
|
||
self.assertIn(os.path.realpath(os.getcwd()), prompt)
|
||
|
||
def test_muse_prompt_forbids_harness_tools(self):
|
||
prompt = tai.muse_prompt([{"role": "user", "content": "hi"}], [])
|
||
self.assertIn("never invoke assistant or harness tools yourself", prompt)
|
||
|
||
|
||
class CapabilityTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "cap-test-1"))
|
||
self.agent = tai.Agent(self.config, self.store, persist=True, quiet=True)
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def test_yolo_variant(self):
|
||
block = tai.capability_block(yolo=True)
|
||
self.assertIn("Approvals are off", block)
|
||
self.assertIn("not read-only", block)
|
||
self.assertIn("only a failed tool call proves a limit", block)
|
||
|
||
def test_default_variant(self):
|
||
block = tai.capability_block()
|
||
self.assertIn("ask for approval", block)
|
||
self.assertIn("not read-only", block)
|
||
|
||
def test_apply_system_injects_capabilities(self):
|
||
system = self.agent.messages[0]["content"]
|
||
self.assertIn("## Session capabilities", system)
|
||
self.assertIn("ask for approval", system)
|
||
self.agent.yolo = True
|
||
self.agent.apply_system()
|
||
self.assertIn("Approvals are off", self.agent.messages[0]["content"])
|
||
|
||
def test_run_turn_refreshes_capabilities(self):
|
||
self.agent.messages[0] = {"role": "system", "content": "stale"}
|
||
reply = {"role": "assistant", "content": "done", "reasoning": "", "tool_calls": [], "backend": "x"}
|
||
with mock.patch.object(self.agent.chat, "complete", return_value=dict(reply)):
|
||
self.agent.run_turn("hi", capture=True)
|
||
system = self.agent.messages[0]["content"]
|
||
self.assertIn("## Session capabilities", system)
|
||
self.assertNotIn("stale", system)
|
||
|
||
|
||
class FTSSearchTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "fts-test-1"))
|
||
self.app = FakeSecretApp(self.store)
|
||
self.app.profile = "t"
|
||
self.app.config = self.config
|
||
self.tools = tai.Tools(self.app)
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def mem_id(self, text):
|
||
return re.search(r"mem:[0-9a-f]{16}", text).group(0)
|
||
|
||
def test_triggers_sync_index(self):
|
||
self.assertTrue(self.store.fts_ok)
|
||
record_id = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "tangerine", "content": "alpha zonk"})))
|
||
self.assertEqual(len(self.store.fts_search("zonk", ("record",), "t", 5)), 1)
|
||
self.store.upsert_file_record("/tmp/x.txt", "beta zonk")
|
||
self.assertEqual(len(self.store.fts_search("zonk", ("record",), "t", 5)), 2)
|
||
self.store.delete_record(record_id, "t")
|
||
hits = self.store.fts_search("zonk", ("record",), "t", 5)
|
||
self.assertEqual(len(hits), 1)
|
||
self.assertNotEqual(hits[0]["item"], record_id)
|
||
self.store.log_event("t", "user", "message", "zonk event here")
|
||
event_hits = self.store.fts_search("zonk", ("event",), "t", 5)
|
||
self.assertEqual(len(event_hits), 1)
|
||
self.assertIn("[zonk]", event_hits[0]["snippet"])
|
||
found_events = self.store.search_events("t", "zonk event")
|
||
self.assertEqual(len(found_events), 1)
|
||
self.assertIn("zonk event here", found_events[0][3])
|
||
self.assertNotIn("tai1$", found_events[0][3])
|
||
path = os.path.join(self.tmp.name, "zonk.txt")
|
||
self.tools.dispatch("write_file", json.dumps({"path": path, "content": "x"}))
|
||
found = self.store.audit_search("zonk", "t", 5)
|
||
self.assertEqual(len(found), 1)
|
||
self.assertEqual(found[0]["path"], path)
|
||
|
||
def test_ranked_order_and_snippet(self):
|
||
self.store.add_record("note", "sparse", "quasar convenes", [], "t")
|
||
self.store.add_record("note", "dense", "quasar " * 20, [], "t")
|
||
hits = self.store.fts_search("quasar", ("record",), "t", 5)
|
||
self.assertEqual(len(hits), 2)
|
||
self.assertEqual(hits[0]["title"], "dense")
|
||
self.assertIn("[quasar]", hits[0]["snippet"])
|
||
|
||
def test_unified_search_mixed_and_expand(self):
|
||
first = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "plan", "content": "harbor launch", "tags": ["harbor"]})))
|
||
second = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "log", "content": "harbor diary", "tags": ["harbor"]})))
|
||
self.store.log_event("t", "user", "message", "harbor standup notes")
|
||
out = self.tools.dispatch("search", json.dumps({"query": "harbor"}))
|
||
self.assertIn("[record/", out)
|
||
self.assertIn(first, out)
|
||
self.assertIn("[event/", out)
|
||
narrowed = self.tools.dispatch("search", json.dumps({"query": "harbor", "kinds": ["event"]}))
|
||
self.assertIn("[event/", narrowed)
|
||
self.assertNotIn("[record/", narrowed)
|
||
expanded = self.tools.dispatch("search", json.dumps({"query": "harbor launch", "kinds": ["record"], "expand": True}))
|
||
self.assertIn(first, expanded)
|
||
self.assertIn("linked:", expanded)
|
||
self.assertIn("log", expanded)
|
||
self.assertIn(second, self.tools.dispatch("search", json.dumps({"query": "harbor", "kinds": ["record"], "expand": False})))
|
||
|
||
def test_record_search_fallback_partial_token(self):
|
||
record_id = self.mem_id(self.tools.dispatch("record_save", json.dumps({"title": "soup", "content": "alphabet soup serving", "tags": ["lunch"]})))
|
||
found = self.tools.dispatch("record_search", json.dumps({"query": "alphab", "tags": ["lunch"]}))
|
||
self.assertIn(record_id, found)
|
||
|
||
def test_recall_sealed_fallback(self):
|
||
other_home = os.path.join(self.tmp.name, "sealed")
|
||
os.makedirs(other_home)
|
||
|
||
class SealedArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
with mock.patch.dict(os.environ, {"TAI_HOME": other_home}):
|
||
config = tai.Config(SealedArgs())
|
||
sealed = tai.Store(config, tai.Seal(config.home, "sealed-pw-3"))
|
||
try:
|
||
sealed.log_event("t", "user", "message", "sealed recall marker words")
|
||
rows = sealed.search_events("t", "recall marker")
|
||
self.assertEqual(len(rows), 1)
|
||
self.assertIn("sealed recall marker", rows[0][3])
|
||
finally:
|
||
sealed.close()
|
||
|
||
def test_backfill_on_reopen(self):
|
||
self.store.add_record("note", "keep", "backfill beacon", [], "t")
|
||
self.store.db.execute("DELETE FROM fts_docs")
|
||
self.store.db.commit()
|
||
self.assertEqual(self.store.fts_search("beacon", ("record",), "t", 5), [])
|
||
self.store.close()
|
||
reopened = tai.Store(self.config, tai.Seal(self.config.home, "fts-test-1"))
|
||
try:
|
||
hits = reopened.fts_search("beacon", ("record",), "t", 5)
|
||
self.assertEqual(len(hits), 1)
|
||
finally:
|
||
reopened.close()
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "fts-test-1"))
|
||
|
||
def test_search_tool_validation(self):
|
||
self.assertIn("empty query", self.tools.dispatch("search", json.dumps({"query": ""})))
|
||
self.assertIn("kinds must be", self.tools.dispatch("search", json.dumps({"query": "x", "kinds": ["nope"]})))
|
||
self.assertIn("invalid limit", self.tools.dispatch("search", json.dumps({"query": "x", "limit": "z"})))
|
||
self.assertIn("no matches", self.tools.dispatch("search", json.dumps({"query": "zzz-nothing-here"})))
|
||
|
||
def test_audit_query(self):
|
||
path = os.path.join(self.tmp.name, "queryme.txt")
|
||
self.tools.dispatch("write_file", json.dumps({"path": path, "content": "x"}))
|
||
out = self.tools.dispatch("audit", json.dumps({"query": "queryme"}))
|
||
self.assertIn("queryme.txt", out)
|
||
|
||
def test_search_lazy_tag(self):
|
||
found = {schema["function"]["name"] for schema in tai.select_tools("find everything about the plan")}
|
||
self.assertIn("search", found)
|
||
|
||
def test_search_repl(self):
|
||
self.tools.dispatch("record_save", json.dumps({"title": "repl", "content": "repl beacon words"}))
|
||
self.app.tools = self.tools
|
||
with mock.patch("builtins.print") as shown:
|
||
tai.handle_command(self.app, "/search repl beacon")
|
||
tai.handle_command(self.app, "/search")
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list)
|
||
self.assertIn("[record/", printed)
|
||
self.assertIn("use /search <query>", printed)
|
||
|
||
def test_mem_index_ranked_sealed(self):
|
||
self.assertIsNotNone(self.store.memdb)
|
||
self.store.log_event("t", "user", "message", "sparse comet sighting")
|
||
self.store.log_event("t", "user", "message", "comet " * 15)
|
||
hits = self.store.fts_search("comet", ("event",), "t", 5)
|
||
self.assertEqual(len(hits), 2)
|
||
self.assertLess(hits[0]["rank"], hits[1]["rank"])
|
||
self.assertIn("[comet]", hits[0]["snippet"])
|
||
self.assertNotIn("tai1$", hits[0]["snippet"])
|
||
stored = self.store.db.execute("SELECT text FROM events WHERE text LIKE 'tai1$%'").fetchall()
|
||
self.assertEqual(len(stored), 2)
|
||
|
||
def test_mem_index_cross_process_sync(self):
|
||
self.store.log_event("t", "user", "message", "first syncable event")
|
||
other = tai.Store(self.config, tai.Seal(self.config.home, "fts-test-1"))
|
||
try:
|
||
other.log_event("t", "user", "message", "second syncable event")
|
||
finally:
|
||
other.close()
|
||
hits = self.store.fts_search("syncable", ("event",), "t", 5)
|
||
self.assertEqual(len(hits), 2)
|
||
|
||
def test_mem_index_profile_scoped(self):
|
||
self.store.log_event("t", "user", "message", "scoped beacon words")
|
||
self.store.log_event("other", "user", "message", "scoped beacon words")
|
||
self.assertEqual(len(self.store.fts_search("beacon", ("event",), "t", 5)), 1)
|
||
self.assertEqual(len(self.store.fts_search("beacon", ("event",), "other", 5)), 1)
|
||
|
||
def test_mem_index_survives_rotation(self):
|
||
self.store.log_event("t", "user", "message", "rotation beacon words")
|
||
old = tai.Seal(self.config.home, "fts-test-1")
|
||
self.store.close()
|
||
tai.rotate_seal(self.config, old, "fts-test-2")
|
||
reopened = tai.Store(self.config, tai.Seal(self.config.home, "fts-test-2"))
|
||
try:
|
||
hits = reopened.fts_search("rotation", ("event",), "t", 5)
|
||
self.assertEqual(len(hits), 1)
|
||
self.assertIn("[rotation]", hits[0]["snippet"])
|
||
finally:
|
||
reopened.close()
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "fts-test-2"))
|
||
|
||
def test_run_search_sealed_events_ranked_and_partial(self):
|
||
self.store.log_event("t", "user", "message", "harbor standup notes")
|
||
out = self.tools.dispatch("search", json.dumps({"query": "harbor", "kinds": ["event"]}))
|
||
self.assertIn("[event/", out)
|
||
self.assertIn("[harbor]", out)
|
||
partial = self.tools.dispatch("search", json.dumps({"query": "standu", "kinds": ["event"]}))
|
||
self.assertIn("[event/", partial)
|
||
self.assertIn("standup", partial)
|
||
|
||
def test_rotate_multi_profile_secrets(self):
|
||
self.store.save_secret("api", "t-secret-1", None, None, "t")
|
||
self.store.save_secret("api", "o-secret-2", None, None, "other")
|
||
old = tai.Seal(self.config.home, "fts-test-1")
|
||
self.store.close()
|
||
tai.rotate_seal(self.config, old, "fts-test-2")
|
||
reopened = tai.Store(self.config, tai.Seal(self.config.home, "fts-test-2"))
|
||
try:
|
||
self.assertEqual(reopened.load_secret("api", "t"), "t-secret-1")
|
||
self.assertEqual(reopened.load_secret("api", "other"), "o-secret-2")
|
||
finally:
|
||
reopened.close()
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "fts-test-2"))
|
||
|
||
|
||
class StreamTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "stream-test-1"))
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def test_truncate_ansi_keeps_short_colors(self):
|
||
text = "\033[31mhi\033[0m"
|
||
self.assertEqual(tai.truncate_ansi(text, 10), text)
|
||
|
||
def test_truncate_ansi_cuts_visible_width(self):
|
||
cut = tai.truncate_ansi("\033[31m" + "x" * 50 + "\033[0m", 10)
|
||
self.assertEqual(tai.strip_ansi(cut), "x" * 10)
|
||
self.assertTrue(cut.endswith(tai.Ansi.RESET))
|
||
self.assertIn("\033[31m", cut)
|
||
|
||
def test_truncate_ansi_plain(self):
|
||
self.assertEqual(tai.truncate_ansi("abcdef", 4), "abcd" + tai.Ansi.RESET)
|
||
|
||
def test_human_size(self):
|
||
self.assertEqual(tai.human_size(0), "0 B")
|
||
self.assertEqual(tai.human_size(512), "512 B")
|
||
self.assertEqual(tai.human_size(1024), "1.0 KB")
|
||
self.assertEqual(tai.human_size(1536), "1.5 KB")
|
||
self.assertEqual(tai.human_size(2097152), "2.0 MB")
|
||
|
||
def test_stream_status(self):
|
||
self.assertEqual(tai.stream_status(0, "a\nb\n", 1.25, False), "exit 0 · 1.2s · 2 lines · 4 B")
|
||
self.assertEqual(tai.stream_status(3, "", 0.0, False), "exit 3 · 0.0s · 0 lines · 0 B")
|
||
self.assertEqual(tai.stream_status(-9, "x\n" * 10, 120.0, True), "timed out after 120s (killed) · 10 lines · 20 B")
|
||
self.assertEqual(tai.stream_status(0, "z" * 2048, 0.5, False), "exit 0 · 0.5s · 1 lines · 2.0 KB")
|
||
|
||
def test_window_rolls_at_height(self):
|
||
buf = io.StringIO()
|
||
stream = tai.LiveStream(height=2, width=40, file=buf, enabled=True)
|
||
stream.feed("one\n")
|
||
stream.feed("two\n")
|
||
stream.feed("three\n")
|
||
stream.close("exit 0")
|
||
out = buf.getvalue()
|
||
self.assertIn("│ one\n│ two\n", out)
|
||
self.assertIn("\033[2A", out)
|
||
self.assertIn("\r\033[K│ three\n", out)
|
||
self.assertTrue(out.endswith("│ exit 0\n"))
|
||
|
||
def test_window_disabled_writes_nothing(self):
|
||
buf = io.StringIO()
|
||
stream = tai.LiveStream(file=buf, enabled=False)
|
||
stream.feed("one\n")
|
||
stream.close("exit 0")
|
||
self.assertEqual(buf.getvalue(), "")
|
||
|
||
def test_window_no_output_status(self):
|
||
buf = io.StringIO()
|
||
stream = tai.LiveStream(file=buf, enabled=True)
|
||
stream.close("exit 3")
|
||
self.assertEqual(buf.getvalue(), "│ (no output)\n│ exit 3\n")
|
||
|
||
def test_window_keeps_progress_tail(self):
|
||
buf = io.StringIO()
|
||
stream = tai.LiveStream(height=4, width=40, file=buf, enabled=True)
|
||
stream.feed("50%\r100%\n")
|
||
self.assertIn("│ 100%\n", buf.getvalue())
|
||
|
||
def test_run_live_captures(self):
|
||
code, out, _elapsed, timed = tai.run_live(["echo", "hi"])
|
||
self.assertEqual((code, out, timed), (0, "hi\n", False))
|
||
|
||
def test_run_live_exit_code(self):
|
||
code, _out, _elapsed, timed = tai.run_live(["sh", "-c", "exit 3"])
|
||
self.assertEqual((code, timed), (3, False))
|
||
|
||
def test_run_live_merges_stderr(self):
|
||
code, out, _elapsed, _timed = tai.run_live(["sh", "-c", "echo out; echo err >&2"])
|
||
self.assertIn("out\n", out)
|
||
self.assertIn("err\n", out)
|
||
self.assertEqual(code, 0)
|
||
|
||
def test_run_live_streams(self):
|
||
buf = io.StringIO()
|
||
stream = tai.LiveStream(file=buf, enabled=True)
|
||
_code, out, _elapsed, timed = tai.run_live(["echo", "streamed"], stream=stream)
|
||
stream.close("exit 0")
|
||
self.assertEqual(out, "streamed\n")
|
||
self.assertFalse(timed)
|
||
self.assertIn("│ streamed\n", buf.getvalue())
|
||
|
||
def test_run_live_timeout_kills(self):
|
||
started = time.time()
|
||
_code, _out, _elapsed, timed = tai.run_live(["sleep", "30"], timeout=1)
|
||
self.assertTrue(timed)
|
||
self.assertLess(time.time() - started, 10)
|
||
|
||
def test_shell_live_flag(self):
|
||
app = FakeSecretApp(self.store)
|
||
app.profile = "t"
|
||
tools = tai.Tools(app)
|
||
self.assertFalse(tools.live)
|
||
tools.live = True
|
||
buf = io.StringIO()
|
||
buf.isatty = lambda: True
|
||
with mock.patch.object(sys, "stdout", buf):
|
||
result = tools.dispatch("shell", json.dumps({"command": "echo live-mark"}))
|
||
self.assertTrue(result.startswith("exit 0"))
|
||
self.assertIn("live-mark", result)
|
||
self.assertIn("│ live-mark\n", buf.getvalue())
|
||
self.assertIn("exit 0 ·", buf.getvalue())
|
||
self.assertIn("1 lines · 10 B", buf.getvalue())
|
||
|
||
def test_shell_silent_without_flag(self):
|
||
app = FakeSecretApp(self.store)
|
||
app.profile = "t"
|
||
tools = tai.Tools(app)
|
||
buf = io.StringIO()
|
||
buf.isatty = lambda: True
|
||
with mock.patch.object(sys, "stdout", buf):
|
||
result = tools.dispatch("shell", json.dumps({"command": "echo quiet-mark"}))
|
||
self.assertTrue(result.startswith("exit 0"))
|
||
self.assertEqual(buf.getvalue(), "")
|
||
|
||
def test_agent_enables_live_for_turn(self):
|
||
agent = tai.Agent(self.config, self.store, persist=False, quiet=True)
|
||
reply = {"role": "assistant", "content": "done", "reasoning": "", "tool_calls": [], "backend": "x"}
|
||
with mock.patch.object(agent.chat, "complete", return_value=dict(reply)):
|
||
agent.run_turn("hi", capture=True)
|
||
self.assertFalse(agent.tools.live)
|
||
agent.quiet = False
|
||
with mock.patch.object(agent.chat, "complete", return_value=dict(reply)):
|
||
with mock.patch("builtins.print"):
|
||
agent.run_turn("hi")
|
||
self.assertTrue(agent.tools.live)
|
||
|
||
|
||
class ProgressTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "progress-test-1"))
|
||
self.agent = tai.Agent(self.config, self.store, persist=False, quiet=True)
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def tool_reply(self, command, call_id="1"):
|
||
return {"role": "assistant", "content": "", "reasoning": "", "tool_calls": [{"id": call_id, "name": "shell", "arguments": json.dumps({"command": command})}], "backend": "x"}
|
||
|
||
def text_reply(self, text="all done"):
|
||
return {"role": "assistant", "content": text, "reasoning": "", "tool_calls": [], "backend": "x"}
|
||
|
||
def tool_count(self):
|
||
return len([item for item in self.agent.messages if item["role"] == "tool"])
|
||
|
||
def test_novelty_resets_stall(self):
|
||
replies = [self.tool_reply("echo novel-%d" % num, str(num)) for num in range(5)] + [self.text_reply()]
|
||
with mock.patch.object(self.agent.chat, "complete", side_effect=[dict(item) for item in replies]):
|
||
answer = self.agent.run_turn("go", capture=True, max_steps=3)
|
||
self.assertEqual(answer, "all done")
|
||
self.assertEqual(self.tool_count(), 5)
|
||
|
||
def test_loop_nudge_then_stop(self):
|
||
replies = [self.tool_reply("echo same", str(num)) for num in range(6)] + [self.text_reply()]
|
||
with mock.patch.object(self.agent.chat, "complete", side_effect=[dict(item) for item in replies]):
|
||
answer = self.agent.run_turn("go", capture=True, max_steps=10)
|
||
self.assertIn("stuck in a loop", answer)
|
||
self.assertEqual(self.tool_count(), 5)
|
||
nudges = [item for item in self.agent.messages if item["role"] == "user" and "loop warning" in item.get("content", "")]
|
||
self.assertEqual(len(nudges), 1)
|
||
|
||
def test_nudge_recovery(self):
|
||
replies = [self.tool_reply("echo same", str(num)) for num in range(3)] + [self.tool_reply("echo different", "9"), self.text_reply()]
|
||
with mock.patch.object(self.agent.chat, "complete", side_effect=[dict(item) for item in replies]):
|
||
answer = self.agent.run_turn("go", capture=True, max_steps=10)
|
||
self.assertEqual(answer, "all done")
|
||
self.assertEqual(self.tool_count(), 4)
|
||
nudges = [item for item in self.agent.messages if item["role"] == "user" and "loop warning" in item.get("content", "")]
|
||
self.assertEqual(len(nudges), 1)
|
||
|
||
def test_stall_oscillation(self):
|
||
replies = []
|
||
for num in range(8):
|
||
replies.append(self.tool_reply("echo osc-%d" % (num % 2), str(num)))
|
||
replies.append(self.text_reply())
|
||
with mock.patch.object(self.agent.chat, "complete", side_effect=[dict(item) for item in replies]):
|
||
answer = self.agent.run_turn("go", capture=True, max_steps=3)
|
||
self.assertIn("no progress for 3 steps", answer)
|
||
self.assertEqual(self.tool_count(), 5)
|
||
|
||
def test_user_interaction_resets(self):
|
||
replies = [self.tool_reply("echo int-%d" % (num % 2), str(num)) for num in range(8)] + [self.text_reply()]
|
||
real_dispatch = self.agent.tools.dispatch
|
||
calls = []
|
||
|
||
def approving(name, raw):
|
||
calls.append(name)
|
||
if len(calls) <= 3:
|
||
self.agent.user_stepped_in = True
|
||
return real_dispatch(name, raw)
|
||
|
||
with mock.patch.object(self.agent.chat, "complete", side_effect=[dict(item) for item in replies]):
|
||
with mock.patch.object(self.agent.tools, "dispatch", side_effect=approving):
|
||
answer = self.agent.run_turn("go", capture=True, max_steps=2)
|
||
self.assertIn("no progress for 2 steps", answer)
|
||
self.assertEqual(self.tool_count(), 5)
|
||
|
||
def test_total_cap_backstop(self):
|
||
replies = [self.tool_reply("echo fresh-%d" % num, str(num)) for num in range(8)] + [self.text_reply()]
|
||
with mock.patch.object(tai.Consts, "TOTAL_STEP_CAP", 6):
|
||
with mock.patch.object(self.agent.chat, "complete", side_effect=[dict(item) for item in replies]):
|
||
answer = self.agent.run_turn("go", capture=True, max_steps=100)
|
||
self.assertEqual(answer, "")
|
||
self.assertEqual(self.tool_count(), 6)
|
||
|
||
def test_ask_approval_marks_interaction(self):
|
||
agent = tai.Agent(self.config, self.store, persist=True, quiet=True)
|
||
self.assertFalse(agent.user_stepped_in)
|
||
fake_stdin = io.StringIO("y\n")
|
||
fake_stdin.isatty = lambda: True
|
||
with mock.patch.object(sys, "stdin", fake_stdin):
|
||
with mock.patch("builtins.print"):
|
||
self.assertTrue(agent.ask_approval("echo hi"))
|
||
self.assertTrue(agent.user_stepped_in)
|
||
auto_agent = tai.Agent(self.config, self.store, persist=True, quiet=True, auto=True)
|
||
self.assertTrue(auto_agent.ask_approval("echo hi"))
|
||
self.assertFalse(auto_agent.user_stepped_in)
|
||
|
||
|
||
class WebFetchTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "fetch-test-1"))
|
||
self.app = FakeSecretApp(self.store)
|
||
self.app.profile = "t"
|
||
self.tools = tai.Tools(self.app)
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def fake_open(self, body, status=200, content_type="text/html; charset=utf-8", seen=None):
|
||
class FakeResponse:
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, *exc):
|
||
return False
|
||
|
||
def read(self, limit=0):
|
||
return body
|
||
|
||
def fake_urlopen(request, timeout=30):
|
||
if seen is not None:
|
||
seen["method"] = request.get_method()
|
||
seen["data"] = request.data
|
||
seen["headers"] = {key.lower(): value for key, value in request.headers.items()}
|
||
seen["timeout"] = timeout
|
||
response = FakeResponse()
|
||
response.status = status
|
||
response.headers = {"Content-Type": content_type}
|
||
return response
|
||
|
||
return fake_urlopen
|
||
|
||
def test_fetch_post_json(self):
|
||
seen = {}
|
||
payload = b'{"ok": true, "tag": "<kept>"}'
|
||
with mock.patch("urllib.request.urlopen", side_effect=self.fake_open(payload, 201, "application/json", seen)):
|
||
result = self.tools.dispatch("web_fetch", json.dumps({"url": "https://api.example.test/items", "method": "post", "headers": {"Content-Type": "application/json", "X-Trace": "1"}, "body": payload.decode("utf-8")}))
|
||
self.assertEqual(seen["method"], "POST")
|
||
self.assertEqual(seen["data"], payload)
|
||
self.assertEqual(seen["headers"]["content-type"], "application/json")
|
||
self.assertEqual(seen["headers"]["x-trace"], "1")
|
||
self.assertTrue(result.startswith("HTTP 201 · application/json · 29 B\n"))
|
||
self.assertIn("<kept>", result)
|
||
|
||
def test_fetch_html_strips(self):
|
||
with mock.patch("urllib.request.urlopen", side_effect=self.fake_open(b"<html><body><h1>Hi</h1><script>var x = 1;</script></body></html>")):
|
||
result = self.tools.dispatch("web_fetch", json.dumps({"url": "https://example.test/"}))
|
||
self.assertTrue(result.startswith("HTTP 200 · text/html · "))
|
||
self.assertIn("Hi", result)
|
||
self.assertNotIn("<h1>", result)
|
||
self.assertNotIn("var x", result)
|
||
|
||
def test_fetch_raw_keeps_tags(self):
|
||
with mock.patch("urllib.request.urlopen", side_effect=self.fake_open(b"<p>Hi</p>")):
|
||
result = self.tools.dispatch("web_fetch", json.dumps({"url": "https://example.test/", "raw": True}))
|
||
self.assertIn("<p>Hi</p>", result)
|
||
|
||
def test_fetch_validation(self):
|
||
base = "https://example.test/"
|
||
self.assertIn("method must be", self.tools.dispatch("web_fetch", json.dumps({"url": base, "method": "BREW"})))
|
||
self.assertIn("headers must be", self.tools.dispatch("web_fetch", json.dumps({"url": base, "headers": ["x"]})))
|
||
self.assertIn("headers must be", self.tools.dispatch("web_fetch", json.dumps({"url": base, "headers": {"X-A": 1}})))
|
||
self.assertIn("invalid header", self.tools.dispatch("web_fetch", json.dumps({"url": base, "headers": {"Bad Name": "x"}})))
|
||
self.assertIn("invalid header", self.tools.dispatch("web_fetch", json.dumps({"url": base, "headers": {"X-A": "a\nb"}})))
|
||
self.assertIn("body must be", self.tools.dispatch("web_fetch", json.dumps({"url": base, "method": "POST", "body": 42})))
|
||
self.assertIn("invalid timeout", self.tools.dispatch("web_fetch", json.dumps({"url": base, "timeout": "x"})))
|
||
self.assertIn("url must start", self.tools.dispatch("web_fetch", json.dumps({"url": "ftp://example.test/"})))
|
||
|
||
def test_fetch_http_error_surfaces_body(self):
|
||
failure = urllib.error.HTTPError("https://example.test/", 404, "Not Found", None, io.BytesIO(b"no such widget"))
|
||
with mock.patch("urllib.request.urlopen", side_effect=failure):
|
||
result = self.tools.dispatch("web_fetch", json.dumps({"url": "https://example.test/"}))
|
||
self.assertEqual(result, "error: HTTP 404: no such widget")
|
||
empty = urllib.error.HTTPError("https://example.test/", 500, "Server Error", None, io.BytesIO(b""))
|
||
with mock.patch("urllib.request.urlopen", side_effect=empty):
|
||
result = self.tools.dispatch("web_fetch", json.dumps({"url": "https://example.test/"}))
|
||
self.assertTrue(result.startswith("error: HTTP 500:"))
|
||
|
||
def test_fetch_connection_error(self):
|
||
with mock.patch("urllib.request.urlopen", side_effect=urllib.error.URLError("refused")):
|
||
result = self.tools.dispatch("web_fetch", json.dumps({"url": "https://example.test/"}))
|
||
self.assertIn("error: fetch failed:", result)
|
||
|
||
def test_shell_curl_hint(self):
|
||
marked = self.tools.dispatch("shell", json.dumps({"command": "echo curl http://example.test"}))
|
||
self.assertIn(tai.CURL_HINT, marked)
|
||
plain = self.tools.dispatch("shell", json.dumps({"command": "echo just-local"}))
|
||
self.assertNotIn("web_fetch", plain)
|
||
self.assertEqual(tai.curl_hint("wget https://example.test/x"), "\n" + tai.CURL_HINT)
|
||
self.assertEqual(tai.curl_hint("curl --version"), "")
|
||
self.assertEqual(tai.curl_hint("echo hi"), "")
|
||
|
||
def test_steering_text(self):
|
||
descs = {schema["function"]["name"]: schema["function"]["description"] for schema in tai.TOOL_SCHEMAS}
|
||
self.assertIn("never curl or wget", descs["web_fetch"])
|
||
self.assertIn("use web_fetch for HTTP(S)", descs["shell"])
|
||
self.assertIn("never curl or wget", tai.DEFAULT_SYSTEM)
|
||
|
||
|
||
class DiffTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "diff-test-1"))
|
||
self.app = FakeSecretApp(self.store)
|
||
self.app.profile = "t"
|
||
self.tools = tai.Tools(self.app)
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def test_highlight_python(self):
|
||
out = tai.highlight_python("def f(x=1): # hi")
|
||
self.assertIn(tai.Ansi.MAGENTA + "def" + tai.Ansi.RESET, out)
|
||
self.assertIn(tai.Ansi.CYAN + "f" + tai.Ansi.RESET, out)
|
||
self.assertIn(tai.FG_ORANGE + "1" + tai.Ansi.RESET, out)
|
||
self.assertIn(tai.Ansi.GRAY + "# hi" + tai.Ansi.RESET, out)
|
||
self.assertEqual(tai.strip_ansi(out), "def f(x=1): # hi")
|
||
|
||
def test_highlight_string_with_hash(self):
|
||
out = tai.highlight_python("x = 'a # b'")
|
||
self.assertIn(tai.Ansi.YELLOW + "'a # b'" + tai.Ansi.RESET, out)
|
||
self.assertNotIn(tai.Ansi.GRAY, out)
|
||
|
||
def test_render_diff_plain(self):
|
||
out = tai.render_diff("a.py", "x = 1\n", "x = 2\n", width=60, colors=False)
|
||
self.assertEqual(out, "│ ── a.py (+1 -1)\n 1 - x = 1\n 1 + x = 2")
|
||
|
||
def test_render_diff_colors(self):
|
||
out = tai.render_diff("a.py", "x = 1\n", "x = 2\n", width=60, colors=True)
|
||
self.assertIn(tai.BG_ADD, out)
|
||
self.assertIn(tai.BG_DEL, out)
|
||
self.assertIn(tai.FG_ORANGE + "2" + tai.Ansi.RESET, out)
|
||
for line in out.splitlines()[1:]:
|
||
self.assertEqual(len(tai.strip_ansi(line)), 60)
|
||
|
||
def test_render_diff_gap_and_cap(self):
|
||
old = "".join("line %d\n" % num for num in range(20))
|
||
new = old.replace("line 0\n", "line zero\n").replace("line 19\n", "line nineteen\n")
|
||
out = tai.render_diff("a.py", old, new, width=60, colors=False)
|
||
self.assertIn("···", out)
|
||
big = "".join("row %d\n" % num for num in range(200))
|
||
capped = tai.render_diff("a.py", "", big, width=60, colors=False)
|
||
self.assertIn("more lines hidden", capped)
|
||
self.assertEqual(len(capped.splitlines()), 1 + tai.DIFF_MAX_LINES + 1)
|
||
|
||
def test_render_diff_shapes(self):
|
||
self.assertEqual(tai.render_diff("x.py", "a\n", "a\n"), "")
|
||
created = tai.render_diff("x.py", None, "a\nb\n", width=60, colors=False)
|
||
self.assertIn("(new file, 2 lines)", created.splitlines()[0])
|
||
removed = tai.render_diff("x.py", "a\nb\n", None, width=60, colors=False)
|
||
self.assertIn("(deleted, 2 lines)", removed.splitlines()[0])
|
||
plain = tai.render_diff("notes.txt", "a\n", "b\n", width=60, colors=True)
|
||
self.assertIn(tai.BG_ADD, plain)
|
||
self.assertNotIn(tai.Ansi.MAGENTA, plain)
|
||
self.assertNotIn(tai.Ansi.YELLOW, plain)
|
||
|
||
def test_render_diff_docstring_fence(self):
|
||
out = tai.render_diff("a.py", "", '"""\ndoc body\n"""\nx = 1\n', width=60, colors=True)
|
||
self.assertIn(tai.Ansi.YELLOW + "doc body" + tai.Ansi.RESET, out)
|
||
|
||
def test_tools_note_preview(self):
|
||
path = os.path.join(self.tmp.name, "note.py")
|
||
self.assertTrue(self.tools.dispatch("write_file", json.dumps({"path": path, "content": "x = 1\n"})).startswith("wrote "))
|
||
preview = self.tools.diff_preview
|
||
self.assertEqual((preview["path"], preview["old"], preview["new"]), (path, None, "x = 1\n"))
|
||
self.tools.dispatch("read_file", json.dumps({"path": path}))
|
||
self.assertTrue(self.tools.dispatch("edit_file", json.dumps({"path": path, "find": "x = 1", "replace": "x = 2"})).startswith("edited "))
|
||
preview = self.tools.diff_preview
|
||
self.assertEqual((preview["old"], preview["new"]), ("x = 1\n", "x = 2\n"))
|
||
self.assertTrue(self.tools.dispatch("delete_file", json.dumps({"path": path})).startswith("deleted "))
|
||
preview = self.tools.diff_preview
|
||
self.assertEqual((preview["old"], preview["new"]), ("x = 2\n", None))
|
||
|
||
def test_show_result_renders_and_clears(self):
|
||
agent = tai.Agent(self.config, self.store, persist=False, quiet=False)
|
||
agent.tools.diff_preview = {"path": "a.py", "old": "x = 1\n", "new": "x = 2\n"}
|
||
with mock.patch.object(sys.stdout, "isatty", return_value=False):
|
||
with mock.patch("builtins.print") as shown:
|
||
agent.show_result("ok", 1)
|
||
agent.show_result("ok", 1)
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list)
|
||
self.assertEqual(printed.count("── a.py"), 1)
|
||
self.assertEqual(printed.count("└─"), 2)
|
||
self.assertNotIn("\033", printed)
|
||
self.assertIsNone(agent.tools.diff_preview)
|
||
agent.quiet = True
|
||
agent.tools.diff_preview = {"path": "a.py", "old": "x = 1\n", "new": "x = 2\n"}
|
||
with mock.patch("builtins.print") as shown:
|
||
agent.show_result("ok", 1)
|
||
shown.assert_not_called()
|
||
|
||
|
||
class GlowRenderingTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "glow-test-1"))
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def test_describe_call_verbs(self):
|
||
self.assertEqual(tai.Glow.describe_call("shell", json.dumps({"command": "make test\nsecond"})),
|
||
("Run", "make test second", "shell", "make test second"))
|
||
self.assertEqual(tai.Glow.describe_call("read_file", json.dumps({"path": "/tmp/x/main.py"})),
|
||
("Read", "main.py", "read_file", "/tmp/x/main.py"))
|
||
self.assertEqual(tai.Glow.describe_call("shell", {"command": "ls"}), ("Run", "ls", "shell", "ls"))
|
||
self.assertEqual(tai.Glow.describe_call("shell", "nope{"), ("Run", "", "shell", "nope{"))
|
||
self.assertEqual(tai.Glow.describe_call("mystery", "{}"), ("mystery", "", "", "{}"))
|
||
self.assertEqual(tai.Glow.describe_call("poll", {"id": 3}), ("Poll", "3", "poll", "agent 3"))
|
||
self.assertEqual(tai.Glow.describe_call("write_file", json.dumps({"path": "n.txt", "content": "x" * 500}))[3], "n.txt")
|
||
verb, target, _extra, preview = tai.Glow.describe_call("store_secret", json.dumps({"name": "api", "value": "s3cr3t"}))
|
||
self.assertEqual(verb, "Seal")
|
||
self.assertNotIn("s3cr3t", target)
|
||
self.assertNotIn("s3cr3t", preview)
|
||
self.assertIn("[hidden]", preview)
|
||
|
||
def test_shimmer_sweeps(self):
|
||
out = tai.Glow.shimmer("thinking", 2)
|
||
self.assertIn(tai.Ansi.BOLD, out)
|
||
self.assertIn(tai.Ansi.DIM, out)
|
||
self.assertEqual(tai.strip_ansi(out), "thinking")
|
||
self.assertEqual(tai.strip_ansi(tai.Glow.shimmer("thinking", 99)), "thinking")
|
||
self.assertEqual(tai.Glow.shimmer("", 5), "")
|
||
|
||
def test_event_and_done_rows(self):
|
||
row = tai.Glow.event_row("◆", "Edited", "a.py", "(+1 -1)")
|
||
self.assertIn("◆", row)
|
||
self.assertIn(tai.FG_FILE, row)
|
||
self.assertEqual(tai.strip_ansi(row), "◆ Edited a.py (+1 -1)")
|
||
done = tai.Glow.done_row("Run", "make test", 12, 234)
|
||
self.assertIn("✓", done)
|
||
self.assertIn("234ms", done)
|
||
self.assertIn("12 chars", done)
|
||
failed = tai.Glow.done_row("Run", "", 5, 1500, failed=True)
|
||
self.assertIn("✗", failed)
|
||
self.assertIn("1.5s", failed)
|
||
short = tai.Glow.event_row("◆", "Edited", "a.py", "(+1 -1)", width=10)
|
||
self.assertEqual(len(tai.strip_ansi(short)), 10)
|
||
|
||
def test_kind_icons_and_ellipsis(self):
|
||
self.assertEqual(tai.Glow.icon_for("shell"), "▶")
|
||
self.assertEqual(tai.Glow.icon_for("fork"), "✦")
|
||
self.assertEqual(tai.Glow.icon_for("poll"), "⟳")
|
||
self.assertEqual(tai.Glow.icon_for("read_file"), "◇")
|
||
self.assertEqual(tai.Glow.icon_for("mystery"), "◇")
|
||
self.assertEqual(tai.Glow.ellipsis("short", 60), "short")
|
||
long = "a" * 30 + "/" + "b" * 30
|
||
cut = tai.Glow.ellipsis(long, 20)
|
||
self.assertEqual(len(cut), 20)
|
||
self.assertIn("…", cut)
|
||
self.assertTrue(cut.startswith("a" * 9))
|
||
self.assertTrue(cut.endswith("b" * 10))
|
||
self.assertEqual(tai.Glow.ellipsis(long, 4), long)
|
||
self.assertEqual(tai.Glow.ellipsis("a" * 60, 20, head=True), "a" * 19 + "…")
|
||
|
||
def test_accents_and_bars(self):
|
||
self.assertEqual(tai.Glow.accent_for("shell"), tai.FG_RUN)
|
||
self.assertEqual(tai.Glow.accent_for("fork"), tai.FG_FORK)
|
||
self.assertEqual(tai.Glow.accent_for("poll"), tai.FG_FILE)
|
||
self.assertEqual(tai.Glow.accent_for("edit_file"), tai.Ansi.GREEN)
|
||
self.assertEqual(tai.Glow.accent_for("ask_user"), tai.FG_ORANGE)
|
||
self.assertEqual(tai.Glow.accent_for("store_secret"), tai.FG_YOLO)
|
||
self.assertEqual(tai.Glow.accent_for("read_file"), tai.Ansi.DIM)
|
||
bar = tai.Glow.bar("shell")
|
||
self.assertIn("▌", bar)
|
||
self.assertIn(tai.FG_RUN, bar)
|
||
tinted = tai.Glow.event_row("▶", "Run", "x", "", 0, tai.FG_RUN)
|
||
self.assertIn(tai.Ansi.BOLD + tai.FG_RUN + "▶", tinted)
|
||
self.assertEqual(tai.strip_ansi(tinted), "▶ Run x")
|
||
|
||
def test_done_row_right_aligns_meta(self):
|
||
row = tai.Glow.done_row("Run", "x", 12, 234, width=60)
|
||
plain = tai.strip_ansi(row)
|
||
self.assertEqual(len(plain), 60)
|
||
self.assertTrue(plain.startswith("◆ Run x"))
|
||
self.assertTrue(plain.endswith("✓ · 234ms · 12 chars"))
|
||
narrow = tai.Glow.done_row("Run", "averylongtargetname", 12, 234, width=10)
|
||
self.assertEqual(len(tai.strip_ansi(narrow)), 10)
|
||
|
||
def test_agent_row_and_lanes(self):
|
||
row = tai.Glow.agent_row({"id": 3, "task": "translate p1 to german", "status": "running", "elapsed": 75})
|
||
self.assertIn("#3", row)
|
||
self.assertIn(tai.Ansi.YELLOW + "●" + tai.Ansi.RESET, row)
|
||
self.assertIn("translate p1 to german", tai.strip_ansi(row))
|
||
self.assertIn("1m", tai.strip_ansi(row))
|
||
done = tai.Glow.agent_row({"id": 4, "task": "t", "status": "done", "elapsed": 3})
|
||
self.assertIn(tai.Ansi.GREEN + "●" + tai.Ansi.RESET, done)
|
||
many = [{"id": num, "task": "t%d" % num, "status": "running", "elapsed": 1} for num in range(8)]
|
||
rows = tai.Glow.lanes(many)
|
||
self.assertEqual(len(rows), 1 + 6 + 1)
|
||
self.assertIn("8 running", tai.strip_ansi(rows[0]))
|
||
self.assertIn("+2 more", tai.strip_ansi(rows[-1]))
|
||
|
||
def test_lane_tasks_keep_distinguishing_head(self):
|
||
task = "Translate README.md lines 94-186 to /tmp/xxx_p2.md. Rules: " + "x" * 80
|
||
row = tai.strip_ansi(tai.Glow.agent_row({"id": 1, "task": task, "status": "running", "elapsed": 0}))
|
||
self.assertIn("lines 94-186", row)
|
||
self.assertIn("…", row)
|
||
wide = tai.Glow.agent_row({"id": 1, "task": task, "status": "running", "elapsed": 0}, width=40)
|
||
self.assertLessEqual(len(tai.strip_ansi(wide)), 40)
|
||
pulsed = tai.Glow.agent_row({"id": 1, "task": "t", "status": "running", "elapsed": 0}, 0, tai.Ansi.BOLD)
|
||
self.assertIn(tai.Ansi.BOLD + tai.Ansi.YELLOW + "●", pulsed)
|
||
|
||
def test_lanes_show_pending_pickups(self):
|
||
records = [{"id": 1, "task": "a", "status": "running", "elapsed": 5, "collected": False},
|
||
{"id": 2, "task": "b", "status": "done", "elapsed": 9, "collected": False},
|
||
{"id": 3, "task": "c", "status": "done", "elapsed": 9, "collected": True}]
|
||
rows = tai.Glow.lanes(records)
|
||
self.assertIn("1 running · 1 to collect", tai.strip_ansi(rows[0]))
|
||
body = "\n".join(tai.strip_ansi(row) for row in rows[1:])
|
||
self.assertIn("#1", body)
|
||
self.assertIn("#2", body)
|
||
self.assertNotIn("#3", body)
|
||
self.assertIn("to collect", body)
|
||
|
||
def test_show_call_truncates_to_width(self):
|
||
agent = tai.Agent(self.config, self.store, persist=False, quiet=False)
|
||
with mock.patch.object(sys.stdout, "isatty", return_value=True):
|
||
with mock.patch.dict(os.environ, {"NO_COLOR": ""}):
|
||
with mock.patch("builtins.print") as shown:
|
||
agent.show_call({"name": "shell", "arguments": json.dumps({"command": "echo " + "x" * 300})})
|
||
width = tai.Glow.width()
|
||
for call in shown.call_args_list:
|
||
if call.args:
|
||
self.assertLessEqual(len(tai.strip_ansi(str(call.args[0]))), width)
|
||
|
||
def test_fork_done_row_carries_agent_id(self):
|
||
agent = tai.Agent(self.config, self.store, persist=False, quiet=False)
|
||
with mock.patch.object(sys.stdout, "isatty", return_value=True):
|
||
with mock.patch.dict(os.environ, {"NO_COLOR": ""}):
|
||
with mock.patch("builtins.print") as shown:
|
||
agent.show_call({"name": "fork", "arguments": json.dumps({"task": "translate p2"})})
|
||
with mock.patch("tai.list_agents", return_value=[]):
|
||
agent.show_result("agent 7 started: translate p2", 5)
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list if call.args)
|
||
self.assertIn("#7", printed)
|
||
|
||
def test_show_lanes_snapshot_and_collect(self):
|
||
agent = tai.Agent(self.config, self.store, persist=False, quiet=False)
|
||
running = [{"id": 3, "task": "translate p1 to german", "status": "running", "elapsed": 5},
|
||
{"id": 4, "task": "translate p2 to german", "status": "running", "elapsed": 5}]
|
||
with mock.patch.object(sys.stdout, "isatty", return_value=True):
|
||
with mock.patch.dict(os.environ, {"NO_COLOR": ""}):
|
||
with mock.patch("builtins.print") as shown:
|
||
agent.show_call({"name": "fork", "arguments": json.dumps({"task": "translate p2"})})
|
||
with mock.patch("tai.list_agents", return_value=running):
|
||
agent.show_result("agent 4 started: translate p2", 5)
|
||
with mock.patch("builtins.print") as collected:
|
||
agent.show_call({"name": "poll", "arguments": json.dumps({"id": 3})})
|
||
with mock.patch("tai.list_agents", return_value=[]):
|
||
agent.show_result("[done] translated text", 5)
|
||
with mock.patch("builtins.print") as quiet:
|
||
agent.show_call({"name": "poll", "arguments": json.dumps({"id": 3})})
|
||
with mock.patch("tai.list_agents", return_value=running[:1]):
|
||
agent.show_result("[done] translated text", 5)
|
||
with mock.patch("builtins.print") as missing:
|
||
agent.show_call({"name": "poll", "arguments": json.dumps({"id": 9})})
|
||
agent.show_result("[missing] no agent 9", 5)
|
||
snap = "\n".join(str(call.args[0]) for call in shown.call_args_list if call.args)
|
||
self.assertIn("✦", snap)
|
||
self.assertIn("lanes", snap)
|
||
self.assertIn("#3", snap)
|
||
self.assertIn("#4", snap)
|
||
done = "\n".join(str(call.args[0]) for call in collected.call_args_list if call.args)
|
||
self.assertIn("⟳", done)
|
||
self.assertIn("all collected", done)
|
||
lone = "\n".join(str(call.args[0]) for call in quiet.call_args_list if call.args)
|
||
self.assertNotIn("lanes", lone)
|
||
gone = "\n".join(str(call.args[0]) for call in missing.call_args_list if call.args)
|
||
self.assertIn("✗", gone)
|
||
|
||
def test_show_lanes_pending_and_boarded(self):
|
||
agent = tai.Agent(self.config, self.store, persist=False, quiet=False)
|
||
leftover = [{"id": 2, "task": "translate p2", "status": "done", "elapsed": 9, "collected": False}]
|
||
picked = [{"id": 2, "task": "translate p2", "status": "done", "elapsed": 9, "collected": True}]
|
||
pair = [{"id": 3, "task": "t3", "status": "running", "elapsed": 5, "collected": False},
|
||
{"id": 4, "task": "t4", "status": "running", "elapsed": 5, "collected": False}]
|
||
with mock.patch.object(sys.stdout, "isatty", return_value=True):
|
||
with mock.patch.dict(os.environ, {"NO_COLOR": ""}):
|
||
with mock.patch("builtins.print") as pending:
|
||
agent.show_call({"name": "poll", "arguments": json.dumps({"id": 1})})
|
||
with mock.patch("tai.list_agents", return_value=leftover):
|
||
agent.show_result("[done] translated text", 5)
|
||
with mock.patch("builtins.print") as clean:
|
||
agent.show_call({"name": "poll", "arguments": json.dumps({"id": 2})})
|
||
with mock.patch("tai.list_agents", return_value=picked):
|
||
agent.show_result("[done] translated text", 5)
|
||
agent.tools.board_shown = True
|
||
with mock.patch("builtins.print") as boarded:
|
||
agent.show_call({"name": "poll", "arguments": json.dumps({"id": 3})})
|
||
with mock.patch("tai.list_agents", return_value=pair):
|
||
agent.show_result("[running] agent 3 still running (5s elapsed)", 5000)
|
||
held = "\n".join(str(call.args[0]) for call in pending.call_args_list if call.args)
|
||
self.assertIn("to collect", held)
|
||
self.assertNotIn("all collected", held)
|
||
swept = "\n".join(str(call.args[0]) for call in clean.call_args_list if call.args)
|
||
self.assertIn("all collected", swept)
|
||
live = "\n".join(str(call.args[0]) for call in boarded.call_args_list if call.args)
|
||
self.assertNotIn("lanes", live)
|
||
self.assertFalse(agent.tools.board_shown)
|
||
|
||
def test_agents_command_lanes(self):
|
||
agent = tai.Agent(self.config, self.store, persist=False, quiet=False)
|
||
records = [{"id": 7, "task": "translate p3", "status": "running", "elapsed": 9}]
|
||
with mock.patch.object(sys.stdout, "isatty", return_value=True):
|
||
with mock.patch.dict(os.environ, {"NO_COLOR": ""}):
|
||
with mock.patch("tai.list_agents", return_value=records):
|
||
with mock.patch("builtins.print") as shown:
|
||
tai.handle_command(agent, "/agents")
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list)
|
||
self.assertIn("agents", printed)
|
||
self.assertIn("#7", printed)
|
||
self.assertIn("translate p3", printed)
|
||
|
||
def test_statusbar_segments(self):
|
||
bar = tai.Glow.statusbar("9.9.9", "work", os.path.expanduser("~/proj"), yolo=True, agents=2, width=200)
|
||
plain = tai.strip_ansi(bar)
|
||
for bit in ("tai 9.9.9", "work", "~/proj", "YOLO", "+2"):
|
||
self.assertIn(bit, plain)
|
||
self.assertIn(tai.FG_YOLO, bar)
|
||
slim = tai.strip_ansi(tai.Glow.statusbar("9.9.9", "work", "/tmp", width=200))
|
||
self.assertNotIn("YOLO", slim)
|
||
self.assertIn("/tmp", slim)
|
||
compact = tai.strip_ansi(tai.Glow.statusbar("1.24.0", "t", "/tmp", yolo=True, agents=2, width=34))
|
||
self.assertEqual(compact, "tai 1.24.0 · t · YOLO · +2")
|
||
tiny = tai.strip_ansi(tai.Glow.statusbar("1.24.0", "t", "/tmp", yolo=True, agents=2, width=20))
|
||
self.assertFalse(tiny.endswith("·") or tiny.endswith(" "))
|
||
|
||
def test_prompt_and_hairline(self):
|
||
self.assertEqual(tai.strip_ansi(tai.Glow.prompt_text()), "› ")
|
||
self.assertIn(tai.FG_ORANGE, tai.Glow.prompt_text())
|
||
self.assertEqual(tai.strip_ansi(tai.Glow.hairline(10)), "─" * 10)
|
||
|
||
def test_render_diff_glow_header_and_rows(self):
|
||
out = tai.render_diff("a.py", "x = 1\n", "x = 2\n", width=60, colors=True)
|
||
self.assertTrue(tai.strip_ansi(out.splitlines()[0]).startswith("◆ Edited a.py (+1 -1)"))
|
||
self.assertIn(tai.FG_ADD, out)
|
||
self.assertIn(tai.FG_DEL, out)
|
||
|
||
def test_show_call_result_glow_branch(self):
|
||
agent = tai.Agent(self.config, self.store, persist=False, quiet=False)
|
||
self.store.save_secret("api", "known-token-5")
|
||
with mock.patch.object(sys.stdout, "isatty", return_value=True):
|
||
with mock.patch.dict(os.environ, {"NO_COLOR": ""}):
|
||
with mock.patch("builtins.print") as shown:
|
||
agent.show_call({"name": "shell", "arguments": json.dumps({"command": "curl known-token-5"})})
|
||
agent.show_call({"name": "read_file", "arguments": json.dumps({"path": "/tmp/x"})})
|
||
agent.show_result("ok", 250)
|
||
with mock.patch("builtins.print") as failed:
|
||
agent.show_call({"name": "shell", "arguments": "{}"})
|
||
agent.show_result("error: boom", 5)
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list if call.args)
|
||
self.assertIn("▶", printed)
|
||
self.assertIn("◇", printed)
|
||
self.assertIn("◆", printed)
|
||
self.assertIn("▌", printed)
|
||
self.assertIn("✓", printed)
|
||
self.assertNotIn('{"command"', printed)
|
||
self.assertNotIn("known-token-5", printed)
|
||
self.assertIn("[redacted:api]", printed)
|
||
busted = "\n".join(str(call.args[0]) for call in failed.call_args_list if call.args)
|
||
self.assertIn("✗", busted)
|
||
|
||
|
||
class LaneBoardTests(unittest.TestCase):
|
||
def test_available(self):
|
||
loud = mock.Mock(quiet=False)
|
||
hushed = mock.Mock(quiet=True)
|
||
with mock.patch.object(sys.stdout, "isatty", return_value=False):
|
||
self.assertFalse(tai.LaneBoard.available(loud))
|
||
with mock.patch.object(sys.stdout, "isatty", return_value=True):
|
||
with mock.patch.dict(os.environ, {"NO_COLOR": ""}):
|
||
self.assertTrue(tai.LaneBoard.available(loud))
|
||
self.assertFalse(tai.LaneBoard.available(hushed))
|
||
with mock.patch.dict(os.environ, {"NO_COLOR": "1"}):
|
||
self.assertFalse(tai.LaneBoard.available(loud))
|
||
|
||
def test_paint_rewrites_in_place(self):
|
||
records = [{"id": 1, "task": "translate p1", "status": "running", "elapsed": 5, "collected": False}]
|
||
buf = io.StringIO()
|
||
board = tai.LaneBoard("t", width=100, file=buf)
|
||
with mock.patch("tai.list_agents", return_value=records):
|
||
board._paint("")
|
||
first = buf.getvalue()
|
||
self.assertNotIn("\033[2A", first)
|
||
self.assertIn("1 running", tai.strip_ansi(first))
|
||
board._paint(tai.Ansi.BOLD)
|
||
second = buf.getvalue()
|
||
self.assertIn("\033[2A", second)
|
||
self.assertIn(tai.Ansi.BOLD + tai.Ansi.YELLOW + "●", second)
|
||
self.assertEqual(board.painted, 2)
|
||
|
||
def test_count_change_degrades_gracefully(self):
|
||
one = [{"id": 1, "task": "t1", "status": "running", "elapsed": 5, "collected": False}]
|
||
two = one + [{"id": 2, "task": "t2", "status": "running", "elapsed": 5, "collected": False}]
|
||
buf = io.StringIO()
|
||
board = tai.LaneBoard("t", width=100, file=buf)
|
||
with mock.patch("tai.list_agents", return_value=one):
|
||
board._paint("")
|
||
mark = len(buf.getvalue())
|
||
with mock.patch("tai.list_agents", return_value=two):
|
||
board._paint("")
|
||
fresh = buf.getvalue()[mark:]
|
||
self.assertNotIn("[A", fresh)
|
||
self.assertIn("#2", fresh)
|
||
self.assertEqual(board.painted, 3)
|
||
|
||
def test_start_stop_lifecycle(self):
|
||
records = [{"id": 1, "task": "translate p1", "status": "running", "elapsed": 5, "collected": False}]
|
||
buf = io.StringIO()
|
||
board = tai.LaneBoard("t", width=100, file=buf)
|
||
with mock.patch("tai.list_agents", return_value=records):
|
||
board.start()
|
||
time.sleep(0.5)
|
||
board.stop()
|
||
self.assertIsNone(board.worker)
|
||
self.assertGreaterEqual(buf.getvalue().count("lanes · 1 running"), 2)
|
||
|
||
|
||
class ResumeTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "resume-test-1"))
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def text_reply(self, text="done"):
|
||
return {"role": "assistant", "content": text, "reasoning": "", "tool_calls": [], "backend": "x"}
|
||
|
||
def test_atomic_session_survives_failed_replace(self):
|
||
first = [{"role": "user", "content": "first"}]
|
||
self.store.save_session("t", first)
|
||
with mock.patch("os.replace", side_effect=OSError("disk full")):
|
||
self.store.save_session("t", [{"role": "user", "content": "second"}])
|
||
self.assertEqual([item["content"] for item in self.store.load_session("t")], ["first"])
|
||
|
||
def test_corrupt_session_quarantined(self):
|
||
path = self.store.session_path("t")
|
||
self.store.save_session("t", [{"role": "user", "content": "good"}])
|
||
with open(path, "w", encoding="utf-8") as handle:
|
||
handle.write("{not json")
|
||
self.assertEqual(self.store.load_session("t"), [])
|
||
leftovers = [entry for entry in os.listdir(self.config.profiles_dir) if ".corrupt-" in entry]
|
||
self.assertEqual(len(leftovers), 1)
|
||
with open(os.path.join(self.config.profiles_dir, leftovers[0]), encoding="utf-8") as handle:
|
||
self.assertEqual(handle.read(), "{not json")
|
||
|
||
def test_normal_turn_closes_state(self):
|
||
agent = tai.Agent(self.config, self.store, persist=True, quiet=True)
|
||
with mock.patch.object(agent.chat, "complete", return_value=self.text_reply()):
|
||
agent.run_turn("hello", capture=True)
|
||
state = self.store.load_turn_state("t", "main")
|
||
self.assertFalse(state["open"])
|
||
self.assertEqual(state["goal"], "hello")
|
||
|
||
def test_crash_resume_end_to_end(self):
|
||
agent = tai.Agent(self.config, self.store, persist=True, quiet=True)
|
||
tool_reply = {"role": "assistant", "content": "", "reasoning": "", "tool_calls": [{"id": "1", "name": "shell", "arguments": json.dumps({"command": "echo resume-mark"})}], "backend": "x"}
|
||
with mock.patch.object(agent.chat, "complete", side_effect=[dict(tool_reply), KeyboardInterrupt()]):
|
||
with self.assertRaises(KeyboardInterrupt):
|
||
agent.run_turn("do the thing", capture=True)
|
||
state = self.store.load_turn_state("t", "main")
|
||
self.assertTrue(state["open"])
|
||
self.assertEqual(state["goal"], "do the thing")
|
||
self.store.close()
|
||
store2 = tai.Store(self.config, tai.Seal(self.config.home, "resume-test-1"))
|
||
self.store = store2
|
||
agent2 = tai.Agent(self.config, store2, persist=True, quiet=True)
|
||
texts = [item.get("content") or "" for item in agent2.messages]
|
||
self.assertTrue(any("do the thing" in text for text in texts))
|
||
self.assertTrue(any("resume-mark" in text for text in texts))
|
||
with mock.patch.object(agent2.chat, "complete", return_value=self.text_reply("finished")):
|
||
with mock.patch("builtins.print"):
|
||
answer = tai.resume_turn(agent2)
|
||
self.assertEqual(answer, "finished")
|
||
self.assertFalse(store2.load_turn_state("t", "main")["open"])
|
||
|
||
def test_resume_nothing_open(self):
|
||
agent = tai.Agent(self.config, self.store, persist=True, quiet=True)
|
||
with mock.patch.object(agent.chat, "complete") as called:
|
||
with mock.patch("builtins.print") as shown:
|
||
self.assertEqual(tai.resume_turn(agent), "")
|
||
called.assert_not_called()
|
||
printed = "\n".join(str(call.args[0]) for call in shown.call_args_list)
|
||
self.assertIn("nothing to resume", printed)
|
||
|
||
def test_checkpoint_idle_writes_no_state(self):
|
||
agent = tai.Agent(self.config, self.store, persist=True, quiet=True)
|
||
agent.checkpoint()
|
||
self.assertEqual(self.store.load_turn_state("t", "main"), {})
|
||
self.assertFalse(os.path.exists(self.store.turn_state_path("t", "main")))
|
||
|
||
def test_parser_continue(self):
|
||
args = tai.build_parser().parse_args(["--continue"])
|
||
self.assertTrue(args.resume)
|
||
self.assertFalse(tai.build_parser().parse_args([]).resume)
|
||
|
||
|
||
class DenialTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = False
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, ""))
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def call_reply(self, *calls):
|
||
return {"role": "assistant", "content": "", "reasoning": "", "tool_calls": list(calls), "backend": "x"}
|
||
|
||
def test_deny_continues_with_guidance(self):
|
||
agent = tai.Agent(self.config, self.store, persist=True, quiet=True)
|
||
first = self.call_reply({"id": "c1", "name": "shell", "arguments": json.dumps({"command": "ssh evil.example.com"})})
|
||
final = {"role": "assistant", "content": "done", "reasoning": "", "tool_calls": [], "backend": "x"}
|
||
with mock.patch.object(agent.chat, "complete", side_effect=[first, final]) as completer:
|
||
with mock.patch("sys.stdin") as fake_stdin:
|
||
fake_stdin.isatty.return_value = True
|
||
with mock.patch("builtins.input", side_effect=["n", "do X instead"]) as asker:
|
||
with mock.patch("builtins.print"):
|
||
result = agent.run_turn("try ssh", capture=True)
|
||
self.assertEqual(result, "done")
|
||
self.assertEqual(asker.call_count, 2)
|
||
second_messages = completer.call_args_list[1].args[0]
|
||
users = [item["content"] for item in second_messages if item["role"] == "user"]
|
||
self.assertIn("do X instead", users)
|
||
tools = [item for item in second_messages if item["role"] == "tool"]
|
||
self.assertEqual(tools[0]["content"], "denied by user")
|
||
|
||
def test_deny_empty_aborts(self):
|
||
agent = tai.Agent(self.config, self.store, persist=True, quiet=True)
|
||
first = self.call_reply(
|
||
{"id": "c1", "name": "shell", "arguments": json.dumps({"command": "ssh a.example.com"})},
|
||
{"id": "c2", "name": "shell", "arguments": json.dumps({"command": "ssh b.example.com"})},
|
||
)
|
||
with mock.patch.object(agent.chat, "complete", return_value=first) as completer:
|
||
with mock.patch("sys.stdin") as fake_stdin:
|
||
fake_stdin.isatty.return_value = True
|
||
with mock.patch("builtins.input", side_effect=["n", ""]):
|
||
with mock.patch("builtins.print"):
|
||
result = agent.run_turn("try ssh", capture=True)
|
||
self.assertEqual(result, "stopped by user")
|
||
self.assertEqual(completer.call_count, 1)
|
||
by_id = {}
|
||
for item in agent.messages:
|
||
for recorded in item.get("tool_calls") or []:
|
||
by_id[recorded["id"]] = None
|
||
for item in agent.messages:
|
||
if item.get("role") == "tool":
|
||
by_id[item["tool_call_id"]] = item["content"]
|
||
self.assertEqual(by_id, {"c1": "denied by user", "c2": "skipped: stopped after denial"})
|
||
|
||
def test_auto_never_asks(self):
|
||
agent = tai.Agent(self.config, self.store, persist=True, quiet=True, auto=True)
|
||
first = self.call_reply({"id": "c1", "name": "shell", "arguments": json.dumps({"command": "echo auto-ok"})})
|
||
final = {"role": "assistant", "content": "done", "reasoning": "", "tool_calls": [], "backend": "x"}
|
||
with mock.patch.object(agent.chat, "complete", side_effect=[first, final]):
|
||
with mock.patch("sys.stdin") as fake_stdin:
|
||
fake_stdin.isatty.return_value = True
|
||
with mock.patch("builtins.input", side_effect=AssertionError("must not ask")) as asker:
|
||
with mock.patch("builtins.print"):
|
||
result = agent.run_turn("try echo", capture=True)
|
||
self.assertEqual(result, "done")
|
||
self.assertEqual(asker.call_count, 0)
|
||
ran = [item for item in agent.messages if item.get("role") == "tool"]
|
||
self.assertIn("auto-ok", ran[0]["content"])
|
||
|
||
def test_dispatch_reraises_denied(self):
|
||
tools = tai.Tools(FakeApp())
|
||
|
||
def raiser(args):
|
||
raise tai.Denied("go left")
|
||
|
||
tools.handlers["boom"] = raiser
|
||
with self.assertRaises(tai.Denied) as caught:
|
||
tools.dispatch("boom", "{}")
|
||
self.assertEqual(caught.exception.guidance, "go left")
|
||
|
||
def test_secret_delete_passes_no_guidance(self):
|
||
agent = mock.Mock()
|
||
seen = []
|
||
|
||
def approver(command, guidance=True):
|
||
seen.append(guidance)
|
||
return False
|
||
|
||
agent.ask_approval = approver
|
||
agent.store = self.store
|
||
agent.profile = "t"
|
||
self.store.save_secret("wifi", "repl-value-9")
|
||
with mock.patch("builtins.print"):
|
||
tai.handle_command(agent, "/secret delete wifi")
|
||
self.assertEqual(seen, [False])
|
||
self.assertEqual(self.store.load_secret("wifi"), "repl-value-9")
|
||
|
||
def test_show_call_masks_secrets(self):
|
||
agent = tai.Agent(self.config, self.store, persist=False, quiet=False)
|
||
self.store.save_secret("api", "known-token-5")
|
||
with mock.patch("builtins.print") as printer:
|
||
agent.show_call({"name": "store_secret", "arguments": json.dumps({"name": "api", "value": "brand-new-1"})})
|
||
agent.show_call({"name": "shell", "arguments": json.dumps({"command": "curl known-token-5"})})
|
||
shown = "\n".join(str(call.args[0]) for call in printer.call_args_list if call.args)
|
||
self.assertNotIn("brand-new-1", shown)
|
||
self.assertIn("[hidden]", shown)
|
||
self.assertNotIn("known-token-5", shown)
|
||
self.assertIn("[redacted:api]", shown)
|
||
|
||
|
||
class FileGuardTests(unittest.TestCase):
|
||
def test_write_requires_read(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
target = os.path.join(tmp, "notes.txt")
|
||
with open(target, "w", encoding="utf-8") as handle:
|
||
handle.write("original one")
|
||
tools = tai.Tools(FakeApp())
|
||
denied = tools.dispatch("write_file", json.dumps({"path": target, "content": "clobber"}))
|
||
self.assertIn("read it first", denied)
|
||
with open(target, encoding="utf-8") as handle:
|
||
self.assertEqual(handle.read(), "original one")
|
||
self.assertIn("original one", tools.dispatch("read_file", json.dumps({"path": target})))
|
||
self.assertIn("wrote", tools.dispatch("write_file", json.dumps({"path": target, "content": "updated two"})))
|
||
with open(target, encoding="utf-8") as handle:
|
||
self.assertEqual(handle.read(), "updated two")
|
||
|
||
def test_new_files_always_writable(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
target = os.path.join(tmp, "sub", "fresh.txt")
|
||
tools = tai.Tools(FakeApp())
|
||
self.assertIn("wrote", tools.dispatch("write_file", json.dumps({"path": target, "content": "hello"})))
|
||
self.assertIn("edited", tools.dispatch("edit_file", json.dumps({"path": target, "find": "hello", "replace": "hi"})))
|
||
|
||
def test_edit_requires_read(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
target = os.path.join(tmp, "code.py")
|
||
with open(target, "w", encoding="utf-8") as handle:
|
||
handle.write("print(1)")
|
||
tools = tai.Tools(FakeApp())
|
||
self.assertIn("read it first", tools.dispatch("edit_file", json.dumps({"path": target, "find": "1", "replace": "2"})))
|
||
tools.dispatch("read_file", json.dumps({"path": target}))
|
||
self.assertIn("edited", tools.dispatch("edit_file", json.dumps({"path": target, "find": "1", "replace": "2"})))
|
||
|
||
def test_guard_is_per_session(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
target = os.path.join(tmp, "data.txt")
|
||
with open(target, "w", encoding="utf-8") as handle:
|
||
handle.write("v1")
|
||
first = tai.Tools(FakeApp())
|
||
first.dispatch("read_file", json.dumps({"path": target}))
|
||
second = tai.Tools(FakeApp())
|
||
self.assertIn("read it first", second.dispatch("write_file", json.dumps({"path": target, "content": "v2"})))
|
||
|
||
|
||
class FakeBot:
|
||
def __init__(self):
|
||
self.sent = []
|
||
self.actions = []
|
||
|
||
def call(self, method, params, timeout=70):
|
||
if method == "sendMessage":
|
||
self.sent.append(params["text"])
|
||
if method == "sendChatAction":
|
||
self.actions.append(params["action"])
|
||
return {"ok": True}
|
||
|
||
|
||
class FakeAgent:
|
||
def __init__(self):
|
||
self.reset = 0
|
||
self.turns = []
|
||
|
||
def reset_history(self):
|
||
self.reset += 1
|
||
|
||
def run_turn(self, text, capture=False):
|
||
self.turns.append((text, capture))
|
||
return "canned reply"
|
||
|
||
|
||
class TelegramTests(unittest.TestCase):
|
||
def test_send_chunks(self):
|
||
bot = FakeBot()
|
||
tai.telegram_send(bot, 7, "x" * 5000)
|
||
self.assertEqual([len(part) for part in bot.sent], [4000, 1000])
|
||
|
||
def test_start_and_new(self):
|
||
agent = FakeAgent()
|
||
bot = FakeBot()
|
||
tai.handle_telegram_update(agent, bot, {"message": {"chat": {"id": 1}, "text": "/start"}})
|
||
tai.handle_telegram_update(agent, bot, {"message": {"chat": {"id": 1}, "text": "/new"}})
|
||
self.assertEqual(agent.reset, 1)
|
||
self.assertEqual(len(bot.sent), 2)
|
||
self.assertEqual(agent.turns, [])
|
||
|
||
def test_text_turn(self):
|
||
agent = FakeAgent()
|
||
bot = FakeBot()
|
||
tai.handle_telegram_update(agent, bot, {"message": {"chat": {"id": 1}, "text": "hello"}})
|
||
self.assertEqual(agent.turns, [("hello", True)])
|
||
self.assertEqual(bot.sent, ["canned reply"])
|
||
self.assertEqual(bot.actions, ["typing"])
|
||
|
||
def test_token_loading(self):
|
||
with tempfile.TemporaryDirectory() as home:
|
||
with mock.patch.dict(os.environ, {}, clear=False):
|
||
os.environ.pop("TELEGRAM_BOT_TOKEN", None)
|
||
self.assertEqual(tai.load_telegram_token(home), "")
|
||
with open(os.path.join(home, "telegram.env"), "w", encoding="utf-8") as handle:
|
||
handle.write("TELEGRAM_BOT_TOKEN=file-token-1\n")
|
||
self.assertEqual(tai.load_telegram_token(home), "file-token-1")
|
||
os.environ["TELEGRAM_BOT_TOKEN"] = "env-token-2"
|
||
self.assertEqual(tai.load_telegram_token(home), "env-token-2")
|
||
os.environ.pop("TELEGRAM_BOT_TOKEN", None)
|
||
|
||
|
||
class InstallTests(unittest.TestCase):
|
||
def test_bashrc_upsert(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
path = os.path.join(tmp, ".bashrc")
|
||
with open(path, "w", encoding="utf-8") as handle:
|
||
handle.write("export PATH=$PATH:/x\n")
|
||
self.assertTrue(tai.upsert_bashrc_block(path))
|
||
self.assertFalse(tai.upsert_bashrc_block(path))
|
||
with open(path, encoding="utf-8") as handle:
|
||
content = handle.read()
|
||
self.assertIn("export PATH=$PATH:/x", content)
|
||
self.assertEqual(content.count(tai.BASHRC_MARK_BEGIN), 1)
|
||
with open(path + ".bak-tai", encoding="utf-8") as handle:
|
||
self.assertNotIn("tai command-not-found", handle.read())
|
||
|
||
def test_bashrc_replaces_stale(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
path = os.path.join(tmp, ".bashrc")
|
||
stale = tai.BASHRC_MARK_BEGIN + "\nold hook\n" + tai.BASHRC_MARK_END + "\n"
|
||
with open(path, "w", encoding="utf-8") as handle:
|
||
handle.write("alias x=y\n" + stale)
|
||
self.assertTrue(tai.upsert_bashrc_block(path))
|
||
with open(path, encoding="utf-8") as handle:
|
||
content = handle.read()
|
||
self.assertNotIn("old hook", content)
|
||
self.assertIn("alias x=y", content)
|
||
self.assertEqual(content.count(tai.BASHRC_MARK_BEGIN), 1)
|
||
|
||
def test_shell_rc_path_bash_zsh(self):
|
||
with mock.patch.dict(os.environ, {"HOME": "/home/x"}, clear=False):
|
||
self.assertEqual(tai.shell_rc_path("bash"), "/home/x/.bashrc")
|
||
self.assertEqual(tai.shell_rc_path("zsh"), "/home/x/.zshrc")
|
||
|
||
def test_shell_rc_path_fish_respects_xdg_config_home(self):
|
||
with mock.patch.dict(os.environ, {"HOME": "/home/x"}, clear=False):
|
||
os.environ.pop("XDG_CONFIG_HOME", None)
|
||
self.assertEqual(tai.shell_rc_path("fish"), "/home/x/.config/fish/config.fish")
|
||
with mock.patch.dict(os.environ, {"XDG_CONFIG_HOME": "/custom/cfg"}):
|
||
self.assertEqual(tai.shell_rc_path("fish"), "/custom/cfg/fish/config.fish")
|
||
|
||
def test_fish_hook_creates_missing_config_directory(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
path = os.path.join(tmp, "fish", "config.fish")
|
||
self.assertFalse(os.path.isdir(os.path.dirname(path)))
|
||
self.assertTrue(tai.upsert_fish_hook_block(path))
|
||
with open(path, encoding="utf-8") as handle:
|
||
content = handle.read()
|
||
self.assertIn("function fish_command_not_found", content)
|
||
self.assertFalse(tai.upsert_fish_hook_block(path))
|
||
|
||
def test_zsh_fish_hook_roundtrip_preserves_other_content(self):
|
||
for shell, upsert, remove, present, needle in (
|
||
("zsh", tai.upsert_zsh_hook_block, tai.remove_zsh_hook_block, tai.zsh_hook_present, "command_not_found_handler"),
|
||
("fish", tai.upsert_fish_hook_block, tai.remove_fish_hook_block, tai.fish_hook_present, "fish_command_not_found"),
|
||
):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
path = os.path.join(tmp, "rc")
|
||
with open(path, "w", encoding="utf-8") as handle:
|
||
handle.write("# pre-existing user content\nexport FOO=bar\n")
|
||
self.assertTrue(upsert(path))
|
||
self.assertTrue(present(path))
|
||
with open(path, encoding="utf-8") as handle:
|
||
content = handle.read()
|
||
self.assertIn("export FOO=bar", content)
|
||
self.assertIn(needle, content)
|
||
self.assertTrue(remove(path))
|
||
self.assertFalse(present(path))
|
||
with open(path, encoding="utf-8") as handle:
|
||
content = handle.read()
|
||
self.assertIn("export FOO=bar", content)
|
||
self.assertNotIn(needle, content, "shell=%s" % shell)
|
||
|
||
def test_parser_prompt(self):
|
||
args = tai.build_parser().parse_args(["what", "is", "(2+3)?"])
|
||
self.assertEqual(args.prompt, ["what", "is", "(2+3)?"])
|
||
self.assertFalse(args.install)
|
||
args = tai.build_parser().parse_args(["--install-telegram"])
|
||
self.assertTrue(args.install_telegram)
|
||
self.assertEqual(args.prompt, [])
|
||
args = tai.build_parser().parse_args(["--profile", "work", "--yes"])
|
||
self.assertEqual(args.profile, "work")
|
||
self.assertTrue(args.yes)
|
||
|
||
|
||
class FakePrincipal:
|
||
def __init__(self):
|
||
self.depth = 0
|
||
self.profile = "t"
|
||
self.config = None
|
||
self.store = mock.Mock()
|
||
self.store.redact = lambda text, profile=None: text
|
||
self.runner_override = None
|
||
|
||
|
||
class OrchestrationTests(unittest.TestCase):
|
||
def setUp(self):
|
||
with tai.AGENTS_LOCK:
|
||
tai.AGENTS.clear()
|
||
tai.AGENTS_NEXT[0] = 1
|
||
|
||
def tearDown(self):
|
||
with tai.AGENTS_LOCK:
|
||
tai.AGENTS.clear()
|
||
tai.AGENTS_NEXT[0] = 1
|
||
|
||
def test_spawn_poll_instant(self):
|
||
agent_id = tai.spawn_agent("do it", "t", 60, None, None, 0, lambda task, profile, timeout: "stub-done")
|
||
self.assertEqual(agent_id, 1)
|
||
status, text = tai.poll_agent(agent_id, wait=5)
|
||
self.assertEqual(status, "done")
|
||
self.assertIn("stub-done", text)
|
||
|
||
def test_poll_tool_marks_collected(self):
|
||
tools = tai.Tools(FakePrincipal())
|
||
first = tai.spawn_agent("quick", "t", 60, None, None, 0, lambda task, profile, timeout: "r")
|
||
self.assertIn("[done]", tools.dispatch("poll", {"id": first, "wait": 5}))
|
||
|
||
def broken(task, profile, timeout):
|
||
raise RuntimeError("boom")
|
||
|
||
bad = tai.spawn_agent("bad", "t", 60, None, None, 0, broken)
|
||
self.assertIn("[error]", tools.dispatch("poll", {"id": bad, "wait": 5}))
|
||
with tai.AGENTS_LOCK:
|
||
self.assertTrue(tai.AGENTS[first]["collected"])
|
||
self.assertTrue(tai.AGENTS[bad]["collected"])
|
||
|
||
def test_poll_running_marks_nothing(self):
|
||
tools = tai.Tools(FakePrincipal())
|
||
agent_id = tai.spawn_agent("slow", "t", 60, None, None, 0, lambda task, profile, timeout: time.sleep(2) or "slow-done")
|
||
self.assertIn("[running]", tools.dispatch("poll", {"id": agent_id}))
|
||
with tai.AGENTS_LOCK:
|
||
self.assertFalse(tai.AGENTS[agent_id]["collected"])
|
||
|
||
def test_poll_wait_runs_lane_board(self):
|
||
tools = tai.Tools(FakePrincipal())
|
||
agent_id = tai.spawn_agent("slow", "t", 60, None, None, 0, lambda task, profile, timeout: time.sleep(2) or "slow-done")
|
||
with mock.patch.object(sys.stdout, "isatty", return_value=True):
|
||
with mock.patch.dict(os.environ, {"NO_COLOR": ""}):
|
||
with mock.patch("tai.LaneBoard") as board_cls:
|
||
result = tools.dispatch("poll", {"id": agent_id, "wait": 1})
|
||
self.assertIn("[running]", result)
|
||
board_cls.return_value.start.assert_called_once_with()
|
||
board_cls.return_value.stop.assert_called_once_with()
|
||
self.assertTrue(tools.board_shown)
|
||
|
||
def test_poll_running_and_missing(self):
|
||
def slow(task, profile, timeout):
|
||
time.sleep(2)
|
||
return "slow-done"
|
||
|
||
agent_id = tai.spawn_agent("slow", "t", 60, None, None, 0, slow)
|
||
status, text = tai.poll_agent(agent_id, wait=0)
|
||
self.assertEqual(status, "running")
|
||
self.assertIn("still running", text)
|
||
status, _text = tai.poll_agent(999, wait=0)
|
||
self.assertEqual(status, "missing")
|
||
status, text = tai.poll_agent(agent_id, wait=5)
|
||
self.assertEqual(status, "done")
|
||
|
||
def test_timeout_and_error_status(self):
|
||
slow_id = tai.spawn_agent("t", "t", 60, None, None, 0, lambda task, profile, timeout: "partial\n[time limit reached]")
|
||
self.assertEqual(tai.poll_agent(slow_id, wait=5)[0], "timeout")
|
||
|
||
def broken(task, profile, timeout):
|
||
raise RuntimeError("boom")
|
||
|
||
bad_id = tai.spawn_agent("t", "t", 60, None, None, 0, broken)
|
||
status, text = tai.poll_agent(bad_id, wait=5)
|
||
self.assertEqual(status, "error")
|
||
self.assertIn("boom", text)
|
||
|
||
def test_list_and_clear(self):
|
||
tai.spawn_agent("one", "t", 60, None, None, 0, lambda task, profile, timeout: "r1")
|
||
tai.spawn_agent("two", "t", 60, None, None, 0, lambda task, profile, timeout: time.sleep(2) or "r2")
|
||
self.assertEqual(tai.poll_agent(1, wait=5)[0], "done")
|
||
records = tai.list_agents()
|
||
self.assertEqual(len(records), 2)
|
||
self.assertEqual(tai.clear_agents(), 1)
|
||
self.assertEqual(len(tai.list_agents()), 1)
|
||
|
||
def test_fork_tool(self):
|
||
app = FakePrincipal()
|
||
app.runner_override = lambda task, profile, timeout: "forked-ok"
|
||
tools = tai.Tools(app)
|
||
started = tools.dispatch("fork", json.dumps({"task": "research x"}))
|
||
self.assertIn("agent 1 started", started)
|
||
result = tools.dispatch("poll", json.dumps({"id": 1, "wait": 5}))
|
||
self.assertIn("forked-ok", result)
|
||
|
||
def test_fork_depth_limit(self):
|
||
app = FakePrincipal()
|
||
app.depth = 2
|
||
tools = tai.Tools(app)
|
||
self.assertIn("depth limit", tools.dispatch("fork", json.dumps({"task": "x"})))
|
||
|
||
def test_reasoning_only_becomes_result(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
old_home = os.environ.get("TAI_HOME")
|
||
os.environ["TAI_HOME"] = tmp
|
||
try:
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
config = tai.Config(FakeArgs())
|
||
store = tai.Store(config, tai.Seal(config.home, ""))
|
||
agent = tai.Agent(config, store, persist=False, quiet=True)
|
||
reply = {"role": "assistant", "content": "", "reasoning": "thought out", "tool_calls": [], "backend": "x"}
|
||
with mock.patch.object(agent.chat, "complete", return_value=reply):
|
||
self.assertEqual(agent.run_turn("hi", capture=True), "thought out")
|
||
store.close()
|
||
finally:
|
||
if old_home is None:
|
||
os.environ.pop("TAI_HOME", None)
|
||
else:
|
||
os.environ["TAI_HOME"] = old_home
|
||
|
||
def test_deadline_shortcircuit(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
old_home = os.environ.get("TAI_HOME")
|
||
os.environ["TAI_HOME"] = tmp
|
||
try:
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
config = tai.Config(FakeArgs())
|
||
store = tai.Store(config, tai.Seal(config.home, ""))
|
||
agent = tai.Agent(config, store, persist=False, quiet=True)
|
||
agent.deadline = time.time() - 1
|
||
with mock.patch.object(agent.chat, "complete", side_effect=AssertionError("network used")):
|
||
result = agent.run_turn("hi", capture=True)
|
||
self.assertTrue(agent.timed_out)
|
||
self.assertIn("time limit reached", result)
|
||
store.close()
|
||
finally:
|
||
if old_home is None:
|
||
os.environ.pop("TAI_HOME", None)
|
||
else:
|
||
os.environ["TAI_HOME"] = old_home
|
||
|
||
|
||
class MarkdownTests(unittest.TestCase):
|
||
def test_passthrough_without_tty(self):
|
||
with mock.patch("sys.stdout") as fake_out:
|
||
fake_out.isatty.return_value = False
|
||
self.assertEqual(tai.render_markdown("**x**"), "**x**")
|
||
with mock.patch("sys.stdout") as fake_out:
|
||
fake_out.isatty.return_value = True
|
||
with mock.patch.dict(os.environ, {"NO_COLOR": "1"}):
|
||
self.assertEqual(tai.render_markdown("**x**"), "**x**")
|
||
|
||
def test_headings(self):
|
||
out = tai.render_markdown("# One\n\n## Two\n\n### Three", width=80, color=True)
|
||
self.assertNotIn("#", out)
|
||
self.assertIn(tai.Ansi.BOLD + tai.Ansi.CYAN + "One", out)
|
||
self.assertIn(tai.Ansi.BOLD + "Two", out)
|
||
self.assertIn("Three", tai.strip_ansi(out))
|
||
|
||
def test_inline_styles(self):
|
||
out = tai.render_markdown("**b** *i* `c` ~~s~~ _u_ __w__", width=80, color=True)
|
||
self.assertIn(tai.style_text("b", tai.Ansi.BOLD), out)
|
||
self.assertIn(tai.style_text("i", tai.Ansi.DIM), out)
|
||
self.assertIn(tai.style_text("c", tai.Ansi.CYAN), out)
|
||
self.assertIn(tai.style_text("s", tai.Ansi.STRIKE), out)
|
||
self.assertIn(tai.style_text("u", tai.Ansi.DIM), out)
|
||
self.assertIn(tai.style_text("w", tai.Ansi.BOLD), out)
|
||
self.assertNotIn("**", out)
|
||
|
||
def test_snake_case_survives(self):
|
||
out = tai.render_markdown("use my_var_name and `__init__` here", width=80, color=True)
|
||
self.assertIn("my_var_name", tai.strip_ansi(out))
|
||
self.assertIn("__init__", tai.strip_ansi(out))
|
||
self.assertNotIn(tai.Ansi.BOLD, out)
|
||
|
||
def test_link_and_image(self):
|
||
out = tai.render_markdown("[docs](https://x.example/d) and ", width=80, color=True)
|
||
self.assertIn("docs", out)
|
||
self.assertIn("https://x.example/d", out)
|
||
self.assertIn("[image: alt]", tai.strip_ansi(out))
|
||
self.assertNotIn("[docs]", out)
|
||
|
||
def test_code_fence_verbatim(self):
|
||
out = tai.render_markdown("```python\nreturn \"**x**\"\n```", width=80, color=True)
|
||
self.assertIn('"**x**"', out)
|
||
self.assertNotIn(tai.Ansi.BOLD, out)
|
||
self.assertNotIn("```", out)
|
||
|
||
def test_unclosed_fence(self):
|
||
out = tai.render_markdown("```\ncode **x**", width=80, color=True)
|
||
self.assertIn("**x**", out)
|
||
self.assertNotIn(tai.Ansi.BOLD, out)
|
||
|
||
def test_table_alignment(self):
|
||
source = "| item | qty |\n|---|---:|\n| apple | 12 |\n| fig | 3 |"
|
||
out = tai.render_markdown(source, width=80, color=True)
|
||
plain = tai.strip_ansi(out)
|
||
self.assertIn("┌───────┬─────┐", plain)
|
||
self.assertIn("│ apple │ 12 │", plain)
|
||
self.assertIn("│ fig │ 3 │", plain)
|
||
self.assertEqual({len(line) for line in plain.splitlines()}, {15})
|
||
self.assertIn(tai.Ansi.BOLD, out)
|
||
|
||
def test_table_escaped_pipe(self):
|
||
out = tai.render_markdown("| a |\n|---|\n| x\\|y |", width=80, color=True)
|
||
plain = tai.strip_ansi(out)
|
||
self.assertIn("x|y", plain)
|
||
self.assertEqual(len(plain.splitlines()), 5)
|
||
|
||
def test_lists(self):
|
||
source = "- a\n - b\n- [x] done\n- [ ] open\n1. x\n1. y"
|
||
plain = tai.strip_ansi(tai.render_markdown(source, width=80, color=True))
|
||
self.assertIn("• a", plain)
|
||
self.assertIn(" ◦ b", plain)
|
||
self.assertIn("☑ done", plain)
|
||
self.assertIn("☐ open", plain)
|
||
self.assertIn("1. x", plain)
|
||
self.assertIn("2. y", plain)
|
||
|
||
def test_quote_and_rule(self):
|
||
out = tai.render_markdown("> wise\n> words\n\n---", width=40, color=True)
|
||
plain = tai.strip_ansi(out)
|
||
self.assertIn("│ wise words", plain)
|
||
self.assertIn("─" * 40, plain)
|
||
|
||
def test_wrap_keeps_styles(self):
|
||
out = tai.render_markdown("**word " + " ".join("w%d" % n for n in range(20)) + " end**", width=30, color=True)
|
||
lines = out.splitlines()
|
||
self.assertGreater(len(lines), 1)
|
||
self.assertNotIn("**", out)
|
||
for line in lines:
|
||
self.assertLessEqual(len(tai.strip_ansi(line)), 30)
|
||
self.assertIn(tai.Ansi.BOLD, out)
|
||
|
||
def test_escapes_and_breaks(self):
|
||
out = tai.render_markdown("a \\*b\\* c<br>d", width=80, color=True)
|
||
plain = tai.strip_ansi(out)
|
||
self.assertIn("*b*", plain)
|
||
self.assertNotIn(tai.Ansi.DIM, out)
|
||
self.assertEqual(plain.splitlines(), ["a *b* c", "d"])
|
||
|
||
def test_run_turn_renders_reply(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
old_home = os.environ.get("TAI_HOME")
|
||
os.environ["TAI_HOME"] = tmp
|
||
try:
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
config = tai.Config(FakeArgs())
|
||
store = tai.Store(config, tai.Seal(config.home, ""))
|
||
agent = tai.Agent(config, store)
|
||
reply = {"role": "assistant", "content": "# Title\n\nHello **bold**", "reasoning": "", "tool_calls": [], "backend": "x"}
|
||
with mock.patch.object(agent.chat, "complete", return_value=reply):
|
||
with mock.patch.dict(os.environ):
|
||
os.environ.pop("NO_COLOR", None)
|
||
with mock.patch("sys.stdout") as fake_out:
|
||
fake_out.isatty.return_value = True
|
||
with mock.patch("builtins.print") as printer:
|
||
with mock.patch("tai.Spinner"):
|
||
result = agent.run_turn("hi")
|
||
self.assertEqual(result, "# Title\n\nHello **bold**")
|
||
printed = "\n".join(str(call.args[0]) for call in printer.call_args_list if call.args)
|
||
self.assertIn(tai.Ansi.BOLD, printed)
|
||
self.assertIn("Title", printed)
|
||
self.assertNotIn("# Title", printed)
|
||
store.close()
|
||
finally:
|
||
if old_home is None:
|
||
os.environ.pop("TAI_HOME", None)
|
||
else:
|
||
os.environ["TAI_HOME"] = old_home
|
||
|
||
def test_run_turn_prints_no_name_banner(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
old_home = os.environ.get("TAI_HOME")
|
||
os.environ["TAI_HOME"] = tmp
|
||
try:
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
config = tai.Config(FakeArgs())
|
||
store = tai.Store(config, tai.Seal(config.home, ""))
|
||
agent = tai.Agent(config, store)
|
||
reply = {"role": "assistant", "content": "Hello there", "reasoning": "", "tool_calls": [], "backend": "x"}
|
||
with mock.patch.object(agent.chat, "complete", return_value=reply):
|
||
with mock.patch.dict(os.environ):
|
||
os.environ.pop("NO_COLOR", None)
|
||
with mock.patch("sys.stdout") as fake_out:
|
||
fake_out.isatty.return_value = True
|
||
with mock.patch("builtins.print") as printer:
|
||
with mock.patch("tai.Spinner"):
|
||
agent.run_turn("hi")
|
||
lines = []
|
||
for call in printer.call_args_list:
|
||
if call.args:
|
||
lines.extend(str(call.args[0]).splitlines())
|
||
banners = [line for line in lines if tai.strip_ansi(line).strip() == "tai"]
|
||
self.assertEqual(banners, [])
|
||
store.close()
|
||
finally:
|
||
if old_home is None:
|
||
os.environ.pop("TAI_HOME", None)
|
||
else:
|
||
os.environ["TAI_HOME"] = old_home
|
||
|
||
|
||
class ToolArgumentsTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
os.environ["TAI_HOME"] = self.tmp.name
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
self.config = tai.Config(FakeArgs())
|
||
self.store = tai.Store(self.config, tai.Seal(self.config.home, "args-test-1"))
|
||
|
||
def tearDown(self):
|
||
self.store.close()
|
||
self.tmp.cleanup()
|
||
os.environ.pop("TAI_HOME", None)
|
||
|
||
def test_tool_name_clean(self):
|
||
self.assertEqual(tai.tool_name_clean("shell<|channel|>commentary"), "shell")
|
||
self.assertEqual(tai.tool_name_clean(" read_file "), "read_file")
|
||
self.assertEqual(tai.tool_name_clean(None), "")
|
||
|
||
def test_tool_args_text(self):
|
||
self.assertEqual(tai.tool_args_text(None), "{}")
|
||
self.assertEqual(tai.tool_args_text(""), "{}")
|
||
self.assertEqual(tai.tool_args_text("{}"), "{}")
|
||
self.assertEqual(tai.tool_args_text({"limit": 50}), '{"limit": 50}')
|
||
self.assertEqual(tai.tool_args_text([]), "{}")
|
||
self.assertEqual(json.loads(tai.tool_args_text({"checks": ["os"]})), {"checks": ["os"]})
|
||
|
||
def test_muse_reply_coerces_dict_arguments(self):
|
||
agent = tai.Agent(self.config, self.store, persist=False, quiet=True)
|
||
reply = agent.chat.muse_reply({"content": "", "tool_calls": [{"name": "get_current_terminal_content", "arguments": {"limit": 50}}]}, "opencode")
|
||
self.assertEqual(reply["tool_calls"][0]["arguments"], '{"limit": 50}')
|
||
|
||
def test_normalize_coerces_dict_arguments(self):
|
||
agent = tai.Agent(self.config, self.store, persist=False, quiet=True)
|
||
message = {"content": "", "tool_calls": [{"id": "c1", "function": {"name": "sysinfo", "arguments": {"checks": ["os"]}}}]}
|
||
reply = agent.chat.normalize(message, "x")
|
||
self.assertEqual(reply["tool_calls"][0]["arguments"], '{"checks": ["os"]}')
|
||
|
||
def test_show_call_dict_arguments(self):
|
||
agent = tai.Agent(self.config, self.store, persist=False, quiet=False)
|
||
with mock.patch("builtins.print") as printer:
|
||
agent.show_call({"id": "c1", "name": "get_current_terminal_content", "arguments": {"limit": 50}})
|
||
printed = "\n".join(str(call.args[0]) for call in printer.call_args_list if call.args)
|
||
self.assertIn("get_current_terminal_content", printed)
|
||
self.assertIn("limit", printed)
|
||
|
||
def test_dispatch_dict_arguments(self):
|
||
tools = tai.Tools(FakeApp())
|
||
subset = tools.dispatch("sysinfo", {"checks": ["python"]})
|
||
self.assertEqual(len(subset.splitlines()), 1)
|
||
self.assertIn("unknown checks: nope", tools.dispatch("sysinfo", {"checks": ["nope"]}))
|
||
self.assertIn("not valid JSON", tools.dispatch("sysinfo", ["nope"]))
|
||
|
||
def test_action_signature_dict_arguments(self):
|
||
self.assertEqual(tai.action_signature("sysinfo", {"checks": ["os"]}), 'sysinfo {"checks": ["os"]}')
|
||
self.assertEqual(tai.action_signature("sysinfo", '{"checks": ["os"]}'), 'sysinfo {"checks": ["os"]}')
|
||
|
||
def test_redact_non_string_passthrough(self):
|
||
marker = {"limit": 50}
|
||
self.assertIs(self.store.redact(marker), marker)
|
||
|
||
def test_run_turn_dict_arguments(self):
|
||
agent = tai.Agent(self.config, self.store, persist=False, quiet=True)
|
||
call_reply = {"role": "assistant", "content": "", "reasoning": "", "tool_calls": [{"id": "c1", "name": "sysinfo", "arguments": {"checks": ["os"]}}], "backend": "x"}
|
||
final_reply = {"role": "assistant", "content": "done", "reasoning": "", "tool_calls": [], "backend": "x"}
|
||
with mock.patch.object(agent.chat, "complete", side_effect=[call_reply, final_reply]):
|
||
self.assertEqual(agent.run_turn("check os", capture=True), "done")
|
||
|
||
|
||
class OpencodeAgentTests(unittest.TestCase):
|
||
def test_agent_is_not_plan(self):
|
||
# "plan" makes opencode inject "Plan mode ACTIVE - READ-ONLY" into the
|
||
# model, which then refuses to emit tai's write tool calls.
|
||
self.assertNotEqual(tai.OPENCODE_AGENT, "plan")
|
||
captured = {}
|
||
|
||
def fake_run(argv, **kwargs):
|
||
captured["argv"] = argv
|
||
return mock.Mock(stdout='{"type":"text","part":{"text":"ok"}}\n', stderr="", returncode=0)
|
||
|
||
with mock.patch.object(tai, "opencode_binary", return_value="opencode"), mock.patch.object(tai.subprocess, "run", side_effect=fake_run):
|
||
self.assertEqual(tai.Opencode.opencode_run("m", "hi", 5), "ok")
|
||
argv = captured["argv"]
|
||
self.assertEqual(argv[argv.index("--agent") + 1], tai.OPENCODE_AGENT)
|
||
|
||
|
||
class WorkerIdleTests(unittest.TestCase):
|
||
def test_idle_worker_is_flagged(self):
|
||
record = {"id": 1, "task": "t", "status": "done", "elapsed": 3, "tools": 0}
|
||
self.assertIn("0 tools", tai.strip_ansi(tai.Glow.agent_row(record)))
|
||
record["tools"] = 2
|
||
self.assertNotIn("0 tools", tai.strip_ansi(tai.Glow.agent_row(record)))
|
||
record["tools"] = None
|
||
self.assertNotIn("0 tools", tai.strip_ansi(tai.Glow.agent_row(record)))
|
||
running = {"id": 1, "task": "t", "status": "running", "elapsed": 3, "tools": 0}
|
||
self.assertNotIn("0 tools", tai.strip_ansi(tai.Glow.agent_row(running)))
|
||
|
||
def test_poll_notes_zero_tool_calls(self):
|
||
agent_id = tai.spawn_agent("idle", "t", 60, None, None, 0, lambda task, profile, timeout: "I am in plan mode")
|
||
tai.poll_agent(agent_id, wait=5)
|
||
with tai.AGENTS_LOCK:
|
||
tai.AGENTS[agent_id]["tools"] = 0
|
||
status, text = tai.poll_agent(agent_id)
|
||
self.assertEqual(status, "done")
|
||
self.assertIn("ran 0 tool calls", text)
|
||
with tai.AGENTS_LOCK:
|
||
tai.AGENTS[agent_id]["tools"] = 3
|
||
self.assertNotIn("0 tool calls", tai.poll_agent(agent_id)[1])
|
||
|
||
def test_tools_counts_dispatches(self):
|
||
tools = tai.Tools.__new__(tai.Tools)
|
||
tools.calls = 0
|
||
tools.handlers = {"echo": lambda args: "x"}
|
||
tools.app = mock.Mock(store=None)
|
||
tools.dispatch("echo", {})
|
||
tools.dispatch("nope", {})
|
||
self.assertEqual(tools.calls, 2)
|
||
|
||
|
||
class OpencodeParseTests(unittest.TestCase):
|
||
def test_single_answer(self):
|
||
answer = {"content": "hi", "tool_calls": []}
|
||
self.assertEqual(tai.opencode_parse(json.dumps(answer)), answer)
|
||
|
||
def test_schema_echo_then_answer(self):
|
||
answer = {"content": "Checking which text files were touched today.", "tool_calls": [{"name": "shell", "arguments": '{"command": "ls -lt"}'}, {"name": "audit", "arguments": '{"limit": 30}'}]}
|
||
text = json.dumps(tai.muse_output_schema()) + " " + json.dumps(answer)
|
||
parsed = tai.opencode_parse(text)
|
||
self.assertEqual(parsed["content"], answer["content"])
|
||
self.assertEqual(len(parsed["tool_calls"]), 2)
|
||
self.assertEqual(parsed["tool_calls"][0]["name"], "shell")
|
||
|
||
def test_history_json_echo_then_answer(self):
|
||
stale = {"content": "stale tool output", "tool_calls": []}
|
||
answer = {"content": "fresh answer", "tool_calls": []}
|
||
text = "TOOL c1:\n" + json.dumps(stale) + "\n\n" + json.dumps(answer)
|
||
self.assertEqual(tai.opencode_parse(text)["content"], "fresh answer")
|
||
|
||
def test_tools_blob_echo_then_answer(self):
|
||
tools = [{"name": "shell", "parameters": {"command": {"type": "string"}}}, {"name": "audit", "parameters": {"limit": {"type": "integer"}}}]
|
||
answer = {"content": "done", "tool_calls": []}
|
||
text = "TOOLS:\n" + json.dumps(tools) + "\n\nAnswer with content text plus tool_calls.\n" + json.dumps(answer)
|
||
parsed = tai.opencode_parse(text)
|
||
self.assertEqual(parsed["content"], "done")
|
||
self.assertEqual(parsed["tool_calls"], [])
|
||
|
||
def test_answer_then_trailing_prose(self):
|
||
answer = {"content": "done", "tool_calls": []}
|
||
parsed = tai.opencode_parse(json.dumps(answer) + "\nHope that helps! } junk {")
|
||
self.assertEqual(parsed["content"], "done")
|
||
self.assertEqual(parsed["tool_calls"], [])
|
||
|
||
def test_fenced_answer(self):
|
||
answer = {"content": "hi", "tool_calls": []}
|
||
self.assertEqual(tai.opencode_parse("```json\n" + json.dumps(answer) + "\n```"), answer)
|
||
|
||
def test_fenced_answer_among_prose(self):
|
||
answer = {"content": "hi", "tool_calls": []}
|
||
text = "Sure, here it is:\n```json\n" + json.dumps(answer) + "\n```\nLet me know."
|
||
self.assertEqual(tai.opencode_parse(text), answer)
|
||
|
||
def test_plain_text_fallback(self):
|
||
self.assertEqual(tai.opencode_parse("just some words"), {"content": "just some words", "tool_calls": []})
|
||
self.assertEqual(tai.opencode_parse(""), {"content": "", "tool_calls": []})
|
||
self.assertEqual(tai.opencode_parse(None), {"content": "", "tool_calls": []})
|
||
|
||
def test_braces_inside_content_string(self):
|
||
answer = {"content": 'use {x} and {"d": 1} wisely', "tool_calls": []}
|
||
parsed = tai.opencode_parse(json.dumps(answer))
|
||
self.assertEqual(parsed["content"], 'use {x} and {"d": 1} wisely')
|
||
|
||
def test_pretty_printed_answer(self):
|
||
answer = {"content": "hi", "tool_calls": [{"name": "shell", "arguments": '{"command": "ls"}'}]}
|
||
parsed = tai.opencode_parse(json.dumps(answer, indent=2))
|
||
self.assertEqual(parsed["content"], "hi")
|
||
self.assertEqual(len(parsed["tool_calls"]), 1)
|
||
|
||
def test_arguments_as_object_preserved(self):
|
||
answer = {"content": "checking", "tool_calls": [{"name": "shell", "arguments": {"command": "ls -lt"}}]}
|
||
parsed = tai.opencode_parse(json.dumps(tai.muse_output_schema()) + " " + json.dumps(answer))
|
||
self.assertEqual(parsed["tool_calls"][0]["arguments"], {"command": "ls -lt"})
|
||
|
||
def test_many_leading_braces_still_finds_answer(self):
|
||
answer = {"content": "found", "tool_calls": []}
|
||
parsed = tai.opencode_parse("{{{not json, {{{also not, " * 40 + json.dumps(answer))
|
||
self.assertEqual(parsed["content"], "found")
|
||
|
||
def test_opencode_turn_runs_tools_after_schema_echo(self):
|
||
answer = {"content": "checking", "tool_calls": [{"name": "shell", "arguments": {"command": "ls -lt"}}]}
|
||
text = json.dumps(tai.muse_output_schema()) + " " + json.dumps(answer)
|
||
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
os.environ["TAI_HOME"] = tmp
|
||
try:
|
||
chat = tai.ChatClient(tai.Config(FakeArgs()))
|
||
with mock.patch.object(tai.Opencode, "opencode_run", return_value=text):
|
||
reply = chat.opencode_turn([{"role": "user", "content": "hi"}], None, None, "opencode", "m")
|
||
finally:
|
||
os.environ.pop("TAI_HOME", None)
|
||
self.assertEqual(len(reply["tool_calls"]), 1)
|
||
self.assertEqual(reply["tool_calls"][0]["name"], "shell")
|
||
self.assertEqual(json.loads(reply["tool_calls"][0]["arguments"]), {"command": "ls -lt"})
|
||
|
||
|
||
class CliJsonRepairTests(unittest.TestCase):
|
||
def make_chat(self, tmp):
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
os.environ["TAI_HOME"] = tmp
|
||
return tai.ChatClient(tai.Config(FakeArgs()))
|
||
|
||
def test_find_answer_none_on_prose(self):
|
||
self.assertIsNone(tai.opencode_find_answer("just some words"))
|
||
self.assertIsNone(tai.opencode_find_answer(""))
|
||
answer = {"content": "hi", "tool_calls": []}
|
||
self.assertEqual(tai.opencode_find_answer(json.dumps(answer)), answer)
|
||
|
||
def test_prose_then_json_repairs_once(self):
|
||
tools = [{"type": "function", "function": {"name": "shell"}}]
|
||
fixed = {"content": "checking", "tool_calls": [{"name": "shell", "arguments": '{"command": "ls"}'}]}
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
try:
|
||
chat = self.make_chat(tmp)
|
||
with mock.patch.object(tai.Opencode, "opencode_run", side_effect=["Ik kijk even mee op je scherm.", json.dumps(fixed)]) as runner:
|
||
reply = chat.opencode_turn([{"role": "user", "content": "hi"}], tools, None, "opencode", "m")
|
||
self.assertEqual(runner.call_count, 2)
|
||
self.assertIn("no JSON object", runner.call_args_list[1][0][1])
|
||
finally:
|
||
os.environ.pop("TAI_HOME", None)
|
||
self.assertEqual(len(reply["tool_calls"]), 1)
|
||
self.assertEqual(reply["tool_calls"][0]["name"], "shell")
|
||
|
||
def test_prose_twice_accepts_second_as_final(self):
|
||
tools = [{"type": "function", "function": {"name": "shell"}}]
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
try:
|
||
chat = self.make_chat(tmp)
|
||
with mock.patch.object(tai.Opencode, "opencode_run", side_effect=["first prose", "second prose"]) as runner:
|
||
reply = chat.opencode_turn([{"role": "user", "content": "hi"}], tools, None, "opencode", "m")
|
||
self.assertEqual(runner.call_count, 2)
|
||
finally:
|
||
os.environ.pop("TAI_HOME", None)
|
||
self.assertEqual(reply["content"], "second prose")
|
||
self.assertEqual(reply["tool_calls"], [])
|
||
|
||
def test_json_first_runs_no_repair(self):
|
||
tools = [{"type": "function", "function": {"name": "shell"}}]
|
||
answer = {"content": "checking", "tool_calls": [{"name": "shell", "arguments": '{"command": "ls"}'}]}
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
try:
|
||
chat = self.make_chat(tmp)
|
||
with mock.patch.object(tai.Opencode, "opencode_run", return_value=json.dumps(answer)) as runner:
|
||
reply = chat.opencode_turn([{"role": "user", "content": "hi"}], tools, None, "opencode", "m")
|
||
self.assertEqual(runner.call_count, 1)
|
||
finally:
|
||
os.environ.pop("TAI_HOME", None)
|
||
self.assertEqual(len(reply["tool_calls"]), 1)
|
||
|
||
def test_no_repair_without_tools(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
try:
|
||
chat = self.make_chat(tmp)
|
||
with mock.patch.object(tai.Opencode, "opencode_run", return_value="just chatting") as runner:
|
||
reply = chat.opencode_turn([{"role": "user", "content": "hi"}], None, None, "opencode", "m")
|
||
self.assertEqual(runner.call_count, 1)
|
||
finally:
|
||
os.environ.pop("TAI_HOME", None)
|
||
self.assertEqual(reply["content"], "just chatting")
|
||
|
||
def test_failed_repair_keeps_first_answer(self):
|
||
tools = [{"type": "function", "function": {"name": "shell"}}]
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
try:
|
||
chat = self.make_chat(tmp)
|
||
with mock.patch.object(tai.Opencode, "opencode_run", side_effect=["usable prose", tai.BackendError("boom")]) as runner:
|
||
reply = chat.opencode_turn([{"role": "user", "content": "hi"}], tools, None, "opencode", "m")
|
||
self.assertEqual(runner.call_count, 2)
|
||
finally:
|
||
os.environ.pop("TAI_HOME", None)
|
||
self.assertEqual(reply["content"], "usable prose")
|
||
|
||
def test_prompt_contract_states_turn_ending(self):
|
||
self.assertIn("text without tool calls ends your turn", tai.DEFAULT_SYSTEM)
|
||
rules = tai.opencode_prompt([{"role": "user", "content": "hi"}], [])
|
||
self.assertIn("ends the turn at once", rules)
|
||
self.assertIn("Never narrate a planned action", rules)
|
||
|
||
|
||
class MuseArgvTests(unittest.TestCase):
|
||
def argv(self, session_id=None):
|
||
class FakeConfig:
|
||
explicit_model = ""
|
||
|
||
return tai.muse_argv(FakeConfig(), "s.json", "p.txt", session_id)
|
||
|
||
def test_ephemeral_omits_workspace_switch(self):
|
||
argv = self.argv(None)
|
||
self.assertIn("--no-session-log", argv)
|
||
self.assertNotIn("--allow-workspace-switch", argv)
|
||
self.assertNotIn("--session-id", argv)
|
||
self.assertIn("--disable-shell", argv)
|
||
self.assertIn("--disable-write", argv)
|
||
|
||
def test_session_keeps_workspace_switch(self):
|
||
argv = self.argv("sess-1")
|
||
self.assertIn("--allow-workspace-switch", argv)
|
||
self.assertIn("--session-id", argv)
|
||
self.assertIn("sess-1", argv)
|
||
self.assertNotIn("--no-session-log", argv)
|
||
|
||
|
||
class RetrySummaryTests(unittest.TestCase):
|
||
def test_empty_errors(self):
|
||
self.assertEqual(tai.ChatClient.summarize_errors([]), "no details")
|
||
|
||
def test_actionable_error_leads_over_missing_cli_tail(self):
|
||
errors = [
|
||
"openrouter/m1: HTTP 401: No auth credentials found",
|
||
"pollinations/openai: HTTP 429: rate limited",
|
||
"muse/muse: muse CLI not found (MUSE_BIN or PATH)",
|
||
]
|
||
summary = tai.ChatClient.summarize_errors(errors)
|
||
self.assertTrue(summary.startswith("openrouter/m1: HTTP 401"), summary)
|
||
self.assertIn("1 missing CLI skipped", summary)
|
||
self.assertNotIn("muse CLI not found", summary.split(";")[0])
|
||
|
||
def test_only_missing_cli_still_reported(self):
|
||
summary = tai.ChatClient.summarize_errors(["muse/muse: muse CLI not found (MUSE_BIN or PATH)"])
|
||
self.assertIn("muse CLI not found", summary)
|
||
self.assertNotIn("skipped", summary)
|
||
|
||
def test_duplicates_collapsed_with_more_count(self):
|
||
errors = ["a/m: boom"] * 5 + ["b/m: bust", "c/m: kaput"]
|
||
summary = tai.ChatClient.summarize_errors(errors)
|
||
self.assertIn("a/m: boom; b/m: bust", summary)
|
||
self.assertIn("+1 more", summary)
|
||
|
||
def test_width_cap(self):
|
||
summary = tai.ChatClient.summarize_errors(["a/m: " + "x" * 500])
|
||
self.assertLessEqual(len(summary), 320)
|
||
|
||
def test_complete_retry_line_leads_with_actionable_error(self):
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
os.environ["TAI_HOME"] = tmp
|
||
try:
|
||
chat = tai.ChatClient(tai.Config(FakeArgs()))
|
||
cands = [
|
||
{"provider": "openrouter", "chat_url": "https://x/chat", "models_url": "https://x/models", "key": "", "model": "m1", "profile": ""},
|
||
{"provider": "muse", "chat_url": "muse://exec", "models_url": "muse://models", "key": "", "model": "muse", "profile": ""},
|
||
]
|
||
calls = []
|
||
|
||
def fake_try(cand, messages, tools, payload, stream_sink, ephemeral, errors):
|
||
calls.append(cand["provider"])
|
||
if len(calls) > len(cands):
|
||
return {"role": "assistant", "content": "ok", "reasoning": "", "tool_calls": [], "backend": cand["provider"]}
|
||
if cand["provider"] == "muse":
|
||
errors.append("muse/muse: muse CLI not found (MUSE_BIN or PATH)")
|
||
else:
|
||
errors.append("openrouter/m1: HTTP 401: No auth credentials found")
|
||
return None
|
||
|
||
with mock.patch.object(tai.ChatClient, "ordered_candidates", return_value=cands), \
|
||
mock.patch.object(tai.ChatClient, "all_candidates", return_value=cands), \
|
||
mock.patch.object(chat, "try_candidate", side_effect=fake_try), \
|
||
mock.patch("time.sleep") as sleeper:
|
||
err = io.StringIO()
|
||
with mock.patch("sys.stderr", err):
|
||
reply = chat.complete([{"role": "user", "content": "hi"}])
|
||
self.assertEqual(reply["content"], "ok")
|
||
self.assertEqual(len(sleeper.call_args_list), 1)
|
||
finally:
|
||
os.environ.pop("TAI_HOME", None)
|
||
out = err.getvalue()
|
||
self.assertIn("all 2 candidates across 2 backends failed (round 1)", out)
|
||
self.assertIn("openrouter/m1: HTTP 401", out)
|
||
self.assertNotIn("retrying in 2s: muse/muse", out)
|
||
|
||
def test_complete_give_up_error_summarizes(self):
|
||
class FakeArgs:
|
||
profile = "t"
|
||
yes = True
|
||
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
os.environ["TAI_HOME"] = tmp
|
||
try:
|
||
config = tai.Config(FakeArgs())
|
||
config.retry_seconds = 0.05
|
||
chat = tai.ChatClient(config)
|
||
cands = [
|
||
{"provider": "openrouter", "chat_url": "https://x/chat", "models_url": "https://x/models", "key": "", "model": "m1", "profile": ""},
|
||
{"provider": "muse", "chat_url": "muse://exec", "models_url": "muse://models", "key": "", "model": "muse", "profile": ""},
|
||
]
|
||
|
||
def fake_try(cand, messages, tools, payload, stream_sink, ephemeral, errors):
|
||
errors.append("%s/%s: %s" % (cand["provider"], cand["model"], "muse CLI not found (MUSE_BIN or PATH)" if cand["provider"] == "muse" else "HTTP 401: No auth"))
|
||
return None
|
||
|
||
with mock.patch.object(tai.ChatClient, "ordered_candidates", return_value=cands), \
|
||
mock.patch.object(chat, "try_candidate", side_effect=fake_try):
|
||
with self.assertRaises(tai.BackendError) as raised:
|
||
chat.complete([{"role": "user", "content": "hi"}])
|
||
finally:
|
||
os.environ.pop("TAI_HOME", None)
|
||
self.assertIn("after 1 rounds", str(raised.exception))
|
||
self.assertTrue(str(raised.exception).index("openrouter/m1") < str(raised.exception).index("missing CLI skipped"), str(raised.exception))
|
||
|
||
|
||
class TuiLayoutTests(unittest.TestCase):
|
||
def make_app(self, rows=10, cols=20):
|
||
app = tai.TuiApp()
|
||
app.rows = rows
|
||
app.cols = cols
|
||
return app
|
||
|
||
def test_rows_fr_split_exact(self):
|
||
app = self.make_app(rows=9, cols=30)
|
||
rows = tai.TuiRows(tai.TuiPanel(h="1fr"), tai.TuiPanel(h="1fr"), tai.TuiPanel(h="1fr"), x=0, y=0, w=30, h=9)
|
||
root = tai.TuiPanel(flow=None)
|
||
root.add(rows)
|
||
app.mount(root)
|
||
self.assertEqual([(child.ay, child.ah) for child in rows.children], [(0, 3), (3, 3), (6, 3)])
|
||
|
||
def test_rows_fr_largest_remainder(self):
|
||
app = self.make_app(rows=10, cols=30)
|
||
rows = tai.TuiRows(tai.TuiPanel(h="1fr"), tai.TuiPanel(h="1fr"), tai.TuiPanel(h="1fr"), x=0, y=0, w=30, h=10)
|
||
root = tai.TuiPanel(flow=None)
|
||
root.add(rows)
|
||
app.mount(root)
|
||
self.assertEqual([child.ah for child in rows.children], [4, 3, 3])
|
||
self.assertEqual(sum(child.ah for child in rows.children), 10)
|
||
|
||
def test_dock_carves_then_fill(self):
|
||
app = self.make_app(rows=10, cols=20)
|
||
root = tai.TuiPanel(flow=None)
|
||
top = tai.TuiPanel(dock="top", h=2)
|
||
left = tai.TuiPanel(dock="left", w=5)
|
||
fill = tai.TuiPanel(x=0, y=0, w="100%", h="100%")
|
||
root.add(top, left, fill)
|
||
app.mount(root)
|
||
self.assertEqual((fill.ax, fill.ay, fill.aw, fill.ah), (5, 2, 15, 8))
|
||
|
||
def test_float_percent_center_negative(self):
|
||
app = self.make_app(rows=10, cols=20)
|
||
root = tai.TuiPanel(flow=None)
|
||
centered = tai.TuiPanel(x="center", y=1, w="50%", h=4)
|
||
neg = tai.TuiPanel(x=-5, y=0, w=4, h=1)
|
||
root.add(centered, neg)
|
||
app.mount(root)
|
||
self.assertEqual((centered.ax, centered.ay, centered.aw, centered.ah), (5, 1, 10, 4))
|
||
self.assertEqual(neg.ax, 11)
|
||
|
||
def test_rows_cross_auto_fills_cols_hugs(self):
|
||
app = self.make_app(rows=10, cols=20)
|
||
root = tai.TuiPanel(flow=None)
|
||
rows = tai.TuiRows(tai.TuiText("hi"), x=0, y=0, w=20, h=5)
|
||
cols = tai.TuiCols(tai.TuiText("hi"), x=0, y=5, w=20, h=5)
|
||
root.add(rows, cols)
|
||
app.mount(root)
|
||
self.assertEqual(rows.children[0].aw, 20)
|
||
self.assertEqual(cols.children[0].ah, 1)
|
||
|
||
def test_border_chrome_shrinks_content(self):
|
||
panel = tai.TuiPanel(border=True, padding=1, w=10, h=6)
|
||
panel.arrange(0, 0, 10, 6, (0, 0, 10, 6))
|
||
self.assertEqual(panel._content_w(), 6)
|
||
self.assertEqual(panel._content_h(), 2)
|
||
painted = panel.paint()
|
||
self.assertEqual(len(painted), 6)
|
||
top = "".join(char for char, _style in tai.tui_flatten(painted[0]))
|
||
self.assertEqual((top[0], top[-1], len(top)), ("┌", "┐", 10))
|
||
|
||
|
||
class TuiCompositorTests(unittest.TestCase):
|
||
def test_z_occlusion_top_wins(self):
|
||
app = tai.TuiApp()
|
||
app.rows = 6
|
||
app.cols = 20
|
||
root = tai.TuiPanel(flow=None)
|
||
root.add(tai.TuiText("base-base-base-base", x=0, y=0, w=20, h=6))
|
||
root.add(tai.TuiText("TOP", x=5, y=2, w=5, h=1, z=5))
|
||
app.mount(root)
|
||
text = app.render_text()
|
||
self.assertEqual(text.split("\n")[2][5:8], "TOP")
|
||
self.assertEqual(text.split("\n")[0][:4], "base")
|
||
|
||
def test_layer_band_beats_z(self):
|
||
low = tai.TuiPanel(z=5)
|
||
high = tai.TuiPanel(layer="dialog", z=5)
|
||
self.assertGreater(high.effective_z(), low.effective_z())
|
||
|
||
def test_child_paints_above_high_z_parent(self):
|
||
app = tai.TuiApp()
|
||
app.rows = 4
|
||
app.cols = 12
|
||
root = tai.TuiPanel(flow=None)
|
||
dialog = tai.TuiPanel(layer="dialog", x=0, y=0, w=12, h=4)
|
||
dialog.add(tai.TuiText("inside", x=0, y=0, w=12, h=4))
|
||
root.add(dialog)
|
||
app.mount(root)
|
||
order = [node for node in app.compositor.flatten(root)]
|
||
self.assertLess(order.index(dialog), order.index(dialog.children[0]))
|
||
self.assertIn("inside", app.render_text())
|
||
|
||
def test_diff_second_frame_short(self):
|
||
app = tai.TuiApp()
|
||
app.rows = 6
|
||
app.cols = 20
|
||
root = tai.TuiPanel(flow=None)
|
||
root.add(tai.TuiText("hello", x=0, y=0, w=20, h=6))
|
||
app.mount(root)
|
||
frame = app.compositor.compose(root, 6, 20)
|
||
first = app.compositor.render(frame)
|
||
frame2 = app.compositor.compose(root, 6, 20)
|
||
second = app.compositor.render(frame2)
|
||
self.assertIn("hello", first)
|
||
self.assertIn(tai.TUI_SYNC_START, second)
|
||
self.assertLess(len(second), len(first))
|
||
|
||
def test_hit_topmost(self):
|
||
app = tai.TuiApp()
|
||
app.rows = 6
|
||
app.cols = 20
|
||
root = tai.TuiPanel(flow=None)
|
||
bottom = tai.TuiPanel(x=0, y=0, w=20, h=6, id="bottom")
|
||
top = tai.TuiPanel(x=5, y=2, w=5, h=2, z=5, id="top")
|
||
root.add(bottom, top)
|
||
app.mount(root)
|
||
self.assertIs(app.hitmap.at(6, 3), top)
|
||
self.assertIs(app.hitmap.at(1, 1), bottom)
|
||
self.assertIs(app.hitmap.at(19, 5), bottom)
|
||
|
||
|
||
class TuiEventTests(unittest.TestCase):
|
||
def test_bubble_order_and_stop(self):
|
||
order = []
|
||
root = tai.TuiPanel()
|
||
mid = tai.TuiPanel()
|
||
leaf = tai.TuiPanel()
|
||
root.add(mid)
|
||
mid.add(leaf)
|
||
root.on_key = lambda event: order.append("root")
|
||
mid.on_key = lambda event: order.append("mid")
|
||
leaf.on_key = lambda event: order.append("leaf")
|
||
leaf.emit(tai.TuiKey(key="x", char="x"))
|
||
self.assertEqual(order, ["leaf", "mid", "root"])
|
||
order.clear()
|
||
leaf.on_key = lambda event: (order.append("leaf"), event.stop())
|
||
leaf.emit(tai.TuiKey(key="x", char="x"))
|
||
self.assertEqual(order, ["leaf"])
|
||
|
||
def test_focus_blur_do_not_bubble(self):
|
||
seen = []
|
||
root = tai.TuiPanel()
|
||
leaf = tai.TuiPanel()
|
||
root.add(leaf)
|
||
root.on_focus = lambda event: seen.append("root")
|
||
leaf.emit(tai.TuiFocus(panel=leaf))
|
||
self.assertEqual(seen, [])
|
||
|
||
def test_handler_name(self):
|
||
self.assertEqual(tai.TuiClick(x=0, y=0, button=0, modifiers=[], pressed=True).handler_name, "on_click")
|
||
|
||
|
||
class TuiFocusTests(unittest.TestCase):
|
||
def make_app(self):
|
||
app = tai.TuiApp()
|
||
app.rows = 10
|
||
app.cols = 30
|
||
return app
|
||
|
||
def test_chain_sorted_by_position(self):
|
||
app = self.make_app()
|
||
root = tai.TuiPanel(flow=None)
|
||
late = tai.TuiButton("late", x=0, y=8, w=10)
|
||
early = tai.TuiButton("early", x=0, y=1, w=10)
|
||
root.add(late, early)
|
||
app.mount(root)
|
||
self.assertEqual(app.focusables(), [early, late])
|
||
|
||
def test_tab_wraps_both_directions(self):
|
||
app = self.make_app()
|
||
root = tai.TuiPanel(flow=None)
|
||
first = tai.TuiButton("a", x=0, y=1, w=8)
|
||
second = tai.TuiButton("b", x=0, y=5, w=8)
|
||
root.add(first, second)
|
||
app.mount(root)
|
||
self.assertIs(app.focused, first)
|
||
app.dispatch(tai.TuiKey(key="tab", char="\t"))
|
||
self.assertIs(app.focused, second)
|
||
app.dispatch(tai.TuiKey(key="tab", char="\t"))
|
||
self.assertIs(app.focused, first)
|
||
app.dispatch(tai.TuiKey(key="backtab", char=""))
|
||
self.assertIs(app.focused, second)
|
||
|
||
def test_tab_skips_disabled_and_hidden(self):
|
||
app = self.make_app()
|
||
root = tai.TuiPanel(flow=None)
|
||
ok = tai.TuiButton("ok", x=0, y=1, w=8)
|
||
off = tai.TuiButton("off", x=0, y=5, w=8, enabled=False)
|
||
hid = tai.TuiButton("hid", x=0, y=7, w=8, visible=False)
|
||
root.add(ok, off, hid)
|
||
app.mount(root)
|
||
self.assertEqual(app.focusables(), [ok])
|
||
|
||
def test_modal_trap_and_outside_click_swallowed(self):
|
||
app = self.make_app()
|
||
root = tai.TuiPanel(flow=None)
|
||
behind = tai.TuiButton("behind", x=0, y=1, w=10)
|
||
dialog = tai.TuiDialog(TuiFocusTests._ok_button(), x=5, y=2, w=20, h=6)
|
||
root.add(behind, dialog)
|
||
app.mount(root)
|
||
self.assertEqual(app.focusables(), [dialog.children[0]])
|
||
app.set_focus(behind)
|
||
self.assertIsNot(app.focused, behind)
|
||
pressed = []
|
||
behind.on_press = pressed.append
|
||
app._dispatch_click(tai.TuiClick(x=1, y=1, button=0, modifiers=[], pressed=True))
|
||
self.assertEqual(pressed, [])
|
||
|
||
@staticmethod
|
||
def _ok_button():
|
||
return tai.TuiButton("ok", w=10)
|
||
|
||
def test_click_focuses_and_presses(self):
|
||
app = self.make_app()
|
||
root = tai.TuiPanel(flow=None)
|
||
pressed = []
|
||
button = tai.TuiButton("go", x=2, y=2, w=12)
|
||
button.on_press = pressed.append
|
||
root.add(button)
|
||
app.mount(root)
|
||
app._dispatch_click(tai.TuiClick(x=4, y=3, button=0, modifiers=[], pressed=True))
|
||
self.assertIs(app.focused, button)
|
||
self.assertEqual(len(pressed), 1)
|
||
self.assertIs(pressed[0].panel, button)
|
||
|
||
def test_tab_escapes_terminal(self):
|
||
app = self.make_app()
|
||
root = tai.TuiPanel(flow=None)
|
||
term = tai.TuiTerminal(x=0, y=0, w=20, h=8)
|
||
other = tai.TuiButton("b", x=0, y=9, w=8, h=1)
|
||
root.add(term, other)
|
||
app.mount(root)
|
||
app.set_focus(term)
|
||
app.dispatch(tai.TuiKey(key="tab", char="\t"))
|
||
self.assertIs(app.focused, other)
|
||
app.dispatch(tai.TuiKey(key="backtab", char=""))
|
||
self.assertIs(app.focused, term)
|
||
|
||
def test_mount_prefers_non_terminal(self):
|
||
app = self.make_app()
|
||
root = tai.TuiPanel(flow=None)
|
||
term = tai.TuiTerminal(x=0, y=0, w=20, h=8)
|
||
button = tai.TuiButton("b", x=0, y=9, w=8, h=1)
|
||
root.add(term, button)
|
||
app.mount(root)
|
||
self.assertIs(app.focused, button)
|
||
|
||
def test_modal_close_restores_previous_focus(self):
|
||
app = self.make_app()
|
||
root = tai.TuiPanel(flow=None)
|
||
behind = tai.TuiButton("behind", x=0, y=1, w=10)
|
||
root.add(behind)
|
||
app.mount(root)
|
||
self.assertIs(app.focused, behind)
|
||
dialog = tai.TuiDialog(TuiFocusTests._ok_button(), x=5, y=2, w=20, h=6)
|
||
root.add(dialog)
|
||
self.assertIs(app.focused, dialog.children[0])
|
||
app.close_panel(dialog)
|
||
self.assertIs(app.focused, behind)
|
||
|
||
def test_demo_answer_then_q_quits(self):
|
||
app = tai.TuiApp()
|
||
app.rows = 30
|
||
app.cols = 100
|
||
root, _refs = tai.tui_build_demo(app, live=False)
|
||
app.mount(root)
|
||
welcome = root.find("welcome")
|
||
self.assertIsNotNone(welcome)
|
||
self.assertIs(app.focused, welcome.picker)
|
||
app.dispatch(tai.TuiKey(key="enter", char="\r"))
|
||
self.assertIsNone(root.find("welcome"))
|
||
actions = root.find("actions")
|
||
self.assertIsNotNone(actions)
|
||
self.assertIs(app.focused, actions)
|
||
app._running = True
|
||
app.dispatch(tai.TuiKey(key="q", char="q"))
|
||
self.assertFalse(app._running)
|
||
|
||
|
||
class TuiInputTests(unittest.TestCase):
|
||
def feed(self, *chunks):
|
||
parser = tai.TuiInput()
|
||
events = []
|
||
for chunk in chunks:
|
||
events.extend(parser.feed(chunk))
|
||
return parser, events
|
||
|
||
def test_sgr_mouse_press_release_wheel(self):
|
||
_parser, events = self.feed(b"\x1b[<0;10;5M\x1b[<0;10;5m\x1b[<64;3;3M")
|
||
self.assertEqual(len(events), 3)
|
||
press, release, wheel = events
|
||
self.assertEqual((press.x, press.y, press.button, press.pressed), (9, 4, 0, True))
|
||
self.assertFalse(release.pressed)
|
||
self.assertIsInstance(wheel, tai.TuiScroll)
|
||
self.assertEqual((wheel.x, wheel.y, wheel.direction), (2, 2, "up"))
|
||
|
||
def test_sgr_modifiers(self):
|
||
_parser, events = self.feed(b"\x1b[<4;1;1M")
|
||
self.assertEqual(events[0].modifiers, ["shift"])
|
||
self.assertEqual(events[0].button, 0)
|
||
|
||
def test_legacy_mouse(self):
|
||
_parser, events = self.feed(b"\x1b[M !!")
|
||
self.assertIsInstance(events[0], tai.TuiClick)
|
||
self.assertEqual((events[0].x, events[0].y), (0, 0))
|
||
_parser, events = self.feed(b"\x1b[M`!!")
|
||
self.assertIsInstance(events[0], tai.TuiScroll)
|
||
|
||
def test_split_report_buffers(self):
|
||
parser, events = self.feed(b"\x1b[<0;1")
|
||
self.assertEqual(events, [])
|
||
events = parser.feed(b"0;5M")
|
||
self.assertEqual(len(events), 1)
|
||
self.assertEqual((events[0].x, events[0].y), (9, 4))
|
||
|
||
def test_keys(self):
|
||
_parser, events = self.feed(b"\x1b[A\x1b[Z\x1bOP\r\t\x7fq")
|
||
self.assertEqual([event.key for event in events], ["up", "backtab", "f1", "enter", "tab", "backspace", "q"])
|
||
|
||
def test_escape_needs_idle(self):
|
||
parser, events = self.feed(b"\x1b")
|
||
self.assertEqual(events, [])
|
||
self.assertEqual([event.key for event in parser.flush()], ["escape"])
|
||
|
||
def test_alt_combo(self):
|
||
_parser, events = self.feed(b"\x1bx")
|
||
self.assertEqual(events[0].key, "alt+x")
|
||
|
||
def test_ctrl_and_utf8(self):
|
||
_parser, events = self.feed(b"\x01\xc3\xa9")
|
||
self.assertEqual(events[0].key, "ctrl+a")
|
||
self.assertEqual(events[1].key, "é")
|
||
|
||
def test_split_utf8(self):
|
||
parser, events = self.feed(b"\xc3")
|
||
self.assertEqual(events, [])
|
||
events = parser.feed(b"\xa9")
|
||
self.assertEqual(events[0].key, "é")
|
||
|
||
def test_modified_arrows(self):
|
||
_parser, events = self.feed(b"\x1b[1;2A")
|
||
self.assertEqual(events[0].key, "shift+up")
|
||
|
||
def test_bracketed_paste(self):
|
||
_parser, events = self.feed(b"\x1b[200~hello\nworld\x1b[201~")
|
||
self.assertEqual(len(events), 1)
|
||
self.assertIsInstance(events[0], tai.TuiPaste)
|
||
self.assertEqual(events[0].text, "hello\nworld")
|
||
|
||
def test_split_paste(self):
|
||
parser, events = self.feed(b"\x1b[200~ab")
|
||
self.assertEqual(events, [])
|
||
events = parser.feed(b"cd\x1b[201~")
|
||
self.assertEqual(len(events), 1)
|
||
self.assertEqual(events[0].text, "abcd")
|
||
|
||
def test_paste_keeps_escapes_literal(self):
|
||
_parser, events = self.feed(b"\x1b[200~\x1b[A\tX\x1b[201~")
|
||
self.assertEqual(len(events), 1)
|
||
self.assertEqual(events[0].text, "\x1b[A\tX")
|
||
|
||
def test_paste_cap_drops_whole(self):
|
||
big = b"x" * (tai.TUI_PASTE_CAP + 1)
|
||
parser, events = self.feed(b"\x1b[200~" + big)
|
||
self.assertEqual(events, [])
|
||
self.assertEqual(parser.buf, b"")
|
||
events = parser.feed(b"\x1b[201~")
|
||
self.assertEqual(events, [])
|
||
|
||
def test_flush_keeps_partial_paste(self):
|
||
parser, events = self.feed(b"\x1b[200~ab")
|
||
self.assertEqual(events, [])
|
||
self.assertEqual(parser.flush(), [])
|
||
events = parser.feed(b"c\x1b[201~")
|
||
self.assertEqual(len(events), 1)
|
||
self.assertEqual(events[0].text, "abc")
|
||
|
||
def test_crlf_single_enter(self):
|
||
_parser, events = self.feed(b"a\r\nb")
|
||
self.assertEqual([event.key for event in events], ["a", "enter", "b"])
|
||
|
||
def test_lone_cr_waits_for_next_byte(self):
|
||
parser, events = self.feed(b"\r")
|
||
self.assertEqual(events, [])
|
||
self.assertEqual([event.key for event in parser.feed(b"x")], ["enter", "x"])
|
||
|
||
def test_lone_cr_resolves_on_flush(self):
|
||
parser, events = self.feed(b"\r")
|
||
self.assertEqual(events, [])
|
||
self.assertEqual([event.key for event in parser.flush()], ["enter"])
|
||
|
||
|
||
class TuiKittyTests(unittest.TestCase):
|
||
def feed(self, *chunks):
|
||
parser = tai.TuiInput()
|
||
events = []
|
||
for chunk in chunks:
|
||
events.extend(parser.feed(chunk))
|
||
return parser, events
|
||
|
||
def test_shift_press(self):
|
||
_parser, events = self.feed(b"\x1b[57441;2u")
|
||
self.assertEqual(len(events), 1)
|
||
self.assertEqual(events[0].key, "shift")
|
||
|
||
def test_shift_right_press(self):
|
||
_parser, events = self.feed(b"\x1b[57447;2u")
|
||
self.assertEqual(events[0].key, "shift")
|
||
|
||
def test_release_ignored(self):
|
||
_parser, events = self.feed(b"\x1b[57441;1:3u")
|
||
self.assertEqual(events, [])
|
||
|
||
def test_repeat_modifier_ignored(self):
|
||
_parser, events = self.feed(b"\x1b[57441;2:2u")
|
||
self.assertEqual(events, [])
|
||
|
||
def test_text_key(self):
|
||
_parser, events = self.feed(b"\x1b[97u")
|
||
self.assertEqual((events[0].key, events[0].char), ("a", "a"))
|
||
|
||
def test_text_repeat_emits(self):
|
||
_parser, events = self.feed(b"\x1b[97;1:2u")
|
||
self.assertEqual(events[0].key, "a")
|
||
|
||
def test_shift_letter_uppercases(self):
|
||
_parser, events = self.feed(b"\x1b[97;2u")
|
||
self.assertEqual((events[0].key, events[0].char), ("A", "A"))
|
||
self.assertEqual(getattr(events[0], "modifiers", []), [])
|
||
|
||
def test_ctrl_letter(self):
|
||
_parser, events = self.feed(b"\x1b[97;5u")
|
||
self.assertEqual(events[0].key, "ctrl+a")
|
||
self.assertEqual(events[0].modifiers, ["ctrl"])
|
||
|
||
def test_functional(self):
|
||
_parser, events = self.feed(b"\x1b[27u\x1b[13u\x1b[9u\x1b[127u")
|
||
self.assertEqual([event.key for event in events], ["escape", "enter", "tab", "backspace"])
|
||
|
||
def test_associated_text_wins(self):
|
||
_parser, events = self.feed(b"\x1b[49;2;33u")
|
||
self.assertEqual((events[0].key, events[0].char), ("!", "!"))
|
||
|
||
def test_split_chunk(self):
|
||
parser, events = self.feed(b"\x1b[57441")
|
||
self.assertEqual(events, [])
|
||
events = parser.feed(b";2u")
|
||
self.assertEqual(len(events), 1)
|
||
self.assertEqual(events[0].key, "shift")
|
||
|
||
def test_unknown_swallowed(self):
|
||
_parser, events = self.feed(b"\x1b[59999uq")
|
||
self.assertEqual([event.key for event in events], ["q"])
|
||
|
||
def test_stack_forms_ignored(self):
|
||
_parser, events = self.feed(b"\x1b[>10u\x1b[?u")
|
||
self.assertEqual(events, [])
|
||
|
||
|
||
class TuiVtTests(unittest.TestCase):
|
||
def test_text_wrap_and_scroll(self):
|
||
screen = tai.TuiVtScreen(2, 4)
|
||
screen.feed(b"abcdef\r\ngh")
|
||
self.assertEqual(screen.text(), "ef\ngh")
|
||
screen.feed(b"\r\nijk")
|
||
self.assertEqual(screen.text(), "gh\nijk")
|
||
|
||
def test_cursor_addressing(self):
|
||
screen = tai.TuiVtScreen(3, 5)
|
||
screen.feed(b"\x1b[2;3HX")
|
||
self.assertEqual(screen.text(), "\n X\n")
|
||
|
||
def test_erase_line_and_display(self):
|
||
screen = tai.TuiVtScreen(2, 5)
|
||
screen.feed(b"hello\r\nworld\x1b[1K")
|
||
self.assertEqual(screen.text(), "hello\n")
|
||
screen.feed(b"\x1b[2J")
|
||
self.assertEqual(screen.text(), "\n")
|
||
|
||
def test_sgr_colors_stored(self):
|
||
screen = tai.TuiVtScreen(1, 3)
|
||
screen.feed(b"\x1b[31mR\x1b[0mN")
|
||
self.assertEqual(screen.grid[0][0][1], "31")
|
||
self.assertIsNone(screen.grid[0][1][1])
|
||
|
||
def test_alt_screen_restores(self):
|
||
screen = tai.TuiVtScreen(2, 4)
|
||
screen.feed(b"main")
|
||
screen.feed(b"\x1b[?1049h")
|
||
screen.feed(b"\x1b[2Jalt!")
|
||
self.assertEqual(screen.text(), "alt!\n")
|
||
screen.feed(b"\x1b[?1049l")
|
||
self.assertEqual(screen.text(), "main\n")
|
||
|
||
def test_resize_keeps_content(self):
|
||
screen = tai.TuiVtScreen(2, 4)
|
||
screen.feed(b"ab\r\ncd")
|
||
screen.resize(2, 6)
|
||
self.assertEqual(screen.text(), "ab\ncd")
|
||
|
||
def test_terminal_key_bytes(self):
|
||
enter = tai.TuiTerminal.key_bytes(tai.TuiKey(key="enter", char="\n"))
|
||
up = tai.TuiTerminal.key_bytes(tai.TuiKey(key="up", char=""))
|
||
char = tai.TuiTerminal.key_bytes(tai.TuiKey(key="a", char="a"))
|
||
ctrl = tai.TuiTerminal.key_bytes(tai.TuiKey(key="ctrl+a", char="\x01"))
|
||
self.assertEqual((enter, up, char, ctrl), (b"\r", b"\x1b[A", b"a", b"\x01"))
|
||
|
||
|
||
class TuiWidgetTests(unittest.TestCase):
|
||
def make_app(self, rows=12, cols=40):
|
||
app = tai.TuiApp()
|
||
app.rows = rows
|
||
app.cols = cols
|
||
return app
|
||
|
||
def test_options_move_clamps_and_chooses(self):
|
||
seen = []
|
||
picker = tai.TuiOptions(["a", "b"], on_select=lambda event: seen.append((event.index, event.value)))
|
||
picker.move(-5)
|
||
self.assertEqual(picker.selected, 0)
|
||
picker.move(5)
|
||
self.assertEqual(picker.selected, 1)
|
||
picker.choose()
|
||
self.assertEqual(seen, [(1, "b")])
|
||
|
||
def test_options_click_selects_row(self):
|
||
app = self.make_app()
|
||
root = tai.TuiPanel(flow=None)
|
||
seen = []
|
||
picker = tai.TuiOptions(["a", "b", "c"], x=0, y=0, w=10, h=5, on_select=lambda event: seen.append(event.value))
|
||
root.add(picker)
|
||
app.mount(root)
|
||
app._dispatch_click(tai.TuiClick(x=1, y=2, button=0, modifiers=[], pressed=True))
|
||
self.assertEqual(seen, ["c"])
|
||
|
||
def test_question_answers_and_closes(self):
|
||
app = self.make_app()
|
||
root = tai.TuiPanel(flow=None)
|
||
seen = []
|
||
root.add(tai.TuiText("behind", x=0, y=0, w=40, h=12))
|
||
question = tai.TuiQuestion("Pick?", ["x", "y"], on_done=seen.append, w=30)
|
||
root.add(question)
|
||
app.mount(root)
|
||
question.picker.choose(1)
|
||
self.assertEqual(seen, ["y"])
|
||
self.assertIsNone(question.parent)
|
||
|
||
def test_dialog_escape_closes(self):
|
||
app = self.make_app()
|
||
root = tai.TuiPanel(flow=None)
|
||
dialog = tai.TuiDialog(TuiWidgetTests._btn(), w=20, h=6)
|
||
root.add(dialog)
|
||
app.mount(root)
|
||
dialog.emit(tai.TuiKey(key="escape", char="\x1b"))
|
||
self.assertIsNone(dialog.parent)
|
||
|
||
def test_terminal_paste_queues_bytes(self):
|
||
term = tai.TuiTerminal()
|
||
term.on_paste(tai.TuiPaste(text="a\nb"))
|
||
self.assertEqual(bytes(term._pending), b"a\nb")
|
||
term.exit_status = 0
|
||
term._pending = bytearray()
|
||
term.on_paste(tai.TuiPaste(text="a\nb"))
|
||
self.assertEqual(bytes(term._pending), b"")
|
||
|
||
@staticmethod
|
||
def _btn():
|
||
return tai.TuiButton("ok", w=8)
|
||
|
||
def test_log_appends_and_pins(self):
|
||
log = tai.TuiLog(w=20, h=4)
|
||
log.arrange(0, 0, 20, 4, (0, 0, 20, 4))
|
||
for pos in range(10):
|
||
log.append("line %d" % pos)
|
||
self.assertEqual(log.scroll_y, 6)
|
||
log.scroll_by(0, -4)
|
||
log.append("line 10")
|
||
self.assertEqual(log.scroll_y, 2)
|
||
|
||
def test_progress_clamps(self):
|
||
bar = tai.TuiProgress(9.9)
|
||
self.assertEqual(bar.value, 1.0)
|
||
bar.set(-1)
|
||
self.assertEqual(bar.value, 0.0)
|
||
|
||
def test_animation_interpolates_and_finishes(self):
|
||
app = self.make_app()
|
||
panel = tai.TuiPanel(x=0, y=0, w=4, h=1)
|
||
app.mount(tai.TuiPanel(flow=None))
|
||
app.root.add(panel)
|
||
done = []
|
||
panel.animate("x", 10, duration=1.0, easing="linear", on_done=done.append)
|
||
app.tick_animations(1000.0)
|
||
self.assertEqual(panel.x, 0)
|
||
app.tick_animations(1000.5)
|
||
self.assertEqual(panel.x, 5.0)
|
||
app.tick_animations(1001.0)
|
||
self.assertEqual(panel.x, 10.0)
|
||
self.assertEqual(done, [panel])
|
||
self.assertEqual(app.animations, [])
|
||
|
||
def test_timers_fire_and_cancel(self):
|
||
app = self.make_app()
|
||
seen = []
|
||
once = app.after(9999, lambda: seen.append("once"))
|
||
once[0] = 0.0
|
||
app._fire_timers(100.0)
|
||
self.assertEqual(seen, ["once"])
|
||
seen.clear()
|
||
app._fire_timers(200.0)
|
||
self.assertEqual(seen, [])
|
||
every = app.every(10, lambda: seen.append("tick"))
|
||
every[0] = 0.0
|
||
app._fire_timers(100.0)
|
||
app._fire_timers(100.0)
|
||
self.assertEqual(seen, ["tick"])
|
||
app.cancel_timer(every)
|
||
every[0] = 0.0
|
||
app._fire_timers(200.0)
|
||
self.assertEqual(seen, ["tick"])
|
||
|
||
def test_demo_scene_composes(self):
|
||
app = tai.TuiApp()
|
||
app.rows = 30
|
||
app.cols = 100
|
||
root, refs = tai.tui_build_demo(app, live=False)
|
||
app.mount(root)
|
||
refs["shell"].screen.feed(b"$ echo hello\r\nhello\r\n$ ")
|
||
refs["stream"].screen.feed(b"stream tick 1\r\nstream tick 2\r\n")
|
||
app.close_panel(root.find("welcome"))
|
||
text = app.render_text()
|
||
for needle in ("Actions", "Open question dialog", "stream tick 2", "hello", "Progress",
|
||
"Events", "overlapping layers", "70%", "demo started"):
|
||
self.assertIn(needle, text)
|
||
|
||
def test_question_renders_modal_content(self):
|
||
app = tai.TuiApp()
|
||
app.rows = 16
|
||
app.cols = 50
|
||
root = tai.TuiPanel(flow=None)
|
||
root.add(tai.TuiText("behind", x=0, y=0, w=50, h=16))
|
||
root.add(tai.TuiQuestion("Pick a theme:", ["Ocean", "Forest"], on_done=lambda value: None, title="Question", w=30))
|
||
app.mount(root)
|
||
text = app.render_text()
|
||
for needle in ("Question", "Pick a theme:", "Ocean", "Forest", "Cancel"):
|
||
self.assertIn(needle, text)
|
||
|
||
|
||
class TuiDoubleTapTests(unittest.TestCase):
|
||
def test_fires_on_second_within_window(self):
|
||
tap = tai.TuiDoubleTap(timeout=0.4)
|
||
self.assertFalse(tap.hit("shift", 100.0))
|
||
self.assertTrue(tap.hit("shift", 100.2))
|
||
|
||
def test_no_fire_across_window(self):
|
||
tap = tai.TuiDoubleTap(timeout=0.4)
|
||
self.assertFalse(tap.hit("shift", 100.0))
|
||
self.assertFalse(tap.hit("shift", 100.5))
|
||
self.assertTrue(tap.hit("shift", 100.6))
|
||
|
||
def test_triple_fires_once(self):
|
||
tap = tai.TuiDoubleTap(timeout=0.4)
|
||
self.assertFalse(tap.hit("shift", 100.0))
|
||
self.assertTrue(tap.hit("shift", 100.1))
|
||
self.assertFalse(tap.hit("shift", 100.2))
|
||
|
||
def test_other_key_resets(self):
|
||
tap = tai.TuiDoubleTap(timeout=0.4)
|
||
self.assertFalse(tap.hit("shift", 100.0))
|
||
self.assertFalse(tap.hit("ctrl", 100.1))
|
||
self.assertFalse(tap.hit("shift", 100.2))
|
||
self.assertTrue(tap.hit("shift", 100.3))
|
||
|
||
def test_default_timeout(self):
|
||
self.assertEqual(tai.TuiDoubleTap().timeout, 0.4)
|
||
|
||
|
||
class TuiTextBoxTests(unittest.TestCase):
|
||
def make_box(self, **kwargs):
|
||
app = tai.TuiApp()
|
||
app.rows = 10
|
||
app.cols = 40
|
||
root = tai.TuiPanel(flow=None)
|
||
kwargs.setdefault("x", 1)
|
||
kwargs.setdefault("y", 1)
|
||
kwargs.setdefault("w", 20)
|
||
kwargs.setdefault("h", 3)
|
||
box = tai.TuiTextBox(**kwargs)
|
||
root.add(box)
|
||
app.mount(root)
|
||
return app, box
|
||
|
||
def key(self, app, key, char=""):
|
||
app.dispatch(tai.TuiKey(key=key, char=char))
|
||
|
||
def test_insert_and_value(self):
|
||
app, box = self.make_box()
|
||
self.assertIs(app.focused, box)
|
||
for char in "abc":
|
||
self.key(app, char, char)
|
||
self.assertEqual(box.value, "abc")
|
||
self.assertEqual((box.cx, box.cy), (3, 0))
|
||
|
||
def test_enter_splits_line(self):
|
||
app, box = self.make_box()
|
||
for char in "ab":
|
||
self.key(app, char, char)
|
||
self.key(app, "left", "")
|
||
self.key(app, "enter", "\n")
|
||
self.assertEqual(box.value, "a\nb")
|
||
self.assertEqual((box.cx, box.cy), (0, 1))
|
||
|
||
def test_backspace_joins_lines(self):
|
||
app, box = self.make_box()
|
||
for char in "a":
|
||
self.key(app, char, char)
|
||
self.key(app, "enter", "\n")
|
||
self.key(app, "b", "b")
|
||
self.key(app, "left", "")
|
||
self.key(app, "left", "")
|
||
self.key(app, "backspace", "\x7f")
|
||
self.assertEqual(box.value, "ab")
|
||
self.assertEqual((box.cx, box.cy), (1, 0))
|
||
|
||
def test_delete_joins_lines(self):
|
||
app, box = self.make_box()
|
||
for char in "a":
|
||
self.key(app, char, char)
|
||
self.key(app, "enter", "\n")
|
||
self.key(app, "b", "b")
|
||
self.key(app, "up", "")
|
||
self.key(app, "delete", "")
|
||
self.assertEqual(box.value, "ab")
|
||
self.assertEqual((box.cx, box.cy), (1, 0))
|
||
|
||
def test_home_end_arrows(self):
|
||
app, box = self.make_box()
|
||
for char in "abc":
|
||
self.key(app, char, char)
|
||
self.key(app, "home", "")
|
||
self.assertEqual(box.cx, 0)
|
||
self.key(app, "end", "")
|
||
self.assertEqual(box.cx, 3)
|
||
self.key(app, "left", "")
|
||
self.assertEqual(box.cx, 2)
|
||
|
||
def test_stop_and_bubble(self):
|
||
_app, box = self.make_box()
|
||
event = tai.TuiKey(key="a", char="a")
|
||
box.emit(event)
|
||
self.assertTrue(event.stopped)
|
||
event = tai.TuiKey(key="escape", char="\x1b")
|
||
box.emit(event)
|
||
self.assertFalse(event.stopped)
|
||
|
||
def test_click_positions_cursor(self):
|
||
app, box = self.make_box()
|
||
for char in "hello":
|
||
self.key(app, char, char)
|
||
app.layout()
|
||
app._dispatch_click(tai.TuiClick(x=4, y=2, button=0, modifiers=[], pressed=True))
|
||
self.assertIs(app.focused, box)
|
||
self.assertEqual((box.cx, box.cy), (2, 0))
|
||
|
||
def test_space_inserts(self):
|
||
app, box = self.make_box()
|
||
self.key(app, "space", " ")
|
||
self.assertEqual(box.value, " ")
|
||
|
||
def test_paste_multiline(self):
|
||
app, box = self.make_box()
|
||
app.dispatch(tai.TuiPaste(text="a\r\nb\rc"))
|
||
self.assertEqual(box.value, "a\nb\nc")
|
||
self.assertEqual((box.cx, box.cy), (1, 2))
|
||
|
||
|
||
class TuiAutocompleteTests(unittest.TestCase):
|
||
def make_box(self, complete=tai.tui_complete_command, **kwargs):
|
||
app = tai.TuiApp()
|
||
app.rows = 12
|
||
app.cols = 50
|
||
root = tai.TuiPanel(flow=None)
|
||
kwargs.setdefault("x", 1)
|
||
kwargs.setdefault("y", 1)
|
||
kwargs.setdefault("w", 30)
|
||
kwargs.setdefault("h", 3)
|
||
box = tai.TuiTextBox(complete=complete, **kwargs)
|
||
root.add(box)
|
||
app.mount(root)
|
||
return app, box
|
||
|
||
def key(self, app, key, char=""):
|
||
app.dispatch(tai.TuiKey(key=key, char=char))
|
||
|
||
def type(self, app, text):
|
||
for char in text:
|
||
self.key(app, "space" if char == " " else char, char)
|
||
|
||
def test_opens_and_filters(self):
|
||
app, box = self.make_box()
|
||
self.type(app, "/bo")
|
||
popup = box.complete_popup
|
||
self.assertIsNotNone(popup)
|
||
self.assertEqual([value for _label, value in popup.options], ["bots", "bot"])
|
||
self.assertIs(popup.parent, box.root())
|
||
self.assertIs(app.focused, box)
|
||
|
||
def test_accept_replaces_word(self):
|
||
app, box = self.make_box()
|
||
self.type(app, "/bo")
|
||
self.key(app, "down", "")
|
||
self.key(app, "enter", "\n")
|
||
self.assertEqual(box.value, "/bot ")
|
||
self.assertEqual(box.cx, 5)
|
||
self.assertIsNone(box.complete_popup)
|
||
|
||
def test_accept_mid_word_replaces_through_end(self):
|
||
app, box = self.make_box()
|
||
self.type(app, "/bot x")
|
||
for _i in range(5):
|
||
self.key(app, "left", "")
|
||
self.key(app, "enter", "\n")
|
||
self.assertEqual(box.value, "/help x")
|
||
self.assertEqual(box.cx, 5)
|
||
self.assertIsNone(box.complete_popup)
|
||
|
||
def test_escape_dismisses_then_bubbles(self):
|
||
app, box = self.make_box()
|
||
self.type(app, "/bo")
|
||
event = tai.TuiKey(key="escape", char="\x1b")
|
||
box.emit(event)
|
||
self.assertTrue(event.stopped)
|
||
self.assertIsNone(box.complete_popup)
|
||
event = tai.TuiKey(key="escape", char="\x1b")
|
||
box.emit(event)
|
||
self.assertFalse(event.stopped)
|
||
|
||
def test_no_match_hides(self):
|
||
app, box = self.make_box()
|
||
self.type(app, "/zzz")
|
||
self.assertIsNone(box.complete_popup)
|
||
|
||
def test_space_hides(self):
|
||
app, box = self.make_box()
|
||
self.type(app, "/bot")
|
||
self.assertIsNotNone(box.complete_popup)
|
||
self.type(app, " ")
|
||
self.assertIsNone(box.complete_popup)
|
||
|
||
def test_up_down_navigate(self):
|
||
app, box = self.make_box()
|
||
self.type(app, "/bo")
|
||
popup = box.complete_popup
|
||
self.assertEqual(popup.selected, 0)
|
||
self.key(app, "down", "")
|
||
self.assertEqual(popup.selected, 1)
|
||
self.key(app, "down", "")
|
||
self.assertEqual(popup.selected, 1)
|
||
self.key(app, "up", "")
|
||
self.assertEqual(popup.selected, 0)
|
||
self.assertEqual(box.value, "/bo")
|
||
|
||
def test_popup_placed_under_cursor(self):
|
||
app, box = self.make_box()
|
||
self.type(app, "/bo")
|
||
popup = box.complete_popup
|
||
self.assertEqual((popup.x, popup.y), (4, 3))
|
||
|
||
def test_popup_flips_above_when_no_room(self):
|
||
app, box = self.make_box(y=8)
|
||
self.type(app, "/bo")
|
||
popup = box.complete_popup
|
||
self.assertEqual(popup.y, 5)
|
||
|
||
def test_click_accepts(self):
|
||
app, box = self.make_box()
|
||
self.type(app, "/bo")
|
||
app.layout()
|
||
popup = box.complete_popup
|
||
app._dispatch_click(tai.TuiClick(x=popup.ax + 1, y=popup.ay + 1, button=0, modifiers=[], pressed=True))
|
||
self.assertEqual(box.value, "/bots ")
|
||
|
||
def test_blur_hides(self):
|
||
app, box = self.make_box()
|
||
self.type(app, "/bo")
|
||
self.assertIsNotNone(box.complete_popup)
|
||
box.emit(tai.TuiBlur(panel=box))
|
||
self.assertIsNone(box.complete_popup)
|
||
|
||
def test_unmount_removes_popup(self):
|
||
app, box = self.make_box()
|
||
self.type(app, "/bo")
|
||
popup = box.complete_popup
|
||
box.parent.remove(box)
|
||
self.assertIsNone(popup.parent)
|
||
|
||
def test_completer_failure_hides(self):
|
||
def broken(_prefix):
|
||
raise ValueError("nope")
|
||
|
||
app, box = self.make_box(complete=broken)
|
||
self.type(app, "/bo")
|
||
self.assertIsNone(box.complete_popup)
|
||
self.assertEqual(box.value, "/bo")
|
||
|
||
def test_tui_complete_command_labels(self):
|
||
items = tai.tui_complete_command("bo")
|
||
self.assertEqual([value for _label, value in items], ["bots", "bot"])
|
||
self.assertIn("list bots", items[0][0])
|
||
plain = [label for label, value in tai.tui_complete_command("exi") if value == "exit"]
|
||
self.assertEqual(plain, ["/exit"])
|
||
|
||
def test_mention_provider_accepts(self):
|
||
app, box = self.make_box(mentions=tai.tui_complete_mention)
|
||
self.type(app, "@op")
|
||
popup = box.complete_popup
|
||
self.assertIsNotNone(popup)
|
||
self.assertEqual(popup.options[0][1], "opencode")
|
||
self.key(app, "enter", "\n")
|
||
self.assertEqual(box.value, "@opencode ")
|
||
|
||
def test_mention_file_accepts(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
with open(os.path.join(tmp, "doc.txt"), "w") as handle:
|
||
handle.write("x")
|
||
app, box = self.make_box(mentions=lambda prefix: tai.tui_complete_mention(prefix, directory=tmp))
|
||
self.type(app, "@do")
|
||
popup = box.complete_popup
|
||
self.assertIsNotNone(popup)
|
||
self.assertEqual([value for _label, value in popup.options], ["doc.txt"])
|
||
self.key(app, "enter", "\n")
|
||
self.assertEqual(box.value, "@doc.txt ")
|
||
|
||
def test_mention_needs_fn(self):
|
||
app, box = self.make_box(mentions=None, complete=None)
|
||
self.type(app, "@op")
|
||
self.assertIsNone(box.complete_popup)
|
||
|
||
|
||
class TuiPromptDialogTests(unittest.TestCase):
|
||
def make_demo(self):
|
||
app = tai.TuiApp()
|
||
app.rows = 30
|
||
app.cols = 100
|
||
root, refs = tai.tui_build_demo(app, live=False)
|
||
app.mount(root)
|
||
app.dispatch(tai.TuiKey(key="enter", char="\r"))
|
||
app.double_tap.timeout = 3600
|
||
return app, root, refs
|
||
|
||
def test_double_shift_opens(self):
|
||
app, root, _refs = self.make_demo()
|
||
app.dispatch(tai.TuiKey(key="shift", char=""))
|
||
self.assertIsNone(root.find("prompt"))
|
||
app.dispatch(tai.TuiKey(key="shift", char=""))
|
||
dialog = root.find("prompt")
|
||
self.assertIsNotNone(dialog)
|
||
self.assertIs(app.focused, dialog.box)
|
||
|
||
def test_ok_executes(self):
|
||
app, root, refs = self.make_demo()
|
||
app.dispatch(tai.TuiKey(key="shift", char=""))
|
||
app.dispatch(tai.TuiKey(key="shift", char=""))
|
||
for char in "hi":
|
||
app.dispatch(tai.TuiKey(key=char, char=char))
|
||
app.dispatch(tai.TuiKey(key="tab", char="\t"))
|
||
app.dispatch(tai.TuiKey(key="enter", char="\r"))
|
||
self.assertIsNone(root.find("prompt"))
|
||
self.assertEqual(bytes(refs["shell"]._pending), b"hi\n")
|
||
|
||
def test_cancel_closes_without_executing(self):
|
||
app, root, refs = self.make_demo()
|
||
app.dispatch(tai.TuiKey(key="shift", char=""))
|
||
app.dispatch(tai.TuiKey(key="shift", char=""))
|
||
app.dispatch(tai.TuiKey(key="x", char="x"))
|
||
app.dispatch(tai.TuiKey(key="escape", char="\x1b"))
|
||
self.assertIsNone(root.find("prompt"))
|
||
self.assertEqual(bytes(refs["shell"]._pending), b"")
|
||
|
||
def test_double_tab_still_cycles(self):
|
||
app, root, _refs = self.make_demo()
|
||
actions = root.find("actions")
|
||
self.assertIs(app.focused, actions)
|
||
app.dispatch(tai.TuiKey(key="tab", char="\t"))
|
||
app.dispatch(tai.TuiKey(key="tab", char="\t"))
|
||
self.assertIs(app.focused, actions)
|
||
self.assertIsNone(root.find("prompt"))
|
||
|
||
def test_complete_passthrough(self):
|
||
dialog = tai.TuiPromptDialog(complete=tai.tui_complete_command)
|
||
self.assertIs(dialog.box.complete_fn, tai.tui_complete_command)
|
||
|
||
def test_demo_prompt_has_command_completion(self):
|
||
app, root, _refs = self.make_demo()
|
||
app.dispatch(tai.TuiKey(key="shift", char=""))
|
||
app.dispatch(tai.TuiKey(key="shift", char=""))
|
||
dialog = root.find("prompt")
|
||
self.assertIs(dialog.box.complete_fn, tai.tui_complete_command)
|
||
|
||
|
||
class TuiCommandTests(unittest.TestCase):
|
||
def test_demo_in_commands_and_help(self):
|
||
self.assertIn("demo", tai.COMMANDS)
|
||
self.assertIn("/demo", tai.HELP_TEXT)
|
||
|
||
def test_parser_demo_flags(self):
|
||
parser = tai.build_parser()
|
||
args = parser.parse_args(["--demo"])
|
||
self.assertTrue(args.demo)
|
||
self.assertFalse(args.demo_shot)
|
||
args = parser.parse_args(["--demo-shot"])
|
||
self.assertTrue(args.demo_shot)
|
||
|
||
|
||
class TuiThemeTests(unittest.TestCase):
|
||
def tearDown(self):
|
||
tai.tui_set_theme("dark")
|
||
|
||
def test_builtin_registry_complete(self):
|
||
self.assertEqual(set(tai.TUI_THEMES), {"dark", "light", "nord", "gruvbox", "ansi"})
|
||
for name, theme in tai.TUI_THEMES.items():
|
||
for role in tai.TUI_THEME_ROLES:
|
||
self.assertTrue(theme.color(role), "%s lacks %s" % (name, role))
|
||
self.assertFalse(tai.TUI_THEMES["light"].dark)
|
||
self.assertTrue(tai.TUI_THEMES["dark"].dark)
|
||
|
||
def test_role_resolves_at_draw(self):
|
||
tai.tui_set_theme("dark")
|
||
self.assertEqual(tai.TuiStyle(fg="role:accent").sgr(), "\x1b[0;38;2;139;92;246m")
|
||
self.assertEqual(tai.TuiStyle(bg="role:bg").sgr(), "\x1b[0;48;2;10;14;26m")
|
||
tai.tui_set_theme("light")
|
||
self.assertEqual(tai.TuiStyle(fg="role:accent").sgr(), "\x1b[0;38;2;109;74;255m")
|
||
|
||
def test_unknown_role_is_empty(self):
|
||
self.assertEqual(tai.TuiStyle(fg="role:nope").sgr(), "\x1b[0m")
|
||
|
||
def test_raw_colors_pass_through(self):
|
||
self.assertEqual(tai.TuiStyle(fg="#ff0000").sgr(), "\x1b[0;38;2;255;0;0m")
|
||
self.assertEqual(tai.TuiStyle(fg="cyan").sgr(), "\x1b[0;36m")
|
||
|
||
def test_register_validates_roles(self):
|
||
with self.assertRaises(ValueError):
|
||
tai.tui_register_theme(tai.TuiTheme("broken", True, {"bg": "black"}))
|
||
roles = {role: "white" for role in tai.TUI_THEME_ROLES}
|
||
tai.tui_register_theme(tai.TuiTheme("custom-1", True, roles))
|
||
try:
|
||
tai.tui_set_theme("custom-1")
|
||
self.assertEqual(tai.TuiStyle(fg="role:fg").sgr(), "\x1b[0;37m")
|
||
finally:
|
||
del tai.TUI_THEMES["custom-1"]
|
||
|
||
def test_resolve_precedence(self):
|
||
self.assertEqual(tai.tui_resolve_theme_name("nord"), "nord")
|
||
old = os.environ.get("TAI_THEME")
|
||
os.environ["TAI_THEME"] = "gruvbox"
|
||
try:
|
||
self.assertEqual(tai.tui_resolve_theme_name(None), "gruvbox")
|
||
self.assertEqual(tai.tui_resolve_theme_name("light"), "light")
|
||
finally:
|
||
if old is None:
|
||
os.environ.pop("TAI_THEME", None)
|
||
else:
|
||
os.environ["TAI_THEME"] = old
|
||
self.assertEqual(tai.tui_resolve_theme_name(None), "dark")
|
||
with self.assertRaises(ValueError) as raised:
|
||
tai.tui_resolve_theme_name("nope")
|
||
self.assertIn("dark", str(raised.exception))
|
||
|
||
def test_set_theme_redraws(self):
|
||
app = tai.TuiApp()
|
||
app._need_paint = False
|
||
app.set_theme("nord")
|
||
self.assertEqual(tai.tui_theme().name, "nord")
|
||
self.assertTrue(app._need_paint)
|
||
with self.assertRaises(ValueError):
|
||
app.set_theme("nope")
|
||
|
||
def test_cycle_order(self):
|
||
self.assertEqual(tai.tui_cycle_theme_name("dark"), "light")
|
||
self.assertEqual(tai.tui_cycle_theme_name("ansi"), "dark")
|
||
self.assertEqual(tai.tui_cycle_theme_name("nope"), "dark")
|
||
|
||
|
||
class FakeChatChat:
|
||
provider_pin = None
|
||
model_pin = None
|
||
|
||
|
||
class FakeChatStore:
|
||
def list_bots(self, profile):
|
||
return [{"name": "main", "nicknames": []}, {"name": "coder", "nicknames": []}]
|
||
|
||
def load_session(self, profile, bot):
|
||
return []
|
||
|
||
|
||
class FakeChatAgent:
|
||
profile = "t"
|
||
bot = "main"
|
||
|
||
def __init__(self):
|
||
self.chat = FakeChatChat()
|
||
self.store = FakeChatStore()
|
||
self.messages = [{"role": "system", "content": "sys"},
|
||
{"role": "user", "content": "hello"},
|
||
{"role": "assistant", "content": "hi there", "tool_calls": []}]
|
||
self.turns = []
|
||
|
||
def switch_bot(self, name, silent=False):
|
||
self.bot = name
|
||
return "switched to bot '%s'" % name
|
||
|
||
def run_turn(self, text, capture=False):
|
||
self.turns.append(text)
|
||
return "echo: %s" % text
|
||
|
||
def checkpoint(self):
|
||
pass
|
||
|
||
|
||
class TuiChatTests(unittest.TestCase):
|
||
def tearDown(self):
|
||
tai.tui_set_theme("dark")
|
||
|
||
def make_chat(self, cols=100, rows=24):
|
||
app = tai.TuiApp()
|
||
app.rows = rows
|
||
app.cols = cols
|
||
agent = FakeChatAgent()
|
||
root, refs = tai.tui_build_chat(app, agent)
|
||
app.mount(root)
|
||
refs["replay"]()
|
||
return app, agent, root, refs
|
||
|
||
def key(self, app, key, char=""):
|
||
app.dispatch(tai.TuiKey(key=key, char=char))
|
||
|
||
def type(self, app, text):
|
||
for char in text:
|
||
self.key(app, char, char)
|
||
|
||
def test_prefer_puts_devplace_molodetz_first(self):
|
||
chat = tai.ChatClient.__new__(tai.ChatClient)
|
||
chat.provider_pin = chat.model_pin = None
|
||
chat.prefer = (tai.DEVPLACE_LABEL, "molodetz")
|
||
chat.config = type("C", (), {"devplace_key": "k", "devplace_base": "https://x/v1"})()
|
||
other = {"provider": "openrouter", "model": "a", "chat_url": "u", "models_url": "m", "key": "", "profile": ""}
|
||
out = chat._apply_pins([other])
|
||
self.assertEqual((out[0]["provider"], out[0]["model"]), (tai.DEVPLACE_LABEL, "molodetz"))
|
||
self.assertIs(out[1], other)
|
||
|
||
def test_sidebar_width_clamps(self):
|
||
self.assertEqual(tai.tui_chat_sidebar_w(100), 30)
|
||
self.assertEqual(tai.tui_chat_sidebar_w(200), 36)
|
||
self.assertEqual(tai.tui_chat_sidebar_w(60), 24)
|
||
self.assertEqual(tai.tui_chat_sidebar_w(80), 24)
|
||
|
||
def test_layout_and_render(self):
|
||
app, _agent, _root, refs = self.make_chat()
|
||
self.assertEqual(refs["sidebar_wrap"].w, 30)
|
||
self.assertEqual(refs["bots"], ["main", "coder"])
|
||
self.assertIn("t/main", refs["status"].text)
|
||
text = app.render_text()
|
||
self.assertIn("Bots", text)
|
||
self.assertIn("coder", text)
|
||
self.assertIn("hello", text)
|
||
self.assertIn("Write a message...", text)
|
||
self.assertIn("main", text)
|
||
self.assertIn("✦ tai", text)
|
||
self.assertIn("online", text)
|
||
|
||
def test_long_transcript_shows_tail_not_blank(self):
|
||
app, _agent, _root, refs = self.make_chat(rows=24)
|
||
refs["box"].focus()
|
||
for n in range(60):
|
||
refs["show"]("bot", "line number %d" % n, None)
|
||
transcript = refs["transcript"]
|
||
self.assertGreater(len(transcript.lines), transcript._content_h())
|
||
text = app.render_text()
|
||
self.assertIn("line number 59", text)
|
||
self.assertNotIn("line number 0 ", text)
|
||
|
||
def test_many_messages_stay_pinned_and_scroll(self):
|
||
app, _agent, _root, refs = self.make_chat(rows=24)
|
||
refs["box"].focus()
|
||
transcript = refs["transcript"]
|
||
for n in range(tai.TUI_CHAT_LINE_CAP):
|
||
self.type(app, "m%d" % n)
|
||
self.key(app, "enter")
|
||
refs["run_pending"]()
|
||
self.assertLessEqual(len(transcript.lines), tai.TUI_CHAT_LINE_CAP + 10)
|
||
self.assertIn("echo: m%d" % (tai.TUI_CHAT_LINE_CAP - 1), app.render_text())
|
||
transcript.scroll_by(0, -10)
|
||
self.assertNotIn("echo: m%d" % (tai.TUI_CHAT_LINE_CAP - 1), app.render_text())
|
||
transcript.scroll_by(0, 10)
|
||
self.assertIn("echo: m%d" % (tai.TUI_CHAT_LINE_CAP - 1), app.render_text())
|
||
|
||
def test_narrow_hides_and_restores_sidebar(self):
|
||
app, _agent, root, refs = self.make_chat(cols=100)
|
||
self.assertTrue(refs["sidebar_wrap"].visible)
|
||
app.cols = 70
|
||
app.dispatch(tai.TuiResize(panel=root, rows=24, cols=70))
|
||
self.assertFalse(refs["sidebar_wrap"].visible)
|
||
app.cols = 100
|
||
app.dispatch(tai.TuiResize(panel=root, rows=24, cols=100))
|
||
self.assertTrue(refs["sidebar_wrap"].visible)
|
||
|
||
def test_hidden_sidebar_frees_its_columns(self):
|
||
app, _agent, _root, refs = self.make_chat(cols=100)
|
||
right = refs["transcript"].parent
|
||
self.assertEqual(right.ax, 30)
|
||
refs["box"].focus()
|
||
self.key(app, "f2")
|
||
app.layout()
|
||
self.assertEqual(right.ax, 0)
|
||
self.assertEqual(right.aw, 100)
|
||
|
||
def test_narrow_start_hides_sidebar_and_compacts_send(self):
|
||
app, _agent, _root, refs = self.make_chat(cols=44)
|
||
self.assertFalse(refs["sidebar_wrap"].visible)
|
||
self.assertEqual(refs["send_btn"].w, 5)
|
||
|
||
def test_f2_manual_toggle_overrides_auto(self):
|
||
app, _agent, root, refs = self.make_chat(cols=100)
|
||
refs["box"].focus()
|
||
self.key(app, "f2")
|
||
self.assertFalse(refs["sidebar_wrap"].visible)
|
||
app.dispatch(tai.TuiResize(panel=root, rows=24, cols=100))
|
||
self.assertFalse(refs["sidebar_wrap"].visible)
|
||
self.key(app, "f2")
|
||
self.assertTrue(refs["sidebar_wrap"].visible)
|
||
|
||
def test_f4_cycles_theme(self):
|
||
app, _agent, _root, refs = self.make_chat()
|
||
refs["box"].focus()
|
||
self.key(app, "f4")
|
||
self.assertEqual(tai.tui_theme().name, "light")
|
||
self.assertIn("theme: light", refs["status"].text)
|
||
|
||
def test_enter_sends_and_turn_runs(self):
|
||
app, agent, _root, refs = self.make_chat()
|
||
refs["box"].focus()
|
||
self.type(app, "hi bot")
|
||
self.key(app, "enter")
|
||
self.assertEqual(refs["pending"], "hi bot")
|
||
self.assertEqual(refs["box"].value, "")
|
||
self.assertIn("thinking", refs["status"].text)
|
||
refs["run_pending"]()
|
||
self.assertEqual(agent.turns, ["hi bot"])
|
||
shown = [line for line, _spans in refs["transcript"].lines]
|
||
self.assertTrue(any("hi bot" in line for line in shown))
|
||
self.assertTrue(any("echo: hi bot" in line for line in shown))
|
||
self.assertNotIn("thinking", refs["status"].text)
|
||
|
||
def test_turn_stdout_is_buried(self):
|
||
_app, agent, _root, refs = self.make_chat()
|
||
|
||
def noisy(text, capture=False):
|
||
print("CARD-JUNK-TO-STDOUT")
|
||
agent.turns.append(text)
|
||
return "done"
|
||
|
||
agent.run_turn = noisy
|
||
refs["pending"] = "go"
|
||
buf = io.StringIO()
|
||
with contextlib.redirect_stdout(buf):
|
||
refs["run_pending"]()
|
||
self.assertEqual(buf.getvalue(), "")
|
||
shown = [line for line, _spans in refs["transcript"].lines]
|
||
self.assertTrue(any("done" in line for line in shown))
|
||
|
||
def test_send_ignores_empty_and_busy(self):
|
||
_app, _agent, _root, refs = self.make_chat()
|
||
refs["do_send"](" ")
|
||
self.assertIsNone(refs["pending"])
|
||
refs["pending"] = "one"
|
||
refs["do_send"]("two")
|
||
self.assertEqual(refs["pending"], "one")
|
||
|
||
def test_newline_keys_insert_in_send_mode(self):
|
||
app, _agent, _root, refs = self.make_chat()
|
||
refs["box"].focus()
|
||
self.key(app, "alt+enter", "\n")
|
||
self.key(app, "shift+enter", "")
|
||
self.assertEqual(len(refs["box"].lines), 3)
|
||
self.assertIsNone(refs["pending"])
|
||
|
||
def test_completion_enter_accepts_instead_of_send(self):
|
||
app, _agent, _root, refs = self.make_chat()
|
||
refs["box"].focus()
|
||
self.type(app, "/he")
|
||
self.assertIsNotNone(refs["box"].complete_popup)
|
||
self.key(app, "enter")
|
||
self.assertIsNone(refs["pending"])
|
||
self.assertTrue(refs["box"].value.startswith("/help"))
|
||
|
||
def test_max_chars_truncates(self):
|
||
_app, _agent, _root, refs = self.make_chat()
|
||
refs["box"].insert("x" * 2005)
|
||
self.assertEqual(len(refs["box"].value), 2000)
|
||
|
||
def test_slash_help_and_quit(self):
|
||
app, agent, _root, refs = self.make_chat()
|
||
refs["do_send"]("/help")
|
||
shown = "\n".join(line for line, _style in refs["transcript"].lines)
|
||
self.assertIn("this text", shown)
|
||
self.assertEqual(agent.turns, [])
|
||
refs["do_send"]("/quit")
|
||
self.assertFalse(app._running)
|
||
|
||
def test_bot_switch_keeps_drafts(self):
|
||
app, agent, _root, refs = self.make_chat()
|
||
refs["box"].focus()
|
||
self.type(app, "draft-1")
|
||
refs["sidebar"].choose(1)
|
||
self.assertEqual(agent.bot, "coder")
|
||
self.assertEqual(refs["box"].value, "")
|
||
self.assertEqual(refs["head_name"].text, "coder")
|
||
refs["sidebar"].choose(0)
|
||
self.assertEqual(agent.bot, "main")
|
||
self.assertEqual(refs["box"].value, "draft-1")
|
||
|
||
def test_strip_ansi(self):
|
||
self.assertEqual(tai.tui_strip_ansi("\x1b[1;32mok\x1b[0m"), "ok")
|
||
self.assertEqual(tai.tui_strip_ansi("plain"), "plain")
|
||
|
||
def bubble_styles(self):
|
||
plain = tai.TuiStyle()
|
||
return {"base": plain, "bubble": plain, "meta": plain, "tick": plain}
|
||
|
||
def row_text(self, row):
|
||
return "".join(span.text for span in row)
|
||
|
||
def test_bubble_rows_align(self):
|
||
rows = tai.tui_chat_bubble_rows("hi", 60, "out", self.bubble_styles(), "12:00", tai.TUI_CHAT_TICK_READ, True)
|
||
texts = [self.row_text(row) for row in rows]
|
||
self.assertTrue(all(len(text) == 60 for text in texts))
|
||
self.assertEqual(len(texts), 3)
|
||
self.assertTrue(texts[0].startswith(" "))
|
||
self.assertIn("hi", texts[1])
|
||
self.assertTrue(texts[2].rstrip().endswith("12:00 " + tai.TUI_CHAT_TICK_READ))
|
||
rows = tai.tui_chat_bubble_rows("hi", 60, "in", self.bubble_styles(), sender="main")
|
||
texts = [self.row_text(row) for row in rows]
|
||
self.assertTrue(texts[0].startswith(" main"))
|
||
self.assertTrue(texts[1].startswith(" hi"))
|
||
self.assertTrue(texts[1].endswith(" "))
|
||
|
||
def test_bubble_rows_cap_width(self):
|
||
rows = tai.tui_chat_bubble_rows("x" * 100, 60, "in", self.bubble_styles())
|
||
self.assertTrue(len(rows) > 3)
|
||
body = [self.row_text(row) for row in rows][1:-1]
|
||
self.assertEqual(len(body[0]), 60)
|
||
self.assertEqual(body[0].strip(), "x" * 38)
|
||
self.assertEqual("".join(line.strip() for line in body), "x" * 100)
|
||
|
||
def test_bubble_rows_meta_row_keeps_body(self):
|
||
rows = tai.tui_chat_bubble_rows("first line here with padding words second line tail end", 60, "out",
|
||
self.bubble_styles(), "12:00", "x")
|
||
texts = [self.row_text(row) for row in rows]
|
||
self.assertTrue(any("first line here" in line for line in texts))
|
||
self.assertTrue(any("tail end" in line for line in texts))
|
||
self.assertIn("12:00", texts[-1])
|
||
|
||
def test_bubble_rows_narrow_degrade(self):
|
||
rows = tai.tui_chat_bubble_rows("hi there", 20, "out", self.bubble_styles(), "12:00", "x")
|
||
text = self.row_text(rows[0])
|
||
self.assertTrue(text.startswith("hi there"))
|
||
self.assertIn("12:00", self.row_text(rows[-1]))
|
||
|
||
def test_markdown_blocks(self):
|
||
blocks = tai.tui_md_blocks("# Title\n\nplain **bold** and `code`\n- one\n2. two\n> quote\n```py\nx = 1\n```\n---")
|
||
kinds = [kind for kind, _segs, _lead in blocks]
|
||
self.assertEqual(kinds, ["h", "blank", "p", "li", "li", "quote", "code", "rule"])
|
||
self.assertEqual(blocks[2][1], [("plain ", ""), ("bold", "b"), (" and ", ""), ("code", "c")])
|
||
self.assertEqual(blocks[3][2], "•")
|
||
self.assertEqual(blocks[4][2], "2.")
|
||
|
||
def test_wrap_segs_chops_long_words_and_keeps_tags(self):
|
||
lines = tai.tui_wrap_segs([("aaaa bbbb ", ""), ("cccccccccc", "b")], 6)
|
||
self.assertEqual(["".join(text for text, _tag in line) for line in lines], ["aaaa", "bbbb", "cccccc", "cccc"])
|
||
self.assertEqual(lines[2][0][1], "b")
|
||
|
||
def test_markdown_markers_not_shown_in_bubble(self):
|
||
rows = tai.tui_chat_bubble_rows("use **bold** and `code`", 60, "in", self.bubble_styles())
|
||
text = " ".join(self.row_text(row) for row in rows)
|
||
self.assertIn("use bold and code", text)
|
||
self.assertNotIn("**", text)
|
||
self.assertNotIn("`", text)
|
||
|
||
def test_tool_chip_summarizes_json_and_truncates(self):
|
||
row = tai.tui_chat_tool_row("fork", '{"task": "do it", "n": 2}', 60, tai.TuiStyle())
|
||
text = self.row_text(row)
|
||
self.assertEqual(len(text), 60)
|
||
self.assertIn("fork", text)
|
||
self.assertIn("task=do it", text)
|
||
self.assertIn("n=2", text)
|
||
short = self.row_text(tai.tui_chat_tool_row("fork", '{"task": "' + "z" * 200 + '"}', 40, tai.TuiStyle()))
|
||
self.assertEqual(len(short), 40)
|
||
self.assertIn("…", short)
|
||
|
||
def test_composer_and_sidebar_fit_screen(self):
|
||
app, _agent, _root, refs = self.make_chat(rows=24)
|
||
lines = app.render_text().split("\n")
|
||
self.assertIn("╰", lines[-2])
|
||
self.assertIn("Send", lines[-3])
|
||
self.assertIn("F2", lines[-1])
|
||
|
||
def test_pill_row_centered(self):
|
||
spans = tai.tui_chat_pill_row("Today", 40, tai.TuiStyle(), tai.TuiStyle())
|
||
text = "".join(span.text for span in spans)
|
||
self.assertEqual(len(text), 40)
|
||
self.assertEqual(text.strip(), "Today")
|
||
self.assertTrue(text.startswith(" "))
|
||
|
||
def test_ticks_settle_on_reply(self):
|
||
app, agent, _root, refs = self.make_chat()
|
||
refs["do_send"]("hi bot")
|
||
shown = " ".join(line for line, _spans in refs["transcript"].lines)
|
||
self.assertIn(tai.TUI_CHAT_TICK_SEND, shown)
|
||
self.assertNotIn(tai.TUI_CHAT_TICK_READ, shown)
|
||
self.assertEqual(refs["head_status"].text, "typing...")
|
||
refs["run_pending"]()
|
||
shown = " ".join(line for line, _spans in refs["transcript"].lines)
|
||
self.assertIn(tai.TUI_CHAT_TICK_READ, shown)
|
||
self.assertNotIn(tai.TUI_CHAT_TICK_SEND, shown)
|
||
self.assertEqual(refs["head_status"].text, "online")
|
||
|
||
def test_date_pill_on_empty_history(self):
|
||
app = tai.TuiApp()
|
||
app.rows = 24
|
||
app.cols = 100
|
||
agent = FakeChatAgent()
|
||
agent.messages = [{"role": "system", "content": "sys"}]
|
||
_root, refs = tai.tui_build_chat(app, agent)
|
||
app.mount(_root)
|
||
refs["replay"]()
|
||
shown = " ".join(line for line, _spans in refs["transcript"].lines)
|
||
import datetime as datetime_mod
|
||
now = datetime_mod.datetime.now()
|
||
self.assertIn("%s %d" % (now.strftime("%B"), now.day), shown)
|
||
|
||
def test_chat_list_two_line_rows(self):
|
||
app, _agent, _root, refs = self.make_chat()
|
||
sidebar = refs["sidebar"]
|
||
rows = sidebar.content_rows()
|
||
texts = ["".join(span.text for span in row) for row in rows]
|
||
self.assertIn("main", texts[0])
|
||
self.assertIn("hi there", texts[1])
|
||
self.assertIn("coder", texts[2])
|
||
refs["box"].focus()
|
||
self.type(app, "draft-1")
|
||
sidebar.choose(1)
|
||
rows = sidebar.content_rows()
|
||
texts = ["".join(span.text for span in row) for row in rows]
|
||
self.assertIn("Draft:", texts[1])
|
||
|
||
def test_chat_list_click_maps_pairs(self):
|
||
_app, agent, _root, refs = self.make_chat()
|
||
sidebar = refs["sidebar"]
|
||
event = tai.TuiClick(panel=sidebar, x=sidebar.ax + 2, y=sidebar.ay + 2, pressed=True)
|
||
sidebar.on_click(event)
|
||
self.assertEqual(agent.bot, "coder")
|
||
event = tai.TuiClick(panel=sidebar, x=sidebar.ax + 2, y=sidebar.ay + 0, pressed=True)
|
||
sidebar.on_click(event)
|
||
self.assertEqual(agent.bot, "main")
|
||
|
||
def test_send_button_press_sends(self):
|
||
_app, _agent, root, refs = self.make_chat()
|
||
refs["box"].focus()
|
||
refs["box"].insert("via button")
|
||
root.find("send").on_press(None)
|
||
self.assertEqual(refs["pending"], "via button")
|
||
|
||
def test_telegram_roles_registered(self):
|
||
self.assertEqual(len(tai.TUI_THEME_ROLES), 27)
|
||
for role in ("bubble_in", "bubble_out", "bubble_fg", "meta_in", "meta_out",
|
||
"tick_read", "service_bg", "service_fg", "unread", "unread_muted"):
|
||
self.assertIn(role, tai.TUI_THEME_ROLES)
|
||
tai.tui_set_theme("dark")
|
||
self.assertEqual(tai.TuiStyle(bg="role:bubble_out").sgr(), "\x1b[0;48;2;67;56;202m")
|
||
|
||
def test_parser_flags(self):
|
||
args = tai.build_parser().parse_args(["--tui", "--theme", "nord"])
|
||
self.assertTrue(args.tui)
|
||
self.assertEqual(args.theme, "nord")
|
||
|
||
|
||
class EmbeddedStringsTests(unittest.TestCase):
|
||
def test_box_scripts_compile(self):
|
||
for name in ("BOX_STT", "BOX_TTS"):
|
||
compile(getattr(tai, name), "<%s>" % name, "exec")
|
||
|
||
def test_containerfile_shape(self):
|
||
self.assertTrue(tai.BOX_CONTAINERFILE.startswith("FROM "))
|
||
self.assertIn("\nRUN ", tai.BOX_CONTAINERFILE)
|
||
self.assertIn("\nCOPY ", tai.BOX_CONTAINERFILE)
|
||
|
||
def test_systemd_units_shape(self):
|
||
for name in ("SCHEDULER_UNIT", "TELEGRAM_UNIT"):
|
||
value = getattr(tai, name)
|
||
self.assertTrue(value.startswith("[Unit]"))
|
||
self.assertIn("\n[Service]\n", value)
|
||
|
||
|
||
class NoOpencodeFlagTests(unittest.TestCase):
|
||
def test_flag_disables_opencode(self):
|
||
args = tai.build_parser().parse_args(["--no-opencode", "hi"])
|
||
self.assertTrue(args.no_opencode)
|
||
self.assertFalse(tai.Config(args).opencode_enabled)
|
||
|
||
def test_default_keeps_opencode(self):
|
||
args = tai.build_parser().parse_args(["hi"])
|
||
with mock.patch.dict(os.environ, {"TAI_OPENCODE": ""}):
|
||
self.assertTrue(tai.Config(args).opencode_enabled)
|
||
|
||
|
||
class UpdateModelsTests(unittest.TestCase):
|
||
def test_replace_list_prunes_only_that_provider(self):
|
||
with tempfile.TemporaryDirectory() as home:
|
||
health = tai.ModelHealth(home)
|
||
health.store_list("opencode:free", ["old-free", "kept-free"])
|
||
health.record("opencode", "old-free", False)
|
||
health.record("opencode", "kept-free", True)
|
||
health.record("openrouter", "old-free", True)
|
||
added, removed = health.replace_list("opencode:free", "opencode", ["kept-free", "new-free"])
|
||
self.assertEqual(added, ["new-free"])
|
||
self.assertEqual(removed, ["old-free"])
|
||
self.assertEqual(health.entry("opencode", "old-free"), {})
|
||
self.assertTrue(health.entry("opencode", "kept-free"))
|
||
self.assertTrue(health.entry("openrouter", "old-free"))
|
||
self.assertEqual(tai.ModelHealth(home).cached_list("opencode:free"), ["kept-free", "new-free"])
|
||
|
||
def test_empty_discovery_keeps_old_list(self):
|
||
with tempfile.TemporaryDirectory() as home:
|
||
health = tai.ModelHealth(home)
|
||
health.store_list("opencode:free", ["a"])
|
||
self.assertEqual(health.replace_list("opencode:free", "opencode", []), ([], []))
|
||
self.assertEqual(health.cached_list("opencode:free"), ["a"])
|
||
|
||
def test_flag_parses(self):
|
||
self.assertTrue(tai.build_parser().parse_args(["--update-models"]).update_models)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|