Initial commit: single-file autonomous agent with sealed storage, sandbox, and telegram bot
This commit is contained in:
@@ -0,0 +1,198 @@
|
||||
retoor <retoor@molodetz.nl>
|
||||
|
||||
# tai
|
||||
|
||||
tai is a single-file autonomous AI agent written in Python. The entire
|
||||
implementation lives in `tai.py` (about 2000 lines) and uses only the Python
|
||||
standard library: no dependencies, no install step, no build system.
|
||||
|
||||
The agent runs as an interactive REPL or as a one-shot command. It reasons
|
||||
through an OpenAI-compatible backend, acts through fourteen tools, keeps
|
||||
per-profile memory, seals its stored state at rest, and can isolate shell and
|
||||
file operations inside a container sandbox.
|
||||
|
||||
## Requirements
|
||||
|
||||
- Python 3.10 or newer, no third-party packages.
|
||||
- Optional: `podman` or `docker` for the sandbox and Telegram voice notes.
|
||||
- Optional: `tmux` for terminal content capture.
|
||||
|
||||
## Quick start
|
||||
|
||||
./tai.py
|
||||
./tai.py --profile work
|
||||
./tai.py --yes
|
||||
./tai.py --version
|
||||
./tai.py what is 2+3, use the shell
|
||||
|
||||
Trailing arguments form a one-shot prompt: the agent answers once and exits
|
||||
with code 0. Without arguments, tai starts an interactive session.
|
||||
|
||||
## REPL commands
|
||||
|
||||
| Command | Effect |
|
||||
|-------------------|---------------------------------------------------|
|
||||
| `/profile [name]` | Show the current profile or switch to it |
|
||||
| `/profiles` | List all profiles, current marked with `*` |
|
||||
| `/env [target]` | Show or switch execution environment |
|
||||
| `/skills` | List loaded skill files |
|
||||
| `/fork <task>` | Spawn a background subagent, REPL stays free |
|
||||
| `/agents` | List background subagents |
|
||||
| `/agent <id>` | Show one subagent result |
|
||||
| `/agent clear` | Purge finished subagents |
|
||||
| `/compact` | Compress history into a summary |
|
||||
| `/clear` | Drop history, keep the system message |
|
||||
| `/help` | Show the command overview |
|
||||
| `/quit` | Exit |
|
||||
|
||||
Any other input is sent to the agent.
|
||||
|
||||
## Installation
|
||||
|
||||
./tai.py --install
|
||||
|
||||
This copies `tai.py` to `~/.local/bin/tai.py` as an executable and registers
|
||||
a guarded `command_not_found_handle` block in `~/.bashrc` (backed up once to
|
||||
`~/.bashrc.bak-tai`, idempotent), so unknown shell commands are answered by
|
||||
the agent instead of failing.
|
||||
|
||||
## Backends
|
||||
|
||||
The primary backend is `model.cloud.pravda.education`, an OpenAI-compatible
|
||||
gateway that needs no API key and selects a free model per request. If a
|
||||
request fails, tai retries it on `devplace.net/openai/v1`, which requires
|
||||
`DEVPLACE_API_KEY`. Both endpoints speak `/chat/completions`, including
|
||||
native tool calls and streaming.
|
||||
|
||||
## Orchestration
|
||||
|
||||
`/fork <task>` spawns a background subagent with its own context while the
|
||||
REPL stays free (the prompt shows a `+N` counter). `/agents` lists workers,
|
||||
`/agent <id>` shows a result, `/agent clear` purges finished ones.
|
||||
|
||||
The model itself orchestrates through the `fork` tool (task, timeout up to
|
||||
one hour, profile) and the `poll` tool (id, wait up to two minutes). Workers
|
||||
get 12 steps, a cooperative deadline, no session writes, and no interactive
|
||||
approval prompts. Nesting is capped at two levels. Timeouts and errors
|
||||
surface as statuses, never silently.
|
||||
|
||||
## Tools
|
||||
|
||||
| Tool | Purpose |
|
||||
|------------------------------|------------------------------------------------------|
|
||||
| `shell` | Run a shell command, output truncated |
|
||||
| `read_file` | Read a text file, large files truncated |
|
||||
| `write_file` | Write content to a file, creating parent directories |
|
||||
| `edit_file` | Replace one unique exact text match in a file |
|
||||
| `web_search` | Search the web, optionally images or page content |
|
||||
| `web_fetch` | Fetch a URL and return its text content |
|
||||
| `speak` | Synthesize speech, save MP3, play when possible |
|
||||
| `listen` | Record from the microphone and transcribe it |
|
||||
| `remember` | Merge knowledge into the profile system message |
|
||||
| `recall` | Search past session memory by keyword |
|
||||
| `load_skill` | Load a skill file by name |
|
||||
| `get_current_terminal_content` | Capture the current tmux pane with scrollback |
|
||||
| `fork` | Spawn a background subagent |
|
||||
| `poll` | Collect a background subagent result |
|
||||
|
||||
Web search runs on `rsearch.app.molodetz.nl`. Destructive shell commands ask
|
||||
for confirmation unless `--yes` is given; read-only commands run directly.
|
||||
|
||||
## Skills
|
||||
|
||||
Standard agent skill files (`SKILL.md` with `name` plus `description`
|
||||
frontmatter, optional `scripts/`, `references/`, `assets/`, per the Agent
|
||||
Skills open format) are discovered in `~/.tai/skills/*/` and
|
||||
`./.tai/skills/*/` (project wins on name collisions). Descriptions stay in
|
||||
context; the agent loads full instructions through `load_skill` only when
|
||||
needed. `/skills` lists what is available.
|
||||
|
||||
## Sandbox
|
||||
|
||||
`/env` shows the execution environment, `/env sandbox` switches shell and
|
||||
file tools into an isolated `tai-box` container (podman or docker, no mounts,
|
||||
no shared filesystem), `/env home` switches back. The image is built
|
||||
automatically on first use from an embedded Containerfile; every pip
|
||||
requirement (`faster-whisper`, `edge-tts`) lives inside the image while
|
||||
`tai.py` itself stays dependency-free. Sandbox commands need no approval
|
||||
because the container is disposable.
|
||||
|
||||
## Telegram
|
||||
|
||||
./tai.py --install-telegram
|
||||
./tai.py --uninstall-telegram
|
||||
|
||||
Install asks for the bot token up front, verifies it against `getMe`, stores
|
||||
it in `~/.tai/telegram.env` (0600), builds the sandbox container (used for
|
||||
voice transcription), and registers a `tai-telegram.service` systemd user
|
||||
unit with linger enabled (failures ignored). The bot long-polls, answers
|
||||
text, transcribes voice notes, and understands `/new`. It runs without
|
||||
`--yes`, so destructive shell commands are denied. Uninstall stops and
|
||||
removes the service and purges the container and image; the token file and
|
||||
data stay.
|
||||
|
||||
## Profiles and memory
|
||||
|
||||
Each profile owns a system message plus session history under
|
||||
`~/.tai/profiles`, stored with mode 0600.
|
||||
|
||||
/profile show current profile
|
||||
/profile [name] switch profile, creating it when missing
|
||||
/profiles list all profiles
|
||||
|
||||
The `remember` tool merges an instruction into the current profile system
|
||||
message through the model itself: it adds facts, updates behavior, or removes
|
||||
forgotten items while preserving the rest. It fires by default on new
|
||||
passwords and behavior changes. `recall` searches the per-profile episodic
|
||||
log in `~/.tai/memory.db` (SQLite). Context is budgeted at roughly 32k
|
||||
tokens with automatic compaction at 80 percent.
|
||||
|
||||
## Sealed storage
|
||||
|
||||
Storage is sealed by default with a built-in key, which stops casual reads
|
||||
but not a determined attacker, since the key ships in the source. Set
|
||||
`TAI_PASSPHRASE` for real protection: a home sealed with the default key is
|
||||
re-sealed to your passphrase automatically on first boot, with a notice.
|
||||
|
||||
The key comes from PBKDF2-SHA256 (200k rounds) over a random salt in
|
||||
`~/.tai/.seal`; values use a per-value nonce with HMAC-SHA256 encrypt-then-
|
||||
MAC. SQLite access goes through custom `tai_enc`/`tai_dec` functions
|
||||
registered with `create_function`, so inserts encrypt inline and recall
|
||||
decrypts before matching. Existing plaintext data is sealed automatically on
|
||||
first sealed start. A wrong passphrase refuses to start with exit code 2.
|
||||
Set the variable empty for plaintext storage.
|
||||
|
||||
This construction uses only the standard library and is honest file-theft
|
||||
protection, not audited cryptography; high-value secrets still belong in a
|
||||
dedicated manager.
|
||||
|
||||
## Voice
|
||||
|
||||
`speak` synthesizes free neural speech via the Microsoft Edge Read Aloud
|
||||
protocol, implemented with `socket` and `ssl` from the standard library. No
|
||||
key, no package. MP3 files land in `~/.tai/audio` and play when an OS player
|
||||
exists. `listen` records and transcribes when a recorder (`arecord`, `sox`,
|
||||
`ffmpeg`) and a transcriber (`whisper-cpp`, `whisper`) are installed, and
|
||||
reports exactly what is missing otherwise.
|
||||
|
||||
## Configuration
|
||||
|
||||
| Variable | Purpose | Default |
|
||||
|---------------------|--------------------------------------|--------------------------------------------|
|
||||
| `TAI_HOME` | State directory | `~/.tai` |
|
||||
| `TAI_MODEL` | Model id, ignored by primary gateway | `openrouter/free` |
|
||||
| `DEVPLACE_API_KEY` | Fallback backend credential | Empty, fallback disabled |
|
||||
| `TAI_VOICE` | Edge voice name | `en-US-EmmaMultilingualNeural` |
|
||||
| `TAI_PASSPHRASE` | Personal seal key | Unset: built-in key; empty: plaintext |
|
||||
| `TELEGRAM_BOT_TOKEN`| Bot token, or set during install | Empty |
|
||||
|
||||
## Testing
|
||||
|
||||
python3 test_seal.py
|
||||
python3 test_tai.py
|
||||
|
||||
## Layout
|
||||
|
||||
tai.py the entire agent
|
||||
test_seal.py seal regression tests
|
||||
test_tai.py skills, telegram, install, and parser tests
|
||||
+156
@@ -0,0 +1,156 @@
|
||||
# retoor <retoor@molodetz.nl>
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
import tai
|
||||
|
||||
|
||||
class FakeArgs:
|
||||
profile = "default"
|
||||
yes = True
|
||||
|
||||
|
||||
class SealTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
os.environ["TAI_HOME"] = self.tmp.name
|
||||
|
||||
def tearDown(self):
|
||||
self.tmp.cleanup()
|
||||
os.environ.pop("TAI_HOME", None)
|
||||
|
||||
def make_store(self, passphrase="test-secret-1"):
|
||||
config = tai.Config(FakeArgs())
|
||||
seal = tai.Seal(config.home, passphrase)
|
||||
return tai.Store(config, seal)
|
||||
|
||||
def test_round_trip(self):
|
||||
config = tai.Config(FakeArgs())
|
||||
seal = tai.Seal(config.home, "test-secret-1")
|
||||
locked = seal.lock("the wifi password is orchid-4421")
|
||||
self.assertTrue(locked.startswith("tai1$"))
|
||||
self.assertNotIn("orchid", locked)
|
||||
self.assertEqual(seal.unlock(locked), "the wifi password is orchid-4421")
|
||||
|
||||
def test_unique_nonce(self):
|
||||
config = tai.Config(FakeArgs())
|
||||
seal = tai.Seal(config.home, "test-secret-1")
|
||||
self.assertNotEqual(seal.lock("same"), seal.lock("same"))
|
||||
|
||||
def test_wrong_passphrase(self):
|
||||
config = tai.Config(FakeArgs())
|
||||
tai.Seal(config.home, "correct-horse")
|
||||
with self.assertRaises(tai.SealError):
|
||||
tai.Seal(config.home, "wrong-guess")
|
||||
|
||||
def test_tamper_detected(self):
|
||||
config = tai.Config(FakeArgs())
|
||||
seal = tai.Seal(config.home, "test-secret-1")
|
||||
locked = seal.lock("precious")
|
||||
tampered = locked[:-4] + ("A" if locked[-4] != "A" else "B") + locked[-3:]
|
||||
with self.assertRaises(tai.SealError):
|
||||
seal.unlock(tampered)
|
||||
|
||||
def test_sql_functions(self):
|
||||
store = self.make_store()
|
||||
row = store.db.execute("SELECT tai_enc('hello-sql')").fetchone()[0]
|
||||
self.assertTrue(row.startswith("tai1$"))
|
||||
plain = store.db.execute("SELECT tai_dec(?)", (row,)).fetchone()[0]
|
||||
self.assertEqual(plain, "hello-sql")
|
||||
store.close()
|
||||
|
||||
def test_sealed_events_recall(self):
|
||||
store = self.make_store()
|
||||
store.log_event("default", "user", "message", "my vault code is zebra-9")
|
||||
raw = store.db.execute("SELECT text FROM events").fetchone()[0]
|
||||
self.assertTrue(raw.startswith("tai1$"))
|
||||
self.assertNotIn("zebra", raw)
|
||||
rows = store.search_events("default", "zebra")
|
||||
self.assertEqual(len(rows), 1)
|
||||
self.assertIn("zebra-9", rows[0][3])
|
||||
store.close()
|
||||
|
||||
def test_sealed_profile_files(self):
|
||||
store = self.make_store()
|
||||
store.save_system("work", "system knows orchid-4421")
|
||||
with open(store.profile_path("work"), encoding="utf-8") as handle:
|
||||
self.assertTrue(handle.read().startswith("tai1$"))
|
||||
self.assertEqual(store.load_system("work"), "system knows orchid-4421")
|
||||
store.save_session("work", [{"role": "user", "content": "secret hello"}])
|
||||
with open(store.session_path("work"), encoding="utf-8") as handle:
|
||||
self.assertTrue(handle.read().startswith("tai1$"))
|
||||
self.assertEqual(store.load_session("work"), [{"role": "user", "content": "secret hello"}])
|
||||
store.close()
|
||||
|
||||
def test_migration_from_plaintext(self):
|
||||
plain = self.make_store(passphrase="")
|
||||
plain.log_event("default", "user", "message", "plain memory")
|
||||
plain.save_system("default", "plain system")
|
||||
plain.close()
|
||||
sealed = self.make_store(passphrase="new-secret-2")
|
||||
raw = sealed.db.execute("SELECT text FROM events").fetchone()[0]
|
||||
self.assertTrue(raw.startswith("tai1$"))
|
||||
with open(sealed.profile_path("default"), encoding="utf-8") as handle:
|
||||
self.assertTrue(handle.read().startswith("tai1$"))
|
||||
self.assertEqual(sealed.load_system("default"), "plain system")
|
||||
sealed.close()
|
||||
|
||||
def test_sealed_home_requires_passphrase(self):
|
||||
sealed = self.make_store()
|
||||
sealed.log_event("default", "user", "message", "locked memory")
|
||||
sealed.close()
|
||||
config = tai.Config(FakeArgs())
|
||||
with self.assertRaises(tai.SealError):
|
||||
tai.Store(config, tai.Seal(config.home, ""))
|
||||
|
||||
def test_passphrase_default(self):
|
||||
os.environ.pop("TAI_PASSPHRASE", None)
|
||||
value, default = tai.resolve_passphrase()
|
||||
self.assertEqual(value, tai.DEFAULT_PASSPHRASE)
|
||||
self.assertTrue(default)
|
||||
os.environ["TAI_PASSPHRASE"] = "personal-1"
|
||||
value, default = tai.resolve_passphrase()
|
||||
self.assertEqual(value, "personal-1")
|
||||
self.assertFalse(default)
|
||||
os.environ["TAI_PASSPHRASE"] = ""
|
||||
value, default = tai.resolve_passphrase()
|
||||
self.assertEqual(value, "")
|
||||
self.assertFalse(default)
|
||||
os.environ.pop("TAI_PASSPHRASE", None)
|
||||
|
||||
def test_rotate_default_to_personal(self):
|
||||
os.environ.pop("TAI_PASSPHRASE", None)
|
||||
config = tai.Config(FakeArgs())
|
||||
old = tai.Seal(config.home, tai.DEFAULT_PASSPHRASE)
|
||||
store = tai.Store(config, old)
|
||||
store.log_event("default", "user", "message", "rotate me")
|
||||
store.save_system("default", "rotate system")
|
||||
store.close()
|
||||
fresh = tai.rotate_seal(config, old, "personal-2")
|
||||
self.assertTrue(fresh.enabled)
|
||||
with self.assertRaises(tai.SealError):
|
||||
tai.Seal(config.home, tai.DEFAULT_PASSPHRASE)
|
||||
reopened = tai.Store(config, tai.Seal(config.home, "personal-2"))
|
||||
self.assertEqual(reopened.load_system("default"), "rotate system")
|
||||
self.assertIn("rotate me", reopened.search_events("default", "rotate")[0][3])
|
||||
raw = reopened.db.execute("SELECT text FROM events").fetchone()[0]
|
||||
self.assertTrue(raw.startswith("tai1$"))
|
||||
reopened.close()
|
||||
|
||||
def test_plaintext_home_unaffected(self):
|
||||
store = self.make_store(passphrase="")
|
||||
store.log_event("default", "user", "message", "open memory")
|
||||
self.assertEqual(store.search_events("default", "open")[0][3], "open memory")
|
||||
store.save_system("default", "open system")
|
||||
with open(store.profile_path("default"), encoding="utf-8") as handle:
|
||||
self.assertEqual(handle.read(), "open system")
|
||||
store.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+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