# retoor import json import os import sys import tempfile import time import unittest 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) self.assertEqual(tai.skill_catalog({}), "") class FakeApp: def __init__(self): self.skills = {} self.env = "home" 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 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_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.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_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 if __name__ == "__main__": unittest.main()