Initial commit: single-file autonomous agent with sealed storage, sandbox, and telegram bot
This commit is contained in:
+324
@@ -0,0 +1,324 @@
|
||||
# retoor <retoor@molodetz.nl>
|
||||
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()
|
||||
Reference in New Issue
Block a user