Files
tai/test_tai.py
T

6278 lines
290 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# 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 ![alt](pic.png)", 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()