Split monolithic tai.py into tai_pkg package, keep tai-orig.py as original
tai.py (8643 lines) is now a thin facade over tai_pkg/ (33 feature modules, one class or feature per file). tai-orig.py preserves the pre-split single-file version byte for byte (sha256 verified against snapshot); import tai and ./tai.py behavior unchanged. Equivalence proof, all run on this tree immediately pre-commit: differential battery tai-orig.py vs facade 3718 checks 0 failures (surface, pure fuzz, tool dispatch, store/audit, release ops, fake-HOME install roundtrip, CLI runs); test_tai.py 261 OK; test_seal.py 16 OK; --version/--help byte-identical; dir(tai) zero removals. Fix included: opencode wire-protocol schema echo leaked into replies; opencode_parse now scans with JSONDecoder.raw_decode capped at 200 scans (13 OpencodeParseTests). Test count 261 = 237 HEAD + 8 pre-existing uncommitted ToolArgumentsTests (now committed with this change, disclosed) + 13 OpencodeParseTests + 3 EmbeddedStringsTests. Patch-target-only test migration (24 hunks), zero assertion changes. tai_pkg clean-regenerated after mid-verification tamper was found and reverted.
This commit is contained in:
@@ -4,7 +4,7 @@
|
||||
|
||||
- Use plain `import` and `from` statements in Python files.
|
||||
- Use the `.py` extension for Python source and test files.
|
||||
- Keep source files at the repository top level.
|
||||
- Implementation lives in `tai_pkg/`, one feature per module, grouped into classes (instance classes for stateful clusters, static utility classes for stateless groups). `tai.py` is the thin facade: CLI entry plus the stable `import tai` surface re-exported from `tai_pkg`. Canonical patch targets are `tai_pkg.<module>.<Class>.<name>` (or the module attribute for module-level data); never patch through the `tai.` facade copies.
|
||||
- There are two kinds of test, in two different places, on purpose — see "Testing" below for the full split. In short: unit tests are top-level `test_*.py`; integration tests are `tests/<topic>.py` (no `test_` prefix) built on `tai_test.py`.
|
||||
|
||||
## Testing
|
||||
|
||||
@@ -2,9 +2,9 @@ retoor <retoor@molodetz.nl>
|
||||
|
||||
# tai
|
||||
|
||||
**One file. Zero dependencies. A real agent.**
|
||||
**Zero dependencies. A real agent.**
|
||||
|
||||
`tai` is a single-file autonomous AI agent written in Python — the entire implementation lives in `tai.py` (~7,600 lines) and imports **only the Python standard library**. No `pip install`, no `requirements.txt`, no build step, no lockfile to rot. It runs as an interactive REPL or a one-shot command, reasons through an OpenAI-compatible backend, acts through 38 tools, keeps per-profile memory in a self-healing SQLite vault, seals that vault at rest, ranks every search with real BM25 (no vector database, no embeddings bill, no GPU), chunks and indexes anything worth remembering for retrieval, and can wall shell/file work off inside a disposable container sandbox.
|
||||
`tai` is an autonomous AI agent written in Python — the implementation lives in the `tai_pkg/` package (one feature per module, grouped into classes) behind the thin `tai.py` command, and imports **only the Python standard library**. No `pip install`, no `requirements.txt`, no build step, no lockfile to rot. It runs as an interactive REPL or a one-shot command, reasons through an OpenAI-compatible backend, acts through 38 tools, keeps per-profile memory in a self-healing SQLite vault, seals that vault at rest, ranks every search with real BM25 (no vector database, no embeddings bill, no GPU), chunks and indexes anything worth remembering for retrieval, and can wall shell/file work off inside a disposable container sandbox.
|
||||
|
||||
Nothing here is decorative. Every mechanism below exists because something specific was expensive — context tokens, API calls, attacker surface, disk writes, or your attention — and got engineered out.
|
||||
|
||||
@@ -185,7 +185,7 @@ This is the part worth being loud about: **tai does retrieval-augmented generati
|
||||
|
||||
A chunk hit from `search` returns **the actual chunk text**, not just FTS's highlighted snippet — the whole point of chunk-level retrieval is to hand back something directly usable as context, and the parent record id rides along so `record_read` can pull more surrounding material on demand (hierarchical retrieval: chunk for precision, parent for context, one hop away, never both loaded by default).
|
||||
|
||||
**Why no vectors.** Embeddings mean a model call (cost, latency, an API dependency) or a local model (a real dependency, a real download, real RAM) for *every single thing written* — directly incompatible with the single-file, stdlib-only, zero-install rule this whole project is built on. BM25 over SQLite FTS5 costs nothing to build, nothing to maintain, and handles exactly what this vault actually needs to search well: command output, notes, transcripts, terminal history — content where the literal words (an error string, a flag name, a path) *are* the signal. This was an explicit, considered tradeoff, not a missing feature.
|
||||
**Why no vectors.** Embeddings mean a model call (cost, latency, an API dependency) or a local model (a real dependency, a real download, real RAM) for *every single thing written* — directly incompatible with the stdlib-only, zero-install rule this whole project is built on. BM25 over SQLite FTS5 costs nothing to build, nothing to maintain, and handles exactly what this vault actually needs to search well: command output, notes, transcripts, terminal history — content where the literal words (an error string, a flag name, a path) *are* the signal. This was an explicit, considered tradeoff, not a missing feature.
|
||||
|
||||
**Chunking is wired into the system, not bolted onto one feature.** `record_save`, the shell/subagent output spill path (`spill_output`), and terminal log export all chunk through the identical `chunk_lines()` call — "RAG" here means a property of the vault, not a special case for logs. It is deliberately *not* wired into `add_record`'s lower-level, highest-frequency caller — `upsert_file_record`, which fires on every single file write and edit for audit bookkeeping — because exploding the chunk table on every keystroke-adjacent save would be pure waste for content `read_file`/`grep` already serves perfectly well. And deleting a record deletes its chunks and their index rows with it (`Store.delete_record`); there is no orphan path to worry about.
|
||||
|
||||
@@ -345,6 +345,10 @@ Install asks for the bot token up front, verifies it against `getMe`, stores it
|
||||
|
||||
## Layout
|
||||
|
||||
tai.py the entire agent — tools, vault, search, skills, orchestration, sandbox, all of it
|
||||
tai.py CLI entry + stable import surface (thin facade, no logic of its own)
|
||||
tai_pkg/ the entire agent, one feature per module: agent loop, tools,
|
||||
vault, backends (muse/opencode/grok/claude/codex/gemini),
|
||||
search, skills, scheduler, telegram, plexer, sandbox,
|
||||
install/release, config, rendering — grouped into classes
|
||||
test_seal.py seal regression tests
|
||||
test_tai.py agent, tools, records, graph, search/RAG, skills, scheduler tests
|
||||
|
||||
Executable
+8643
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,7 @@
|
||||
"""tai_pkg: implementation package for the tai agent.
|
||||
|
||||
The stable import surface is the root tai.py facade; import that, or
|
||||
import tai_pkg.<module> directly for canonical patch targets.
|
||||
"""
|
||||
|
||||
__all__ = []
|
||||
@@ -0,0 +1,451 @@
|
||||
"""Main agent loop, context accounting, and turn resume."""
|
||||
|
||||
import collections
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import uuid
|
||||
|
||||
from tai_pkg import render
|
||||
from tai_pkg.chat import ChatClient
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.errors import BackendError
|
||||
from tai_pkg.errors import Denied
|
||||
from tai_pkg.render import Ansi
|
||||
from tai_pkg.render import Render
|
||||
from tai_pkg.skills import Skills
|
||||
from tai_pkg.text import Text
|
||||
from tai_pkg.tools import Tools
|
||||
from tai_pkg.toolschema import ToolRegistry
|
||||
|
||||
__all__ = ["Agent"]
|
||||
|
||||
|
||||
class Agent:
|
||||
def __init__(self, config, store, persist=True, quiet=False, depth=0, auto=False, yolo=False):
|
||||
self.config = config
|
||||
self.auto = auto
|
||||
self.yolo = yolo
|
||||
self.store = store
|
||||
self.chat = ChatClient(config)
|
||||
if not persist or depth > 0:
|
||||
self.chat.fixed_key = "worker/%s" % uuid.uuid4().hex
|
||||
self.tools = Tools(self)
|
||||
self.profile = ""
|
||||
self.bot = "main"
|
||||
self.system_message = ""
|
||||
self.messages = []
|
||||
self.skills = Skills.discover_skills(config.home, os.getcwd())
|
||||
self.env = "home"
|
||||
self.persist = persist
|
||||
self.quiet = quiet
|
||||
self.depth = depth
|
||||
self.deadline = None
|
||||
self.timed_out = False
|
||||
self.runner_override = None
|
||||
self.user_stepped_in = False
|
||||
self._turn_open = False
|
||||
self._turn_goal = ""
|
||||
self._turn_steps = 0
|
||||
self._turn_tools = 0
|
||||
self.switch_profile(config.profile, silent=True)
|
||||
|
||||
def apply_system(self):
|
||||
content = self.system_message + Skills.skill_catalog(self.skills) + ToolRegistry.tool_catalog()
|
||||
if self.auto:
|
||||
content += "\n\n" + Consts.AUTONOMOUS_NOTE
|
||||
self.messages[0] = {"role": "system", "content": content}
|
||||
|
||||
def sync_chat_scope(self):
|
||||
self.chat.scope = (self.profile, self.bot)
|
||||
|
||||
def reset_history(self):
|
||||
self.messages = [{"role": "system", "content": self.system_message}]
|
||||
self.apply_system()
|
||||
|
||||
def switch_profile(self, name, silent=False):
|
||||
if self.messages and self.persist:
|
||||
self.store.save_session(self.profile, self.messages, self.bot)
|
||||
self.profile = name
|
||||
self.bot = "main"
|
||||
self.store.profile = name
|
||||
self.tools.read_files.clear()
|
||||
self.tools.secret_grants.clear()
|
||||
loaded = self.store.load_system(name)
|
||||
if loaded is None:
|
||||
loaded = Consts.DEFAULT_SYSTEM
|
||||
self.store.save_system(name, loaded)
|
||||
if Consts.LEGACY_PASSWORD_NOTE in loaded:
|
||||
loaded = loaded.replace(Consts.LEGACY_PASSWORD_NOTE, Consts.VAULT_MEMORY_NOTE)
|
||||
self.store.save_system(name, loaded)
|
||||
self.system_message = loaded
|
||||
restored = self.store.load_session(name) if self.persist else []
|
||||
self.messages = [{"role": "system", "content": self.system_message}] + restored
|
||||
self.apply_system()
|
||||
self.sync_chat_scope()
|
||||
if not silent:
|
||||
print(Render.paint("profile: %s (%d chars knowledge, %d restored messages)" % (name, len(self.system_message), len(restored)), Ansi.GREEN))
|
||||
|
||||
def update_system(self, text):
|
||||
self.system_message = text
|
||||
if self.bot == "main":
|
||||
self.store.save_system(self.profile, text)
|
||||
else:
|
||||
self.store.save_bot_system(self.profile, self.bot, text)
|
||||
self.apply_system()
|
||||
self.store.log_event(self.profile, "system", "remember", "system message updated (%d chars)" % len(text))
|
||||
|
||||
def switch_bot(self, name, silent=False):
|
||||
target = self.store.resolve_bot(self.profile, name)
|
||||
if target is None:
|
||||
return "unknown bot @%s (available: %s)" % (name, ", ".join(item["name"] for item in self.store.list_bots(self.profile)))
|
||||
if self.messages and self.persist:
|
||||
self.store.save_session(self.profile, self.messages, self.bot)
|
||||
self.bot = target
|
||||
loaded = self.store.load_bot_system(self.profile, target)
|
||||
if loaded is None:
|
||||
loaded = Consts.DEFAULT_SYSTEM
|
||||
if target == "main":
|
||||
self.store.save_system(self.profile, loaded)
|
||||
else:
|
||||
self.store.save_bot_system(self.profile, target, loaded)
|
||||
self.system_message = loaded
|
||||
restored = self.store.load_session(self.profile, target) if self.persist else []
|
||||
self.messages = [{"role": "system", "content": self.system_message}] + restored
|
||||
self.apply_system()
|
||||
self.sync_chat_scope()
|
||||
if not silent:
|
||||
print(Render.paint("bot: %s (%d restored messages)" % (target, len(restored)), Ansi.GREEN))
|
||||
return "switched to bot '%s'" % target
|
||||
|
||||
def scrub_secret_from_history(self):
|
||||
for item in self.messages:
|
||||
content = item.get("content")
|
||||
if isinstance(content, str) and content:
|
||||
item["content"] = self.store.redact(content, self.profile)
|
||||
for recorded in item.get("tool_calls") or []:
|
||||
args = (recorded.get("function") or {}).get("arguments")
|
||||
if isinstance(args, str) and args:
|
||||
recorded["function"]["arguments"] = self.store.redact(args, self.profile)
|
||||
self.chat.reset_muse()
|
||||
|
||||
def ask_approval(self, command, guidance=True):
|
||||
if self.yolo or self.auto:
|
||||
return True
|
||||
if not self.persist:
|
||||
return self.config.auto_approve
|
||||
if not sys.stdin.isatty():
|
||||
return False
|
||||
print(Render.paint("run: %s" % command, Ansi.YELLOW))
|
||||
try:
|
||||
answer = input(Render.paint("allow? [y]once [Y]always [n]o ", Ansi.YELLOW)).strip()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
return False
|
||||
self.user_stepped_in = True
|
||||
if answer == "Y" or answer.lower() in ("yolo", "always"):
|
||||
self.yolo = True
|
||||
self.config.auto_approve = True
|
||||
print(Render.paint("yolo mode on, approvals off for this session", Ansi.YELLOW))
|
||||
return True
|
||||
if answer.lower() in ("y", "yes"):
|
||||
return True
|
||||
if not guidance or self.auto:
|
||||
return False
|
||||
try:
|
||||
typed = input(Render.paint("denied. What should I do instead? (empty aborts this turn) ", Ansi.YELLOW)).strip()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
typed = ""
|
||||
raise Denied(typed)
|
||||
|
||||
def show_call(self, call):
|
||||
if self.quiet:
|
||||
return
|
||||
print(Render.paint("┌─ %s" % call["name"], Ansi.MAGENTA))
|
||||
preview_args = Text.tool_args_text(call["arguments"])
|
||||
if call["name"] == "store_secret":
|
||||
try:
|
||||
masked = json.loads(preview_args or "{}")
|
||||
masked["value"] = "[hidden]"
|
||||
preview_args = json.dumps(masked)
|
||||
except ValueError:
|
||||
preview_args = "[hidden]"
|
||||
preview = self.store.redact(preview_args, self.profile)[:400].replace("\n", " ")
|
||||
print(Render.paint("│ %s" % preview, Ansi.DIM))
|
||||
|
||||
def show_result(self, result, elapsed_ms):
|
||||
if self.quiet:
|
||||
return
|
||||
self.show_diff()
|
||||
print(Render.paint("└─ %d chars · %dms" % (len(result), elapsed_ms), Ansi.DIM))
|
||||
|
||||
def show_diff(self):
|
||||
preview = getattr(self.tools, "diff_preview", None)
|
||||
if preview is None:
|
||||
return
|
||||
self.tools.diff_preview = None
|
||||
rendered = Render.render_diff(preview["path"], preview["old"], preview["new"], colors=Render.color_enabled())
|
||||
if rendered:
|
||||
print(rendered)
|
||||
|
||||
def checkpoint(self):
|
||||
if not self.persist:
|
||||
return
|
||||
self.store.save_session(self.profile, self.messages, self.bot)
|
||||
if self._turn_open:
|
||||
self.store.save_turn_state(self.profile, self.bot, {"open": True, "goal": self._turn_goal, "steps": self._turn_steps, "tools": self._turn_tools, "updated": Text.now_iso()})
|
||||
|
||||
def close_turn(self):
|
||||
if not self.persist:
|
||||
self._turn_open = False
|
||||
return
|
||||
self.store.save_session(self.profile, self.messages, self.bot)
|
||||
self._turn_open = False
|
||||
if self._turn_goal:
|
||||
self.store.save_turn_state(self.profile, self.bot, {"open": False, "goal": self._turn_goal, "steps": self._turn_steps, "tools": self._turn_tools, "updated": Text.now_iso()})
|
||||
|
||||
def run_turn(self, text, capture=False, max_steps=Consts.MAX_STEPS, _routed=False):
|
||||
if not _routed:
|
||||
mentioned = Consts.BOT_MENTION_RE.match(text or "")
|
||||
if mentioned:
|
||||
target = self.store.resolve_bot(self.profile, mentioned.group(1))
|
||||
if target is None:
|
||||
return "unknown bot @%s (available: %s)" % (mentioned.group(1), ", ".join(item["name"] for item in self.store.list_bots(self.profile)))
|
||||
if target != self.bot:
|
||||
return self.run_mention(target, mentioned.group(2), text, capture, max_steps)
|
||||
text = mentioned.group(2)
|
||||
self.messages.append({"role": "user", "content": text})
|
||||
self.store.log_event(self.profile, "user", "message", self.store.redact(text, self.profile))
|
||||
self._turn_open = True
|
||||
self._turn_goal = text[:500]
|
||||
self._turn_steps = 0
|
||||
self._turn_tools = 0
|
||||
self.checkpoint()
|
||||
started = time.time()
|
||||
backends = []
|
||||
pieces = []
|
||||
last_text = ""
|
||||
sink = pieces.append
|
||||
loud = not capture and not self.quiet
|
||||
self.tools.live = loud
|
||||
seen_actions = collections.deque(maxlen=400)
|
||||
seen_results = collections.deque(maxlen=400)
|
||||
stall = 0
|
||||
last_pair = None
|
||||
run_len = 0
|
||||
nudged = False
|
||||
for _step in range(Consts.TOTAL_STEP_CAP):
|
||||
if self.deadline is not None and time.time() > self.deadline:
|
||||
self.timed_out = True
|
||||
last_text = (last_text + "\n[time limit reached]").strip()
|
||||
break
|
||||
payload = ToolRegistry.select_tools(Agent.conversation_text(self.messages))
|
||||
tokens = Agent.conversation_tokens(self.messages, payload)
|
||||
if tokens > Consts.CONTEXT_CAP * Consts.COMPACT_RATIO:
|
||||
if loud:
|
||||
print(Render.paint("compacting context (%d tokens)..." % tokens, Ansi.DIM))
|
||||
self.messages, _changed = Agent.compact_messages(self.messages, self.chat)
|
||||
self.apply_system()
|
||||
spinner = None
|
||||
if loud:
|
||||
print(Render.paint("tai", Ansi.BOLD, Ansi.CYAN))
|
||||
spinner = render.Spinner("thinking")
|
||||
spinner.start()
|
||||
try:
|
||||
reply = self.chat.complete(self.messages, payload, stream_sink=sink, deadline=self.deadline)
|
||||
except BackendError as exc:
|
||||
if loud:
|
||||
print(Render.paint("backend error: %s" % exc, Ansi.RED))
|
||||
self.messages.pop()
|
||||
self.checkpoint()
|
||||
return last_text or "backend error: %s" % exc
|
||||
finally:
|
||||
if spinner is not None:
|
||||
spinner.stop()
|
||||
reply["content"] = self.store.redact(reply["content"] or "", self.profile)
|
||||
if loud and reply["content"].strip():
|
||||
print(Render.render_markdown(reply["content"]))
|
||||
if reply["backend"] not in backends:
|
||||
backends.append(reply["backend"])
|
||||
if loud and (len(backends) > 1 or reply["backend"] != Consts.MUSE_LABEL):
|
||||
print(Render.paint("(via %s)" % reply["backend"], Ansi.DIM))
|
||||
if loud and reply["reasoning"] and not reply["content"] and not reply["tool_calls"]:
|
||||
print(Render.paint(reply["reasoning"][:2000], Ansi.GRAY))
|
||||
history_calls = [{"id": call["id"], "type": "function", "function": {"name": call["name"], "arguments": call["arguments"]}} for call in reply["tool_calls"]]
|
||||
self.messages.append({"role": "assistant", "content": reply["content"] or None, "tool_calls": history_calls or None})
|
||||
if not reply["tool_calls"]:
|
||||
if reply["content"]:
|
||||
last_text = reply["content"]
|
||||
elif reply["reasoning"]:
|
||||
last_text = reply["reasoning"]
|
||||
if last_text:
|
||||
self.store.log_event(self.profile, "assistant", "message", last_text)
|
||||
break
|
||||
aborted = False
|
||||
step_pairs = []
|
||||
for pos, call in enumerate(reply["tool_calls"]):
|
||||
self.show_call(call)
|
||||
call_started = time.time()
|
||||
try:
|
||||
result = self.tools.dispatch(call["name"], call["arguments"])
|
||||
except Denied as denied:
|
||||
result = "denied by user"
|
||||
self.show_result(result, (time.time() - call_started) * 1000)
|
||||
self.messages.append({"role": "tool", "tool_call_id": call["id"], "content": result})
|
||||
self.store.log_event(self.profile, "tool", call["name"], result)
|
||||
step_pairs.append((Text.action_signature(call["name"], call["arguments"]), Text.result_digest(result)))
|
||||
for skipped in reply["tool_calls"][pos + 1:]:
|
||||
self.messages.append({"role": "tool", "tool_call_id": skipped["id"], "content": "skipped: stopped after denial"})
|
||||
if denied.guidance:
|
||||
self.messages.append({"role": "user", "content": denied.guidance})
|
||||
self.store.log_event(self.profile, "user", "guidance", self.store.redact(denied.guidance, self.profile))
|
||||
else:
|
||||
aborted = True
|
||||
last_text = "stopped by user"
|
||||
break
|
||||
self.show_result(result, (time.time() - call_started) * 1000)
|
||||
self.messages.append({"role": "tool", "tool_call_id": call["id"], "content": result})
|
||||
self.store.log_event(self.profile, "tool", call["name"], result)
|
||||
step_pairs.append((Text.action_signature(call["name"], call["arguments"]), Text.result_digest(result)))
|
||||
if call["name"] == "store_secret":
|
||||
self.scrub_secret_from_history()
|
||||
if aborted:
|
||||
break
|
||||
productive = self.user_stepped_in
|
||||
self.user_stepped_in = False
|
||||
for pair in step_pairs:
|
||||
if pair[0] not in seen_actions or pair[1] not in seen_results:
|
||||
productive = True
|
||||
seen_actions.append(pair[0])
|
||||
seen_results.append(pair[1])
|
||||
if pair == last_pair:
|
||||
run_len += 1
|
||||
else:
|
||||
last_pair = pair
|
||||
run_len = 1
|
||||
nudged = False
|
||||
self._turn_steps += 1
|
||||
self._turn_tools += len(step_pairs)
|
||||
self.checkpoint()
|
||||
if productive:
|
||||
stall = 0
|
||||
else:
|
||||
stall += 1
|
||||
if run_len == Consts.LOOP_NUDGE_AT and not nudged:
|
||||
nudged = True
|
||||
nudge = "loop warning: you called %s %d times in a row with identical results. Change approach, try different arguments, or stop instead of repeating it." % (last_pair[0][:200], run_len)
|
||||
self.messages.append({"role": "user", "content": nudge})
|
||||
self.store.log_event(self.profile, "user", "guidance", nudge)
|
||||
if loud:
|
||||
print(Render.paint(nudge, Ansi.YELLOW))
|
||||
if run_len >= Consts.LOOP_STOP_AT:
|
||||
last_text = "stuck in a loop: %s repeated %d times with identical results" % (last_pair[0][:200], run_len)
|
||||
if loud:
|
||||
print(Render.paint(last_text, Ansi.YELLOW))
|
||||
self.store.log_event(self.profile, "system", "stall", last_text)
|
||||
break
|
||||
if stall >= max_steps:
|
||||
last_text = "no progress for %d steps: actions and results keep repeating. Say 'continue' to resume with fresh direction." % max_steps
|
||||
if loud:
|
||||
print(Render.paint(last_text, Ansi.YELLOW))
|
||||
self.store.log_event(self.profile, "system", "stall", last_text)
|
||||
break
|
||||
else:
|
||||
if loud:
|
||||
print(Render.paint("step budget exhausted (%d total steps)" % Consts.TOTAL_STEP_CAP, Ansi.YELLOW))
|
||||
self.close_turn()
|
||||
total = time.time() - started
|
||||
if loud:
|
||||
print(Render.paint("tokens≈%d · %s · %.1fs" % (Agent.conversation_tokens(self.messages), "+".join(backends), total), Ansi.DIM))
|
||||
return last_text
|
||||
|
||||
def run_mention(self, target, rest, original, capture=False, max_steps=Consts.MAX_STEPS):
|
||||
if self.persist:
|
||||
self.store.save_session(self.profile, self.messages, self.bot)
|
||||
other = self.store.load_session(self.profile, target)
|
||||
else:
|
||||
other = []
|
||||
system = self.store.load_bot_system(self.profile, target) or Consts.DEFAULT_SYSTEM
|
||||
saved_messages, saved_system, saved_bot = self.messages, self.system_message, self.bot
|
||||
self.messages = [{"role": "system", "content": system}] + other
|
||||
self.system_message = system
|
||||
self.bot = target
|
||||
self.apply_system()
|
||||
self.sync_chat_scope()
|
||||
try:
|
||||
answer = self.run_turn(rest, capture=capture, max_steps=max_steps, _routed=True)
|
||||
finally:
|
||||
if self.persist:
|
||||
self.store.save_session(self.profile, self.messages, target)
|
||||
self.messages, self.system_message, self.bot = saved_messages, saved_system, saved_bot
|
||||
self.apply_system()
|
||||
self.sync_chat_scope()
|
||||
self.store.log_event(self.profile, "user", "message", self.store.redact(original, self.profile))
|
||||
self.messages.append({"role": "user", "content": original})
|
||||
self.messages.append({"role": "assistant", "content": answer, "reasoning": "", "tool_calls": [], "backend": "mention"})
|
||||
if self.persist:
|
||||
self.store.save_session(self.profile, self.messages, self.bot)
|
||||
return answer
|
||||
|
||||
@staticmethod
|
||||
def conversation_text(messages):
|
||||
parts = []
|
||||
for item in messages:
|
||||
content = item.get("content")
|
||||
if item.get("role") in ("user", "assistant", "tool") and isinstance(content, str) and content:
|
||||
parts.append(content)
|
||||
return "\n".join(parts)[-6000:]
|
||||
|
||||
@staticmethod
|
||||
def conversation_tokens(messages, tools=None):
|
||||
total = 0
|
||||
for item in messages:
|
||||
total += Text.estimate_tokens(item.get("content"))
|
||||
for call in item.get("tool_calls") or []:
|
||||
total += Text.estimate_tokens(json.dumps(call))
|
||||
if tools:
|
||||
total += Text.estimate_tokens(json.dumps(tools))
|
||||
return total
|
||||
|
||||
@staticmethod
|
||||
def compact_messages(messages, chat, keep=Consts.KEEP_TURNS):
|
||||
if len(messages) <= keep + 1:
|
||||
return messages, False
|
||||
system = messages[0] if messages[0].get("role") == "system" else {"role": "system", "content": Consts.DEFAULT_SYSTEM}
|
||||
recent = messages[-keep:]
|
||||
while recent and recent[0].get("role") != "user":
|
||||
recent = recent[1:]
|
||||
if not recent:
|
||||
recent = messages[-keep:]
|
||||
middle = messages[1:len(messages) - len(recent)]
|
||||
if not middle:
|
||||
return messages, False
|
||||
digest = []
|
||||
for item in middle:
|
||||
role = item.get("role")
|
||||
if role == "assistant" and item.get("tool_calls"):
|
||||
names = ", ".join(call.get("function", {}).get("name", "?") for call in item["tool_calls"])
|
||||
digest.append("assistant tool calls: " + names)
|
||||
digest.append("%s: %s" % (role, (item.get("content") or "")[:1500]))
|
||||
request = [
|
||||
{"role": "system", "content": Consts.COMPACT_SYSTEM},
|
||||
{"role": "user", "content": "Compress this history:\n" + "\n".join(digest)},
|
||||
]
|
||||
summary = chat.complete(request, ephemeral=True)["content"].strip()
|
||||
rebuilt = [
|
||||
system,
|
||||
{"role": "user", "content": "Previous session summary:\n" + summary},
|
||||
{"role": "assistant", "content": "Understood. Continuing with summarized context."},
|
||||
] + recent
|
||||
return rebuilt, True
|
||||
|
||||
@staticmethod
|
||||
def resume_turn(agent):
|
||||
state = agent.store.load_turn_state(agent.profile, agent.bot)
|
||||
if not state.get("open"):
|
||||
print("nothing to resume: last turn for %s/%s completed normally" % (agent.profile, agent.bot))
|
||||
return ""
|
||||
goal = state.get("goal") or ""
|
||||
print("resuming interrupted turn for %s/%s: %d steps, %d tool calls, goal: %s" % (agent.profile, agent.bot, state.get("steps", 0), state.get("tools", 0), goal[:200]))
|
||||
message = "Continue where you left off%s. Do not repeat steps already completed above; carry on with the next step." % (": the interrupted request was: %s" % goal if goal else "")
|
||||
return agent.run_turn(message)
|
||||
@@ -0,0 +1,98 @@
|
||||
"""Background subagent registry and lifecycle."""
|
||||
|
||||
import threading
|
||||
import time
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.store import Store
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["AgentRegistry", "REGISTRY", "default_runner"]
|
||||
|
||||
|
||||
class AgentRegistry:
|
||||
"""Background subagent registry."""
|
||||
|
||||
def __init__(self):
|
||||
self.agents = {}
|
||||
self.lock = threading.Lock()
|
||||
self.next_id = [1]
|
||||
|
||||
def spawn_agent(self, task, profile, timeout, config, seal, depth=0, runner=None):
|
||||
with self.lock:
|
||||
agent_id = self.next_id[0]
|
||||
self.next_id[0] += 1
|
||||
record = {"id": agent_id, "task": task, "profile": profile, "status": "running", "result": "", "started": time.time(), "ended": None, "thread": None, "spilled": None}
|
||||
self.agents[agent_id] = record
|
||||
|
||||
def target():
|
||||
try:
|
||||
if runner is None:
|
||||
result = default_runner(config, seal, profile, depth, task, timeout)
|
||||
else:
|
||||
result = runner(task, profile, timeout)
|
||||
status = "timeout" if "[time limit reached]" in result else "done"
|
||||
except Exception as exc:
|
||||
result = "worker failed: " + Text.short_error(exc)
|
||||
status = "error"
|
||||
with self.lock:
|
||||
record["result"] = result
|
||||
record["status"] = status
|
||||
record["ended"] = time.time()
|
||||
|
||||
worker_thread = threading.Thread(target=target, daemon=True)
|
||||
with self.lock:
|
||||
record["thread"] = worker_thread
|
||||
worker_thread.start()
|
||||
return agent_id
|
||||
|
||||
def poll_agent(self, agent_id, wait=0, profile=None):
|
||||
with self.lock:
|
||||
record = self.agents.get(agent_id)
|
||||
if profile is not None and record is not None and record.get("profile") != profile:
|
||||
record = None
|
||||
worker_thread = record["thread"] if record else None
|
||||
if record is None:
|
||||
return "missing", "no agent %d" % agent_id
|
||||
if worker_thread is not None:
|
||||
worker_thread.join(timeout=max(0, wait))
|
||||
with self.lock:
|
||||
status = record["status"]
|
||||
result = record["result"]
|
||||
elapsed = (record["ended"] or time.time()) - record["started"]
|
||||
if status == "running":
|
||||
return status, "agent %d still running (%ds elapsed)" % (agent_id, elapsed)
|
||||
if status == "error":
|
||||
return status, result
|
||||
return status, Text.truncate(result, 6000, 2000) or "(empty result)"
|
||||
|
||||
def list_agents(self, profile=None):
|
||||
with self.lock:
|
||||
return [{"id": record["id"], "task": record["task"], "status": record["status"], "elapsed": int((record["ended"] or time.time()) - record["started"])} for record in self.agents.values() if profile is None or record.get("profile") == profile]
|
||||
|
||||
def clear_agents(self, profile=None):
|
||||
with self.lock:
|
||||
finished = [key for key, record in self.agents.items() if record["status"] != "running" and (profile is None or record.get("profile") == profile)]
|
||||
for key in finished:
|
||||
del self.agents[key]
|
||||
return len(finished)
|
||||
|
||||
def running_agents(self):
|
||||
with self.lock:
|
||||
return sum(1 for record in self.agents.values() if record["status"] == "running")
|
||||
|
||||
|
||||
REGISTRY = AgentRegistry()
|
||||
|
||||
|
||||
def default_runner(config, seal, profile, depth, task, timeout, max_steps=Consts.WORKER_STEPS):
|
||||
from tai_pkg.agent import Agent # deferred: breaks tools/agents/agent cycle
|
||||
store = Store(config, seal)
|
||||
try:
|
||||
worker = Agent(config, store, persist=False, quiet=True, depth=depth)
|
||||
if profile != worker.profile:
|
||||
worker.switch_profile(profile, silent=True)
|
||||
worker.deadline = time.time() + timeout
|
||||
return worker.run_turn(task, capture=True, max_steps=max_steps)
|
||||
finally:
|
||||
store.close()
|
||||
+224
@@ -0,0 +1,224 @@
|
||||
"""Startup, CLI parsing, banner, REPL, and main entry."""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import threading
|
||||
|
||||
from tai_pkg.agent import Agent
|
||||
from tai_pkg.agents import REGISTRY
|
||||
from tai_pkg.commands import Commands
|
||||
from tai_pkg.compat import readline
|
||||
from tai_pkg.config import Config
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.errors import BackendError
|
||||
from tai_pkg.install import Install
|
||||
from tai_pkg.models import ModelScoring
|
||||
from tai_pkg.plexer import NativePlexerBackend
|
||||
from tai_pkg.release import Release
|
||||
from tai_pkg.render import Ansi
|
||||
from tai_pkg.render import Render
|
||||
from tai_pkg.sandbox import Sandbox
|
||||
from tai_pkg.scheduler import Scheduler
|
||||
from tai_pkg.seal import Seal
|
||||
from tai_pkg.seal import SealError
|
||||
from tai_pkg.seal import SealManager
|
||||
from tai_pkg.store import Store
|
||||
from tai_pkg.telegram import TelegramBot
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["App"]
|
||||
|
||||
|
||||
class App:
|
||||
"""App."""
|
||||
|
||||
@staticmethod
|
||||
def boot(args):
|
||||
args.yes = args.yes or args.yolo or args.auto
|
||||
config = Config(args)
|
||||
passphrase, default_key = SealManager.resolve_passphrase()
|
||||
try:
|
||||
seal = Seal(config.home, passphrase)
|
||||
except SealError as exc:
|
||||
if default_key:
|
||||
raise
|
||||
try:
|
||||
old = Seal(config.home, Consts.DEFAULT_PASSPHRASE)
|
||||
except SealError:
|
||||
raise exc
|
||||
try:
|
||||
seal = SealManager.rotate_seal(config, old, passphrase)
|
||||
except (OSError, sqlite3.Error) as fail:
|
||||
raise SealError(str(fail))
|
||||
print("seal upgraded from default key to TAI_PASSPHRASE")
|
||||
seal.default_key = default_key and seal.enabled
|
||||
store = Store(config, seal)
|
||||
if store.schema_notes:
|
||||
print("vault schema upgraded: %s" % ", ".join(store.schema_notes))
|
||||
try:
|
||||
backed_up = Release.ensure_self_backup(config, store)
|
||||
except (OSError, sqlite3.Error) as exc:
|
||||
print("warning: self-backup failed: %s" % Text.short_error(exc))
|
||||
backed_up = None
|
||||
if backed_up is not None:
|
||||
print("self-backup: %s" % os.path.basename(backed_up))
|
||||
return Agent(config, store, yolo=args.yolo or args.auto, auto=args.auto)
|
||||
|
||||
@staticmethod
|
||||
def build_parser():
|
||||
parser = argparse.ArgumentParser(prog="tai", description="single-file autonomous agent, standard library only")
|
||||
parser.add_argument("--profile", default=None)
|
||||
parser.add_argument("--yes", action="store_true")
|
||||
parser.add_argument("--yolo", action="store_true")
|
||||
parser.add_argument("--auto", action="store_true")
|
||||
parser.add_argument("--continue", dest="resume", action="store_true")
|
||||
parser.add_argument("--env", default=None)
|
||||
parser.add_argument("--version", action="store_true")
|
||||
parser.add_argument("--telegram", action="store_true")
|
||||
parser.add_argument("--scheduler", action="store_true")
|
||||
parser.add_argument("--install", action="store_true")
|
||||
parser.add_argument("--install-telegram", action="store_true")
|
||||
parser.add_argument("--uninstall-telegram", action="store_true")
|
||||
parser.add_argument("--plexer-wrap", action="store_true", help="internal: invoked by the bashrc plexer hook, not meant to be run by hand")
|
||||
parser.add_argument("prompt", nargs="*")
|
||||
return parser
|
||||
|
||||
@staticmethod
|
||||
def banner(agent):
|
||||
print(Render.paint("tai %s · profile %s" % (Consts.VERSION, agent.profile), Ansi.BOLD, Ansi.CYAN))
|
||||
if agent.yolo:
|
||||
print(Render.paint("YOLO MODE: all approvals off", Ansi.BOLD, Ansi.RED))
|
||||
if agent.auto:
|
||||
print(Render.paint("AUTO MODE: never asks, researches instead", Ansi.BOLD, Ansi.YELLOW))
|
||||
seal_state = "off"
|
||||
if agent.store.seal.enabled:
|
||||
seal_state = "on (default key)" if agent.store.seal.default_key else "on (TAI_PASSPHRASE)"
|
||||
print(" seal: %s" % seal_state)
|
||||
for label, _chat_url, models_url, key, _model in agent.chat.backends():
|
||||
print(" %s: %s" % (label, ModelScoring.probe_backend(models_url, key)))
|
||||
if agent.store.load_turn_state(agent.profile, agent.bot).get("open"):
|
||||
print(Render.paint("interrupted turn available, restart with --continue to pick it up", Ansi.YELLOW))
|
||||
print(Render.paint("type /help for commands", Ansi.DIM))
|
||||
|
||||
@staticmethod
|
||||
def repl(agent):
|
||||
if readline is not None:
|
||||
try:
|
||||
readline.read_history_file(agent.config.history_path)
|
||||
except OSError:
|
||||
pass
|
||||
readline.set_completer(Commands.complete_command)
|
||||
readline.parse_and_bind("tab: complete")
|
||||
scheduler_stop = Scheduler.start_scheduler(agent.config, agent.store.seal)
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
busy = REGISTRY.running_agents()
|
||||
base = agent.profile + ("!" if agent.yolo else "")
|
||||
tag = "%s+%d" % (base, busy) if busy else base
|
||||
line = input(Render.paint("tai[%s]› " % tag, Ansi.BOLD, Ansi.GREEN))
|
||||
except EOFError:
|
||||
print()
|
||||
break
|
||||
except KeyboardInterrupt:
|
||||
print()
|
||||
continue
|
||||
text = line.strip()
|
||||
if not text:
|
||||
continue
|
||||
if text.startswith("/"):
|
||||
if not Commands.handle_command(agent, text):
|
||||
break
|
||||
continue
|
||||
try:
|
||||
agent.run_turn(text)
|
||||
except KeyboardInterrupt:
|
||||
print(Render.paint("\ninterrupted", Ansi.YELLOW))
|
||||
agent.checkpoint()
|
||||
finally:
|
||||
scheduler_stop.set()
|
||||
if readline is not None:
|
||||
try:
|
||||
readline.write_history_file(agent.config.history_path)
|
||||
except OSError:
|
||||
pass
|
||||
agent.store.save_session(agent.profile, agent.messages, agent.bot)
|
||||
|
||||
@staticmethod
|
||||
def main(argv=None):
|
||||
args = App.build_parser().parse_args(argv)
|
||||
if args.install:
|
||||
return Install.install_self()
|
||||
if args.install_telegram:
|
||||
return Install.install_telegram()
|
||||
if args.uninstall_telegram:
|
||||
return Install.uninstall_telegram()
|
||||
if args.version:
|
||||
print("tai %s" % Consts.VERSION)
|
||||
return 0
|
||||
if args.plexer_wrap:
|
||||
home = os.environ.get("TAI_HOME") or os.path.join(os.path.expanduser("~"), ".tai")
|
||||
profile = args.profile or "default"
|
||||
backend = NativePlexerBackend(home, profile)
|
||||
return backend.run()
|
||||
if args.profile is None:
|
||||
args.profile = "telegram" if args.telegram else "default"
|
||||
if not Consts.PROFILE_RE.match(args.profile):
|
||||
print("invalid profile name", file=sys.stderr)
|
||||
return 1
|
||||
try:
|
||||
agent = App.boot(args)
|
||||
except KeyboardInterrupt:
|
||||
print(Render.paint("\ninterrupted", Ansi.YELLOW))
|
||||
return 130
|
||||
except SealError as exc:
|
||||
print("seal error: %s" % exc, file=sys.stderr)
|
||||
return 2
|
||||
if args.env is not None:
|
||||
if args.env not in ("home", "sandbox"):
|
||||
print("use --env home or --env sandbox", file=sys.stderr)
|
||||
return 1
|
||||
if args.env == "sandbox":
|
||||
try:
|
||||
Sandbox.ensure_box()
|
||||
except BackendError as exc:
|
||||
print("sandbox unavailable: %s" % exc, file=sys.stderr)
|
||||
return 2
|
||||
agent.env = args.env
|
||||
if args.scheduler:
|
||||
print("scheduler online, press Ctrl-C to stop")
|
||||
try:
|
||||
Scheduler.run_scheduler_loop(agent.config, agent.store.seal, threading.Event())
|
||||
except KeyboardInterrupt:
|
||||
print()
|
||||
finally:
|
||||
agent.store.close()
|
||||
return 0
|
||||
if args.telegram:
|
||||
return TelegramBot.run_telegram_bot(agent)
|
||||
if args.prompt:
|
||||
try:
|
||||
agent.run_turn(" ".join(args.prompt))
|
||||
except KeyboardInterrupt:
|
||||
print(Render.paint("\ninterrupted", Ansi.YELLOW))
|
||||
return 130
|
||||
finally:
|
||||
try:
|
||||
agent.checkpoint()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
agent.store.close()
|
||||
except Exception:
|
||||
pass
|
||||
return 0
|
||||
if args.resume and not args.telegram and not args.scheduler:
|
||||
Agent.resume_turn(agent)
|
||||
App.banner(agent)
|
||||
try:
|
||||
App.repl(agent)
|
||||
finally:
|
||||
agent.store.close()
|
||||
return 0
|
||||
+479
@@ -0,0 +1,479 @@
|
||||
"""Multi-backend chat router."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
from tai_pkg.claude import Claude
|
||||
from tai_pkg.codex import Codex
|
||||
from tai_pkg.config import Config
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.errors import BackendError
|
||||
from tai_pkg.gemini import Gemini
|
||||
from tai_pkg.grok import Grok
|
||||
from tai_pkg.models import ModelHealth
|
||||
from tai_pkg.models import ModelScoring
|
||||
from tai_pkg.muse import Muse
|
||||
from tai_pkg.opencode import Opencode
|
||||
from tai_pkg.render import Ansi
|
||||
from tai_pkg.render import Render
|
||||
from tai_pkg.sandbox import Sandbox
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["ChatClient"]
|
||||
|
||||
|
||||
class ChatClient:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
self.scope = ("default", "main")
|
||||
self.fixed_key = None
|
||||
self.health = ModelHealth(getattr(config, "home", "") or "")
|
||||
|
||||
def muse_key(self):
|
||||
if self.fixed_key is not None:
|
||||
return self.fixed_key
|
||||
return "%s/%s" % (self.scope[0], self.scope[1])
|
||||
|
||||
def reset_muse(self):
|
||||
home = getattr(self.config, "home", None)
|
||||
if home is None:
|
||||
return
|
||||
key = self.muse_key()
|
||||
|
||||
def drop(data):
|
||||
data.pop(key, None)
|
||||
|
||||
Muse.muse_transact(home, drop)
|
||||
|
||||
def openrouter_models(self):
|
||||
models = [item for item in getattr(self.config, "openrouter_models", []) if item]
|
||||
explicit = getattr(self.config, "explicit_model", "")
|
||||
if explicit and explicit not in models:
|
||||
models.insert(0, explicit)
|
||||
free = self.health.cached_list("openrouter:free")
|
||||
if free is None:
|
||||
fetched = ModelScoring.fetch_model_ids(Consts.OPENROUTER_BASE + "/models", self.config.openrouter_key)
|
||||
free = [item for item in fetched if item.endswith(":free")] if fetched else []
|
||||
if free:
|
||||
self.health.store_list("openrouter:free", free)
|
||||
free = ModelScoring.prefer_muse(free or [])
|
||||
models = [item for item in free if item not in models] + models
|
||||
if not models:
|
||||
models = [Consts.OPENROUTER_MODEL]
|
||||
return models
|
||||
|
||||
def opencode_models(self):
|
||||
models = list(getattr(self.config, "opencode_models", []) or [])
|
||||
if not models:
|
||||
models = self.health.cached_list("opencode:free")
|
||||
if models is None:
|
||||
models = Opencode.opencode_free_models() or [item for item in Consts.ZEN_FREE_MODELS if Text.is_free_model(item)]
|
||||
self.health.store_list("opencode:free", models)
|
||||
return ModelScoring.prefer_muse(models)
|
||||
|
||||
def ollama_models(self):
|
||||
if getattr(self.config, "ollama_enabled", True) is False:
|
||||
return []
|
||||
models = list(getattr(self.config, "ollama_models", []) or [])
|
||||
if not models:
|
||||
models = self.health.cached_list("ollama:local", Consts.OLLAMA_LIST_TTL)
|
||||
if models is None:
|
||||
models = ModelScoring.fetch_model_ids(Consts.OLLAMA_BASE + "/models", "", timeout=3) or []
|
||||
if models:
|
||||
self.health.store_list("ollama:local", models)
|
||||
return models
|
||||
|
||||
def devplace_models(self):
|
||||
models = list(getattr(self.config, "devplace_models", []) or [])
|
||||
if models:
|
||||
return ModelScoring.prefer_muse(models)
|
||||
free = self.health.cached_list("devplace:free")
|
||||
if free is None:
|
||||
fetched = ModelScoring.fetch_model_ids(self.config.devplace_base + "/models", self.config.devplace_key)
|
||||
if fetched is None:
|
||||
return list(Consts.DEVPLACE_FREE_ROUTES)
|
||||
free = [item for item in fetched if item in Consts.DEVPLACE_FREE_ROUTES or (Consts.MUSE_LABEL in item.lower() and Text.is_free_model(item))]
|
||||
if free:
|
||||
self.health.store_list("devplace:free", free)
|
||||
return ModelScoring.prefer_muse(free or [])
|
||||
|
||||
def candidates(self):
|
||||
found = []
|
||||
if getattr(self.config, "opencode_enabled", True) and Opencode.opencode_binary():
|
||||
for model in self.opencode_models():
|
||||
found.append({"provider": Consts.OPENCODE_LABEL, "chat_url": Consts.OPENCODE_CHAT_URL, "models_url": Consts.OPENCODE_MODELS_URL, "key": "", "model": model, "profile": ""})
|
||||
for model in self.ollama_models():
|
||||
found.append({"provider": Consts.OLLAMA_LABEL, "chat_url": Consts.OLLAMA_BASE + "/chat/completions", "models_url": Consts.OLLAMA_BASE + "/models", "key": "", "model": model, "profile": ""})
|
||||
if getattr(self.config, "devplace_key", ""):
|
||||
for model in self.devplace_models():
|
||||
found.append({"provider": Consts.DEVPLACE_LABEL, "chat_url": self.config.devplace_base + "/chat/completions", "models_url": self.config.devplace_base + "/models", "key": self.config.devplace_key, "model": model, "profile": ""})
|
||||
for entry in self.config.backends:
|
||||
for model in Config.split_models(entry.get("model")) or [self.config.model]:
|
||||
found.append({"provider": entry["label"], "chat_url": entry["base"] + "/chat/completions", "models_url": entry["base"] + "/models", "key": entry["key"], "model": model, "profile": entry.get("profile") or ""})
|
||||
if getattr(self.config, "openrouter_enabled", True):
|
||||
for model in self.openrouter_models():
|
||||
found.append({"provider": Consts.OPENROUTER_LABEL, "chat_url": Consts.OPENROUTER_BASE + "/chat/completions", "models_url": Consts.OPENROUTER_BASE + "/models", "key": self.config.openrouter_key, "model": model, "profile": ""})
|
||||
if getattr(self.config, "pollinations_enabled", True):
|
||||
for model in getattr(self.config, "pollinations_models", []) or list(Consts.POLLINATIONS_MODELS):
|
||||
found.append({"provider": Consts.POLLINATIONS_LABEL, "chat_url": Consts.POLLINATIONS_BASE + "/chat/completions", "models_url": Consts.POLLINATIONS_BASE + "/models", "key": "", "model": model, "profile": ""})
|
||||
if getattr(self.config, "zen_enabled", False):
|
||||
for model in getattr(self.config, "zen_models", []) or list(Consts.ZEN_FREE_MODELS):
|
||||
found.append({"provider": Consts.ZEN_LABEL, "chat_url": Consts.ZEN_BASE + "/chat/completions", "models_url": Consts.ZEN_BASE + "/models", "key": "", "model": model, "profile": "opencode"})
|
||||
found.append({"provider": Consts.MUSE_LABEL, "chat_url": Consts.MUSE_CHAT_URL, "models_url": Consts.MUSE_MODELS_URL, "key": "", "model": self.config.model, "profile": ""})
|
||||
if Grok.grok_binary():
|
||||
found.append({"provider": Consts.GROK_LABEL, "chat_url": Consts.GROK_CHAT_URL, "models_url": Consts.GROK_MODELS_URL, "key": "", "model": self.config.model, "profile": ""})
|
||||
if Claude.claude_binary():
|
||||
found.append({"provider": Consts.CLAUDE_LABEL, "chat_url": Consts.CLAUDE_CHAT_URL, "models_url": Consts.CLAUDE_MODELS_URL, "key": "", "model": self.config.model, "profile": ""})
|
||||
if Codex.codex_binary():
|
||||
found.append({"provider": Consts.CODEX_LABEL, "chat_url": Consts.CODEX_CHAT_URL, "models_url": Consts.CODEX_MODELS_URL, "key": "", "model": self.config.model, "profile": ""})
|
||||
if Gemini.gemini_binary():
|
||||
found.append({"provider": Consts.GEMINI_LABEL, "chat_url": Consts.GEMINI_CHAT_URL, "models_url": Consts.GEMINI_MODELS_URL, "key": "", "model": self.config.model, "profile": ""})
|
||||
return found
|
||||
|
||||
def ordered_candidates(self):
|
||||
http = [cand for cand in self.candidates() if cand["chat_url"] not in Consts.CLI_CHAT_URLS]
|
||||
cli = [cand for cand in self.candidates() if cand["chat_url"] in Consts.CLI_CHAT_URLS]
|
||||
ranked = sorted(enumerate(http), key=lambda pair: (ModelScoring.openrouter_free_tier(pair[1]), -self.health.weight(pair[1]["provider"], pair[1]["model"]), pair[0]))
|
||||
usable = [cand for _, cand in ranked if not self.health.circuit_open(cand["provider"], cand["model"])]
|
||||
return (usable or [cand for _, cand in ranked]) + cli
|
||||
|
||||
def backends(self):
|
||||
rows = []
|
||||
seen = set()
|
||||
for cand in self.candidates():
|
||||
key = (cand["provider"], cand["chat_url"])
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
rows.append((cand["provider"], cand["chat_url"], cand["models_url"], cand["key"], cand["model"]))
|
||||
return rows
|
||||
|
||||
def headers(self, key, extra=None):
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if key:
|
||||
headers["Authorization"] = "Bearer " + key
|
||||
if extra:
|
||||
headers.update(extra)
|
||||
return headers
|
||||
|
||||
def complete(self, messages, tools=None, stream_sink=None, ephemeral=False, deadline=None):
|
||||
payload = {"model": self.config.model, "messages": messages}
|
||||
if tools:
|
||||
payload["tools"] = tools
|
||||
payload["tool_choice"] = "auto"
|
||||
limit = getattr(self.config, "retry_seconds", 0)
|
||||
give_up = time.time() + limit if limit > 0 else None
|
||||
if deadline is not None:
|
||||
give_up = deadline if give_up is None else min(give_up, deadline)
|
||||
round_no = 0
|
||||
while True:
|
||||
errors = []
|
||||
cands = self.ordered_candidates() if round_no == 0 else self.all_candidates()
|
||||
for cand in cands:
|
||||
reply = self.try_candidate(cand, messages, tools, payload, stream_sink, ephemeral, errors)
|
||||
if reply is not None:
|
||||
return reply
|
||||
pause = Consts.RETRY_BACKOFF[min(round_no, len(Consts.RETRY_BACKOFF) - 1)]
|
||||
round_no += 1
|
||||
if give_up is not None and time.time() + pause > give_up:
|
||||
raise BackendError("all backends failed after %d rounds (%s)" % (round_no, "; ".join(errors[-6:])))
|
||||
print(Render.paint("all %d backends failed (round %d), retrying in %ds: %s" % (len(cands), round_no, pause, (errors[-1] if errors else "")[:160]), Ansi.DIM), file=sys.stderr)
|
||||
if round_no == 1 and not Opencode.opencode_binary():
|
||||
engine = Sandbox.container_engine()
|
||||
if engine and Sandbox.box_state(engine) == "missing":
|
||||
print(Render.paint(
|
||||
"no local opencode CLI, and no sandbox box set up yet either -- that box "
|
||||
"would also work as a last-resort backend (see box_opencode_candidates()). "
|
||||
"Run `tai.py --env sandbox` once to build it: a small image, podman or "
|
||||
"docker, and it only ever touches that one container, never the rest of "
|
||||
"this machine.", Ansi.DIM
|
||||
), file=sys.stderr)
|
||||
time.sleep(pause)
|
||||
|
||||
def all_candidates(self):
|
||||
http = [cand for cand in self.candidates() if cand["chat_url"] not in Consts.CLI_CHAT_URLS]
|
||||
cli = [cand for cand in self.candidates() if cand["chat_url"] in Consts.CLI_CHAT_URLS]
|
||||
ranked = sorted(enumerate(http), key=lambda pair: (ModelScoring.openrouter_free_tier(pair[1]), -self.health.weight(pair[1]["provider"], pair[1]["model"]), pair[0]))
|
||||
return [cand for _, cand in ranked] + cli + self.box_opencode_candidates()
|
||||
|
||||
def box_opencode_candidates(self):
|
||||
"""Last-resort only -- complete() reaches this via all_candidates(),
|
||||
never via ordered_candidates() (the round-0 fast path), so a normal
|
||||
turn that succeeds through a real backend never touches the box at
|
||||
all. If there's already a local opencode CLI, self.candidates()
|
||||
covers it directly and this stays empty; it only fires when that's
|
||||
missing AND the sandbox box is already running (never auto-created
|
||||
here -- see ensure_box(), which stays an explicit, approval-gated
|
||||
action via --env sandbox or the install tool)."""
|
||||
if not getattr(self.config, "opencode_enabled", True) or Opencode.opencode_binary():
|
||||
return []
|
||||
engine = Sandbox.container_engine()
|
||||
if not engine or Sandbox.box_state(engine) != "running":
|
||||
return []
|
||||
if not Sandbox.box_opencode_binary(engine):
|
||||
return []
|
||||
models = self.box_opencode_models(engine)
|
||||
return [
|
||||
{"provider": Consts.OPENCODE_LABEL, "chat_url": Consts.OPENCODE_CHAT_URL, "models_url": Consts.OPENCODE_MODELS_URL,
|
||||
"key": "", "model": model, "profile": "", "runner": "box", "engine": engine}
|
||||
for model in models
|
||||
]
|
||||
|
||||
def box_opencode_models(self, engine):
|
||||
cached = self.health.cached_list("opencode:box:free")
|
||||
if cached is not None:
|
||||
return ModelScoring.prefer_muse(cached)
|
||||
try:
|
||||
done = Sandbox.box_exec(engine, ["opencode", "models", "opencode"], timeout=Consts.MODEL_LIST_TIMEOUT * 4)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
return []
|
||||
text = done.stdout.decode("utf-8", "replace")
|
||||
models = [line.strip().split("/", 1)[1] for line in text.splitlines() if line.strip().startswith("opencode/")]
|
||||
models = [item for item in models if Text.is_free_model(item)]
|
||||
if models:
|
||||
self.health.store_list("opencode:box:free", models)
|
||||
return ModelScoring.prefer_muse(models)
|
||||
|
||||
def call_candidate(self, cand, messages, tools, attempt, stream_sink, ephemeral):
|
||||
label, chat_url, key, model = cand["provider"], cand["chat_url"], cand["key"], cand["model"]
|
||||
extra = Opencode.opencode_headers() if cand["profile"] == "opencode" else None
|
||||
timeout = Consts.STREAM_TIMEOUT if stream_sink is not None else Consts.HTTP_TIMEOUT
|
||||
if chat_url == Consts.MUSE_CHAT_URL:
|
||||
return self.muse_turn(messages, tools, stream_sink, label, timeout=timeout, ephemeral=ephemeral)
|
||||
if chat_url == Consts.GROK_CHAT_URL:
|
||||
return self.grok_turn(messages, tools, stream_sink, label, timeout=timeout)
|
||||
if chat_url == Consts.CLAUDE_CHAT_URL:
|
||||
return self.claude_turn(messages, tools, stream_sink, label, timeout=timeout)
|
||||
if chat_url == Consts.CODEX_CHAT_URL:
|
||||
return self.codex_turn(messages, tools, stream_sink, label, timeout=timeout)
|
||||
if chat_url == Consts.GEMINI_CHAT_URL:
|
||||
return self.gemini_turn(messages, tools, stream_sink, label, timeout=timeout)
|
||||
if chat_url == Consts.OPENCODE_CHAT_URL:
|
||||
if cand.get("runner") == "box":
|
||||
return self.box_opencode_turn(messages, tools, stream_sink, label, model, cand["engine"], timeout=timeout)
|
||||
return self.opencode_turn(messages, tools, stream_sink, label, model, timeout=timeout)
|
||||
if stream_sink is None:
|
||||
return self.single_shot(chat_url, key, attempt, label, extra)
|
||||
return self.streaming(chat_url, key, attempt, label, stream_sink, extra)
|
||||
|
||||
def try_candidate(self, cand, messages, tools, payload, stream_sink, ephemeral, errors):
|
||||
label, model = cand["provider"], cand["model"]
|
||||
attempt = dict(payload, model=model)
|
||||
tries = [(tools, attempt)]
|
||||
if attempt.get("tools"):
|
||||
tries.append((None, {name: value for name, value in attempt.items() if name not in ("tools", "tool_choice")}))
|
||||
for pos, (use_tools, body) in enumerate(tries):
|
||||
started = time.time()
|
||||
try:
|
||||
reply = self.call_candidate(cand, messages, use_tools, body, stream_sink, ephemeral)
|
||||
if not isinstance(reply, dict):
|
||||
raise BackendError("backend returned no reply")
|
||||
self.health.record(label, model, True, (time.time() - started) * 1000, max(1, len(reply.get("content") or "") // 4) / max(time.time() - started, 0.001))
|
||||
return reply
|
||||
except BackendError as exc:
|
||||
if pos == 0 and exc.status == 400 and len(tries) > 1:
|
||||
continue
|
||||
self.health.record(label, model, False)
|
||||
errors.append("%s/%s: %s" % (label, model, exc))
|
||||
return None
|
||||
except Exception as exc:
|
||||
self.health.record(label, model, False)
|
||||
errors.append("%s/%s: %s" % (label, model, Text.short_error(exc)))
|
||||
return None
|
||||
return None
|
||||
|
||||
def single_shot(self, url, key, payload, label, extra=None):
|
||||
body = json.dumps(payload).encode("utf-8")
|
||||
request = urllib.request.Request(url, data=body, headers=self.headers(key, extra), method="POST")
|
||||
try:
|
||||
with urllib.request.urlopen(request, timeout=Consts.HTTP_TIMEOUT) as response:
|
||||
data = json.loads(response.read().decode("utf-8", "replace"))
|
||||
except urllib.error.HTTPError as exc:
|
||||
raise BackendError("HTTP %d: %s" % (exc.code, exc.read().decode("utf-8", "replace")[:300]), exc.code)
|
||||
except urllib.error.URLError as exc:
|
||||
raise BackendError(Text.short_error(exc))
|
||||
failure = ModelScoring.embedded_failure_message(data)
|
||||
if failure:
|
||||
raise BackendError("upstream error: %s" % failure[:300])
|
||||
try:
|
||||
message = data["choices"][0]["message"]
|
||||
except (KeyError, IndexError, TypeError):
|
||||
raise BackendError("upstream returned no choices: %s" % json.dumps(data)[:300])
|
||||
return self.normalize(message, label)
|
||||
|
||||
def normalize(self, message, label):
|
||||
calls = []
|
||||
for call in message.get("tool_calls") or []:
|
||||
func = call.get("function") or {}
|
||||
calls.append({"id": call.get("id") or uuid.uuid4().hex, "name": func.get("name") or "", "arguments": Text.tool_args_text(func.get("arguments"))})
|
||||
return {"role": "assistant", "content": message.get("content") or "", "reasoning": message.get("reasoning") or message.get("reasoning_content") or "", "tool_calls": calls, "backend": label}
|
||||
|
||||
def streaming(self, url, key, payload, label, sink, extra=None):
|
||||
body = json.dumps(dict(payload, stream=True)).encode("utf-8")
|
||||
request = urllib.request.Request(url, data=body, headers=self.headers(key, extra), method="POST")
|
||||
try:
|
||||
response = urllib.request.urlopen(request, timeout=Consts.STREAM_TIMEOUT)
|
||||
except urllib.error.HTTPError as exc:
|
||||
raise BackendError("HTTP %d: %s" % (exc.code, exc.read().decode("utf-8", "replace")[:300]), exc.code)
|
||||
except urllib.error.URLError as exc:
|
||||
raise BackendError(Text.short_error(exc))
|
||||
content_parts = []
|
||||
reasoning_parts = []
|
||||
calls = {}
|
||||
try:
|
||||
for raw in response:
|
||||
line = raw.decode("utf-8", "replace").strip()
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
data = line[5:].strip()
|
||||
if data == "[DONE]":
|
||||
break
|
||||
try:
|
||||
event = json.loads(data)
|
||||
except ValueError:
|
||||
continue
|
||||
choices = event.get("choices") or []
|
||||
if not choices:
|
||||
continue
|
||||
delta = choices[0].get("delta") or {}
|
||||
piece = delta.get("content")
|
||||
if piece:
|
||||
content_parts.append(piece)
|
||||
sink(piece)
|
||||
think = delta.get("reasoning") or delta.get("reasoning_content")
|
||||
if think:
|
||||
reasoning_parts.append(think)
|
||||
for call in delta.get("tool_calls") or []:
|
||||
slot = calls.setdefault(call.get("index", 0), {"id": "", "name": "", "arguments": ""})
|
||||
if call.get("id"):
|
||||
slot["id"] = call["id"]
|
||||
func = call.get("function") or {}
|
||||
if func.get("name"):
|
||||
slot["name"] += func["name"]
|
||||
if func.get("arguments"):
|
||||
slot["arguments"] += Text.tool_args_text(func["arguments"])
|
||||
except OSError as exc:
|
||||
if not content_parts and not calls:
|
||||
raise BackendError(Text.short_error(exc))
|
||||
finally:
|
||||
response.close()
|
||||
ordered = [calls[pos] for pos in sorted(calls)]
|
||||
for slot in ordered:
|
||||
if not slot["id"]:
|
||||
slot["id"] = uuid.uuid4().hex
|
||||
return {"role": "assistant", "content": "".join(content_parts), "reasoning": "".join(reasoning_parts), "tool_calls": ordered, "backend": label}
|
||||
|
||||
def muse_turn(self, messages, tools, stream_sink, label, timeout=Consts.HTTP_TIMEOUT, ephemeral=False):
|
||||
binary = Muse.muse_binary()
|
||||
if not binary:
|
||||
raise BackendError("muse CLI not found (MUSE_BIN or PATH)")
|
||||
home = getattr(self.config, "home", None)
|
||||
if ephemeral or home is None:
|
||||
with tempfile.TemporaryDirectory(prefix="tai-muse-") as folder:
|
||||
schema_path = os.path.join(folder, "schema.json")
|
||||
prompt_path = os.path.join(folder, "prompt.txt")
|
||||
Muse.muse_write_file(schema_path, json.dumps(Muse.muse_output_schema()))
|
||||
Muse.muse_write_file(prompt_path, Muse.muse_prompt(messages, tools))
|
||||
answer = Muse.muse_run(Muse.muse_argv(self.config, schema_path, prompt_path), folder, stream_sink, timeout)
|
||||
return self.muse_reply(answer, label)
|
||||
key = self.muse_key()
|
||||
release = Muse.muse_acquire(home, "session:" + key)
|
||||
|
||||
def on_reap(ok):
|
||||
try:
|
||||
if not ok:
|
||||
self.reset_muse()
|
||||
finally:
|
||||
release()
|
||||
|
||||
gated = Muse.muse_gate(on_reap)
|
||||
try:
|
||||
return self.muse_persistent(messages, tools, stream_sink, label, timeout, home, key, gated)
|
||||
except BaseException:
|
||||
if not gated.armed.is_set():
|
||||
release()
|
||||
raise
|
||||
|
||||
def muse_persistent(self, messages, tools, stream_sink, label, timeout, home, key, gated):
|
||||
def plan(data):
|
||||
entry = data.get(key)
|
||||
pending = None
|
||||
session_id = None
|
||||
if isinstance(entry, dict) and isinstance(entry.get("id"), str) and isinstance(entry.get("hashes"), list):
|
||||
pending = Muse.muse_session_split(messages, entry["hashes"])
|
||||
if pending is not None and not Muse.muse_session_alive(entry["id"]):
|
||||
pending = None
|
||||
if pending is not None:
|
||||
session_id = entry["id"]
|
||||
if pending is None or not pending:
|
||||
session_id = str(uuid.uuid4())
|
||||
pending = list(messages)
|
||||
return session_id, pending
|
||||
|
||||
session_id, pending = Muse.muse_transact(home, plan)
|
||||
try:
|
||||
with tempfile.TemporaryDirectory(prefix="tai-muse-") as folder:
|
||||
schema_path = os.path.join(folder, "schema.json")
|
||||
prompt_path = os.path.join(folder, "prompt.txt")
|
||||
Muse.muse_write_file(schema_path, json.dumps(Muse.muse_output_schema()))
|
||||
Muse.muse_write_file(prompt_path, Muse.muse_prompt(pending, tools))
|
||||
answer = Muse.muse_run(Muse.muse_argv(self.config, schema_path, prompt_path, session_id), folder, stream_sink, timeout, gated)
|
||||
except BackendError:
|
||||
self.reset_muse()
|
||||
raise
|
||||
|
||||
def commit(data):
|
||||
data[key] = {"id": session_id, "hashes": [Muse.muse_message_hash(item) for item in messages], "updated": Text.now_iso()}
|
||||
|
||||
Muse.muse_transact(home, commit)
|
||||
return self.muse_reply(answer, label)
|
||||
|
||||
def opencode_turn(self, messages, tools, stream_sink, label, model, timeout=Consts.HTTP_TIMEOUT):
|
||||
answer = Opencode.opencode_parse(Opencode.opencode_run(model, Opencode.opencode_prompt(messages, tools), timeout))
|
||||
if stream_sink is not None and answer.get("content"):
|
||||
stream_sink(answer["content"])
|
||||
return self.muse_reply(answer, label)
|
||||
|
||||
def box_opencode_turn(self, messages, tools, stream_sink, label, model, engine, timeout=Consts.HTTP_TIMEOUT):
|
||||
answer = Opencode.opencode_parse(Sandbox.box_opencode_run(engine, model, Opencode.opencode_prompt(messages, tools), timeout))
|
||||
if stream_sink is not None and answer.get("content"):
|
||||
stream_sink(answer["content"])
|
||||
return self.muse_reply(answer, label + " (box)")
|
||||
|
||||
def grok_turn(self, messages, tools, stream_sink, label, timeout=Consts.HTTP_TIMEOUT):
|
||||
answer = Grok.grok_parse_output(Grok.grok_run(Opencode.opencode_prompt(messages, tools), timeout))
|
||||
if stream_sink is not None and answer.get("content"):
|
||||
stream_sink(answer["content"])
|
||||
return self.muse_reply(answer, label)
|
||||
|
||||
def claude_turn(self, messages, tools, stream_sink, label, timeout=Consts.HTTP_TIMEOUT):
|
||||
answer = Claude.claude_parse_output(Claude.claude_run(Opencode.opencode_prompt(messages, tools), timeout))
|
||||
if stream_sink is not None and answer.get("content"):
|
||||
stream_sink(answer["content"])
|
||||
return self.muse_reply(answer, label)
|
||||
|
||||
def codex_turn(self, messages, tools, stream_sink, label, timeout=Consts.HTTP_TIMEOUT):
|
||||
stdout, outfile_text, used_schema = Codex.codex_run(Opencode.opencode_prompt(messages, tools), timeout)
|
||||
answer = Codex.codex_parse_output(stdout, outfile_text, used_schema)
|
||||
if stream_sink is not None and answer.get("content"):
|
||||
stream_sink(answer["content"])
|
||||
return self.muse_reply(answer, label)
|
||||
|
||||
def gemini_turn(self, messages, tools, stream_sink, label, timeout=Consts.HTTP_TIMEOUT):
|
||||
answer = Gemini.gemini_parse_output(Gemini.gemini_run(Opencode.opencode_prompt(messages, tools), timeout))
|
||||
if stream_sink is not None and answer.get("content"):
|
||||
stream_sink(answer["content"])
|
||||
return self.muse_reply(answer, label)
|
||||
|
||||
def muse_reply(self, answer, label):
|
||||
calls = []
|
||||
for call in answer.get("tool_calls") or []:
|
||||
calls.append({"id": uuid.uuid4().hex, "name": call.get("name") or "", "arguments": Text.tool_args_text(call.get("arguments"))})
|
||||
return {"role": "assistant", "content": answer.get("content") or "", "reasoning": "", "tool_calls": calls, "backend": label}
|
||||
@@ -0,0 +1,97 @@
|
||||
"""Claude CLI backend."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
|
||||
from tai_pkg.clirun import Cli
|
||||
from tai_pkg.errors import BackendError
|
||||
from tai_pkg.muse import Muse
|
||||
from tai_pkg.opencode import Opencode
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["Claude"]
|
||||
|
||||
|
||||
class Claude:
|
||||
"""Claude."""
|
||||
|
||||
@staticmethod
|
||||
def claude_binary():
|
||||
return os.environ.get("CLAUDE_BIN") or shutil.which("claude")
|
||||
|
||||
@staticmethod
|
||||
def claude_parse_output(stdout):
|
||||
try:
|
||||
data = json.loads(stdout or "")
|
||||
except ValueError:
|
||||
return Opencode.opencode_parse(stdout)
|
||||
if not isinstance(data, dict):
|
||||
return Opencode.opencode_parse(stdout)
|
||||
if data.get("is_error"):
|
||||
raise BackendError("claude: %s" % str(data.get("result") or "unknown error")[:300])
|
||||
structured = data.get("structured_output")
|
||||
if isinstance(structured, dict) and ("content" in structured or "tool_calls" in structured):
|
||||
return structured
|
||||
result = data.get("result")
|
||||
if isinstance(result, str) and result.strip():
|
||||
return Opencode.opencode_parse(result)
|
||||
return Opencode.opencode_parse(stdout)
|
||||
|
||||
@staticmethod
|
||||
def claude_run(prompt, timeout):
|
||||
binary = Claude.claude_binary()
|
||||
if not binary:
|
||||
raise BackendError("claude CLI not found (CLAUDE_BIN or PATH)")
|
||||
schema = json.dumps(Muse.muse_output_schema())
|
||||
base = [binary, "-p"]
|
||||
if len(prompt.encode("utf-8", "replace")) > 120000:
|
||||
variants = [
|
||||
(base + ["--output-format", "json", "--json-schema", schema], prompt),
|
||||
(base + ["--output-format", "json"], prompt),
|
||||
]
|
||||
else:
|
||||
variants = [
|
||||
(base + [prompt, "--output-format", "json", "--json-schema", schema], None),
|
||||
(base + [prompt, "--output-format", "json"], None),
|
||||
(base + [prompt], None),
|
||||
]
|
||||
last_error = ""
|
||||
for argv, prompt_input in variants:
|
||||
done = Cli._cli_run(argv, prompt_input, timeout, "claude")
|
||||
if done.returncode != 0 and Cli._cli_unknown_flag(done.stderr) and (argv, prompt_input) != variants[-1]:
|
||||
last_error = (done.stderr or "").strip()[:200]
|
||||
continue
|
||||
if done.returncode != 0 and not (done.stdout or "").strip():
|
||||
blob = ((done.stderr or "") + "\n" + (done.stdout or "")).strip()
|
||||
if Cli._cli_auth_failure(blob):
|
||||
raise BackendError("claude not authenticated (run claude login): %s" % blob[:200])
|
||||
raise BackendError("claude exit %d: %s" % (done.returncode, blob[:300] or last_error or "no output"))
|
||||
if not (done.stdout or "").strip():
|
||||
blob = (done.stderr or "").strip()
|
||||
raise BackendError("claude returned no output: %s" % (blob[:300] or last_error or "empty stdout"))
|
||||
return done.stdout or ""
|
||||
raise BackendError("claude failed: %s" % (last_error or "no usable output"))
|
||||
|
||||
@staticmethod
|
||||
def claude_health():
|
||||
binary = Claude.claude_binary()
|
||||
if not binary:
|
||||
return "missing: claude CLI not found"
|
||||
try:
|
||||
done = subprocess.run([binary, "auth", "status"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=15)
|
||||
except (OSError, subprocess.SubprocessError) as exc:
|
||||
return "fail: %s" % Text.short_error(exc)[:60]
|
||||
try:
|
||||
info = json.loads(done.stdout)
|
||||
except ValueError:
|
||||
info = None
|
||||
if isinstance(info, dict) and "loggedIn" in info:
|
||||
if info["loggedIn"]:
|
||||
return "ok, subscription (%s)" % (info.get("authMethod") or "logged in")
|
||||
return "not logged in"
|
||||
# `claude auth status` didn't return the expected shape (older/newer
|
||||
# CLI) -- fall back to the honest version-only signal instead of
|
||||
# claiming a subscription this check couldn't actually verify.
|
||||
return Cli._cli_version_only_health(binary, "claude")
|
||||
@@ -0,0 +1,53 @@
|
||||
"""Shared subprocess runners for CLI backends."""
|
||||
|
||||
import re
|
||||
import subprocess
|
||||
|
||||
from tai_pkg.errors import BackendError
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["Cli"]
|
||||
|
||||
|
||||
class Cli:
|
||||
"""Cli."""
|
||||
|
||||
@staticmethod
|
||||
def _cli_unknown_flag(stderr):
|
||||
return bool(re.search(r"unknown (option|flag|command)|unrecognised|unrecognized|invalid flag|not found|no such option", stderr or "", re.IGNORECASE))
|
||||
|
||||
@staticmethod
|
||||
def _cli_auth_failure(text):
|
||||
return bool(re.search(r"not (logged in|authenticated)|auth|api key|unauthorized|401|login required|please (login|log in|authenticate)", text or "", re.IGNORECASE))
|
||||
|
||||
@staticmethod
|
||||
def _cli_run(argv, prompt_input, timeout, label):
|
||||
try:
|
||||
done = subprocess.run(argv, input=prompt_input, stdin=subprocess.DEVNULL if prompt_input is None else None, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=timeout)
|
||||
except subprocess.TimeoutExpired:
|
||||
raise BackendError("%s run timed out after %ds" % (label, timeout))
|
||||
except OSError as exc:
|
||||
raise BackendError("%s run failed: %s" % (label, Text.short_error(exc)))
|
||||
return done
|
||||
|
||||
@staticmethod
|
||||
def _cli_version_detail(done, fallback):
|
||||
return (done.stdout.strip() or done.stderr.strip()).splitlines()[0][:60] if (done.stdout.strip() or done.stderr.strip()) else fallback
|
||||
|
||||
@staticmethod
|
||||
def _cli_version_only_health(binary, fallback_name):
|
||||
"""Shared by the CLIs with no cheap, read-only way to check real login
|
||||
state (muse, grok): `--version` only proves the binary runs, never that
|
||||
it's actually authenticated, so this deliberately never claims "ok" or
|
||||
"subscription" -- a claim this check cannot back up (see codex_health/
|
||||
claude_health for the two CLIs that DO expose a real status command,
|
||||
used instead of this). Caught a real bug: this used to print "ok,
|
||||
subscription" unconditionally on a successful --version, which is false
|
||||
whenever the CLI is installed but never logged in."""
|
||||
try:
|
||||
done = subprocess.run([binary, "--version"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=15)
|
||||
except (OSError, subprocess.SubprocessError) as exc:
|
||||
return "fail: %s" % Text.short_error(exc)[:60]
|
||||
if done.returncode != 0:
|
||||
return "fail: exit %d" % done.returncode
|
||||
return "cli found (%s), login unverified: no status command" % Cli._cli_version_detail(done, fallback_name)
|
||||
@@ -0,0 +1,134 @@
|
||||
"""Codex CLI backend."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
|
||||
from tai_pkg.clirun import Cli
|
||||
from tai_pkg.errors import BackendError
|
||||
from tai_pkg.muse import Muse
|
||||
from tai_pkg.opencode import Opencode
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["Codex"]
|
||||
|
||||
|
||||
class Codex:
|
||||
"""Codex."""
|
||||
|
||||
@staticmethod
|
||||
def codex_binary():
|
||||
return os.environ.get("CODEX_BIN") or shutil.which("codex")
|
||||
|
||||
@staticmethod
|
||||
def codex_jsonl_texts(stdout):
|
||||
texts = []
|
||||
last_error = ""
|
||||
for line in (stdout or "").splitlines():
|
||||
line = line.strip()
|
||||
if not line.startswith("{"):
|
||||
continue
|
||||
try:
|
||||
event = json.loads(line)
|
||||
except ValueError:
|
||||
continue
|
||||
if not isinstance(event, dict):
|
||||
continue
|
||||
etype = event.get("type")
|
||||
if etype == "item.completed":
|
||||
item = event.get("item") or {}
|
||||
if item.get("type") == "agent_message" and item.get("text"):
|
||||
texts.append(item["text"])
|
||||
elif etype in ("turn.failed", "error"):
|
||||
msg = event.get("message") or event.get("error") or ""
|
||||
if isinstance(msg, dict):
|
||||
msg = msg.get("message") or json.dumps(msg)
|
||||
if msg:
|
||||
last_error = str(msg)[:300]
|
||||
return texts, last_error
|
||||
|
||||
@staticmethod
|
||||
def codex_parse_output(stdout, outfile_text="", used_schema=False):
|
||||
if isinstance(outfile_text, str) and outfile_text.strip():
|
||||
if used_schema:
|
||||
try:
|
||||
answer = json.loads(outfile_text)
|
||||
except ValueError:
|
||||
answer = None
|
||||
if isinstance(answer, dict) and ("content" in answer or "tool_calls" in answer):
|
||||
return answer
|
||||
return Opencode.opencode_parse(outfile_text)
|
||||
texts, last_error = Codex.codex_jsonl_texts(stdout)
|
||||
if texts:
|
||||
return Opencode.opencode_parse("".join(texts))
|
||||
if last_error:
|
||||
raise BackendError("codex: %s" % last_error)
|
||||
return Opencode.opencode_parse(stdout)
|
||||
|
||||
@staticmethod
|
||||
def codex_run(prompt, timeout):
|
||||
binary = Codex.codex_binary()
|
||||
if not binary:
|
||||
raise BackendError("codex CLI not found (CODEX_BIN or PATH)")
|
||||
schema = json.dumps(Muse.muse_output_schema())
|
||||
with tempfile.TemporaryDirectory(prefix="tai-codex-") as folder:
|
||||
schema_path = os.path.join(folder, "schema.json")
|
||||
out_path = os.path.join(folder, "last.txt")
|
||||
with open(schema_path, "w", encoding="utf-8") as handle:
|
||||
handle.write(schema)
|
||||
base = [binary, "exec", "--json", "--ephemeral", "--sandbox", "read-only", "--skip-git-repo-check"]
|
||||
if len(prompt.encode("utf-8", "replace")) > 120000:
|
||||
variants = [
|
||||
(base + ["--output-schema", schema_path, "-o", out_path, "-"], prompt, True),
|
||||
(base + ["-o", out_path, "-"], prompt, False),
|
||||
]
|
||||
else:
|
||||
variants = [
|
||||
(base + ["--output-schema", schema_path, "-o", out_path, prompt], None, True),
|
||||
(base + ["-o", out_path, prompt], None, False),
|
||||
(base + [prompt], None, False),
|
||||
]
|
||||
last_error = ""
|
||||
for argv, prompt_input, used_schema in variants:
|
||||
try:
|
||||
done = subprocess.run(argv, input=prompt_input, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, cwd=folder, timeout=timeout)
|
||||
except subprocess.TimeoutExpired:
|
||||
raise BackendError("codex run timed out after %ds" % timeout)
|
||||
except OSError as exc:
|
||||
raise BackendError("codex run failed: %s" % Text.short_error(exc))
|
||||
if done.returncode != 0 and Cli._cli_unknown_flag(done.stderr) and (argv, prompt_input, used_schema) != variants[-1]:
|
||||
last_error = (done.stderr or "").strip()[:200]
|
||||
continue
|
||||
outfile_text = ""
|
||||
if "-o" in argv:
|
||||
try:
|
||||
with open(out_path, "r", encoding="utf-8", errors="replace") as handle:
|
||||
outfile_text = handle.read()
|
||||
except OSError:
|
||||
outfile_text = ""
|
||||
if done.returncode != 0 and not (outfile_text.strip() or (done.stdout or "").strip()):
|
||||
blob = ((done.stderr or "") + "\n" + (done.stdout or "")).strip()
|
||||
if Cli._cli_auth_failure(blob):
|
||||
raise BackendError("codex not authenticated (run codex login): %s" % blob[:200])
|
||||
raise BackendError("codex exit %d: %s" % (done.returncode, blob[:300] or last_error or "no output"))
|
||||
if not (outfile_text.strip() or (done.stdout or "").strip()):
|
||||
blob = (done.stderr or "").strip()
|
||||
raise BackendError("codex returned no output: %s" % (blob[:300] or last_error or "empty stdout"))
|
||||
return done.stdout or "", outfile_text, used_schema
|
||||
raise BackendError("codex failed: %s" % (last_error or "no usable output"))
|
||||
|
||||
@staticmethod
|
||||
def codex_health():
|
||||
binary = Codex.codex_binary()
|
||||
if not binary:
|
||||
return "missing: codex CLI not found"
|
||||
try:
|
||||
done = subprocess.run([binary, "login", "status"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=15)
|
||||
except (OSError, subprocess.SubprocessError) as exc:
|
||||
return "fail: %s" % Text.short_error(exc)[:60]
|
||||
detail = Cli._cli_version_detail(done, "codex")
|
||||
if done.returncode != 0:
|
||||
return "not logged in: %s" % detail
|
||||
return "ok, subscription (%s)" % detail
|
||||
@@ -0,0 +1,270 @@
|
||||
"""REPL slash-command completion and dispatch."""
|
||||
|
||||
import getpass
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
from tai_pkg import toolschema
|
||||
from tai_pkg.agent import Agent
|
||||
from tai_pkg.agents import REGISTRY
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.errors import BackendError
|
||||
from tai_pkg.models import ModelScoring
|
||||
from tai_pkg.render import Ansi
|
||||
from tai_pkg.render import Render
|
||||
from tai_pkg.sandbox import Sandbox
|
||||
from tai_pkg.scheduler import Scheduler
|
||||
from tai_pkg.skills import Skills
|
||||
from tai_pkg.sysinfo import SysInfo
|
||||
from tai_pkg.toolschema import ToolRegistry
|
||||
|
||||
__all__ = ["Commands"]
|
||||
|
||||
|
||||
class Commands:
|
||||
"""Commands."""
|
||||
|
||||
@staticmethod
|
||||
def complete_command(text, state):
|
||||
options = ["/" + name for name in Consts.COMMANDS if name.startswith(text.lstrip("/"))]
|
||||
return options[state] if state < len(options) else None
|
||||
|
||||
@staticmethod
|
||||
def handle_command(agent, text):
|
||||
parts = text[1:].split(None, 1)
|
||||
name = parts[0].lower()
|
||||
arg = parts[1].strip() if len(parts) > 1 else ""
|
||||
if name in ("quit", "exit"):
|
||||
return False
|
||||
if name == "help":
|
||||
print(Consts.HELP_TEXT)
|
||||
elif name == "profiles":
|
||||
names = agent.store.list_profiles()
|
||||
for item in names:
|
||||
print("%s %s" % ("*" if item == agent.profile else " ", item))
|
||||
if not names:
|
||||
print("(no profiles)")
|
||||
elif name == "bots":
|
||||
for item in agent.store.list_bots(agent.profile):
|
||||
nicks = " aka %s" % ", ".join(item["nicknames"]) if item["nicknames"] else ""
|
||||
print("%s %s%s" % ("*" if item["name"] == agent.bot else " ", item["name"], nicks))
|
||||
elif name == "bot":
|
||||
if not arg:
|
||||
print(Render.paint("use /bot <name>, see /bots", Ansi.RED))
|
||||
else:
|
||||
print(agent.switch_bot(arg))
|
||||
elif name == "profile":
|
||||
if not arg:
|
||||
print("current profile: %s (%d chars)" % (agent.profile, len(agent.system_message)))
|
||||
elif not Consts.PROFILE_RE.match(arg):
|
||||
print(Render.paint("invalid profile name", Ansi.RED))
|
||||
else:
|
||||
agent.switch_profile(arg)
|
||||
elif name == "env":
|
||||
if not arg:
|
||||
print("current environment: %s" % agent.env)
|
||||
elif arg not in ("home", "sandbox"):
|
||||
print(Render.paint("use /env home or /env sandbox", Ansi.RED))
|
||||
elif arg == "sandbox":
|
||||
try:
|
||||
Sandbox.ensure_box()
|
||||
except BackendError as exc:
|
||||
print(Render.paint("sandbox unavailable: %s" % exc, Ansi.RED))
|
||||
else:
|
||||
agent.env = "sandbox"
|
||||
print("environment: sandbox (isolated container)")
|
||||
else:
|
||||
agent.env = "home"
|
||||
print("environment: home")
|
||||
elif name == "skills":
|
||||
bits = arg.split()
|
||||
if bits and bits[0] == "install":
|
||||
if len(bits) < 2:
|
||||
print(Render.paint("use /skills install <source> [name=x] [scope=project|home] [mode=copy|symlink] [force]", Ansi.RED))
|
||||
return True
|
||||
call_args = {"source": bits[1]}
|
||||
for token in bits[2:]:
|
||||
if token == "force":
|
||||
call_args["force"] = True
|
||||
elif "=" in token:
|
||||
key, _, value = token.partition("=")
|
||||
if key in ("name", "scope", "mode"):
|
||||
call_args[key] = value
|
||||
print(agent.tools.dispatch("install_skill", json.dumps(call_args)))
|
||||
return True
|
||||
agent.skills = Skills.discover_skills(agent.config.home, os.getcwd())
|
||||
agent.apply_system()
|
||||
if not agent.skills:
|
||||
print("(no skills)")
|
||||
for skill_name in sorted(agent.skills):
|
||||
skill = agent.skills[skill_name]
|
||||
print("- %s: %s %s" % (skill_name, skill["description"][:200], Skills.skill_provenance_note(skill)))
|
||||
wanted = sorted(item for item in Consts.SKILL_BLUEPRINTS if item not in agent.skills)
|
||||
if wanted:
|
||||
print("blueprints, built on first load:")
|
||||
for skill_name in wanted:
|
||||
print("- %s: %s" % (skill_name, Consts.SKILL_BLUEPRINTS[skill_name]["hint"]))
|
||||
elif name == "sysinfo":
|
||||
print(SysInfo.collect_sysinfo())
|
||||
elif name == "models":
|
||||
print(ModelScoring.models_report(agent.chat))
|
||||
elif name == "secret":
|
||||
sub = arg.split(None, 1)
|
||||
action = sub[0].lower() if sub else ""
|
||||
rest = sub[1].strip() if len(sub) > 1 else ""
|
||||
if action == "set" and rest:
|
||||
if not Consts.SECRET_RE.match(rest):
|
||||
print(Render.paint("invalid secret name, use letters, digits, dot, underscore, hyphen", Ansi.RED))
|
||||
else:
|
||||
try:
|
||||
value = getpass.getpass("value for '%s': " % rest)
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
print("\naborted")
|
||||
return True
|
||||
if not value:
|
||||
print("empty value, aborted")
|
||||
else:
|
||||
agent.store.save_secret(rest, value, None, None, agent.profile)
|
||||
print("stored secret '%s'" % rest)
|
||||
elif action in ("list", "") and not rest:
|
||||
print(agent.tools.dispatch("list_secrets", "{}"))
|
||||
elif action == "delete" and rest:
|
||||
if agent.store.load_secret(rest, agent.profile) is None:
|
||||
print(Render.paint("unknown secret '%s'" % rest, Ansi.RED))
|
||||
elif agent.ask_approval("delete secret '%s'" % rest, guidance=False):
|
||||
agent.store.delete_secret(rest, agent.profile)
|
||||
print("deleted secret '%s'" % rest)
|
||||
else:
|
||||
print("aborted")
|
||||
else:
|
||||
print(Render.paint("use /secret set <name> or /secret list or /secret delete <name>", Ansi.RED))
|
||||
elif name == "schedule":
|
||||
parts = arg.split(None, 2)
|
||||
if len(parts) != 3 or parts[0] not in ("at", "every"):
|
||||
print(Render.paint("use /schedule at <ISO datetime> <prompt> or /schedule every <30s|10m|2h|1d> <prompt>", Ansi.RED))
|
||||
else:
|
||||
mode, when, prompt = parts
|
||||
if mode == "at":
|
||||
try:
|
||||
next_run = Scheduler.parse_schedule_at(when)
|
||||
except ValueError as exc:
|
||||
print(Render.paint("error: %s" % exc, Ansi.RED))
|
||||
return True
|
||||
every_sec = 0
|
||||
else:
|
||||
every_sec = Scheduler.parse_duration(when)
|
||||
if every_sec is None or every_sec < 60:
|
||||
print(Render.paint("error: every must be a duration of at least 60s", Ansi.RED))
|
||||
return True
|
||||
next_run = (datetime.now(timezone.utc) + timedelta(seconds=every_sec)).isoformat()
|
||||
row_id = agent.store.add_schedule("", prompt, agent.profile, every_sec, next_run, 600)
|
||||
print("scheduled #%d" % row_id)
|
||||
elif name == "schedules":
|
||||
print(Scheduler.format_schedule_lines(agent.store.list_schedules(agent.profile), datetime.now(timezone.utc)))
|
||||
elif name == "unschedule":
|
||||
try:
|
||||
row_id = int(arg or "0")
|
||||
except ValueError:
|
||||
print(Render.paint("use /unschedule <id>", Ansi.RED))
|
||||
return True
|
||||
if agent.store.remove_schedule(row_id, agent.profile):
|
||||
print("deleted schedule #%d" % row_id)
|
||||
else:
|
||||
print(Render.paint("no schedule #%d" % row_id, Ansi.RED))
|
||||
elif name == "records":
|
||||
print(agent.tools.dispatch("record_search", json.dumps({"query": arg})))
|
||||
elif name == "record":
|
||||
bits = arg.split()
|
||||
if not bits:
|
||||
print(Render.paint("use /record <mem:id> [offset]", Ansi.RED))
|
||||
else:
|
||||
try:
|
||||
offset = int(bits[1]) if len(bits) > 1 else 0
|
||||
except ValueError:
|
||||
print(Render.paint("offset must be a number", Ansi.RED))
|
||||
return True
|
||||
print(agent.tools.dispatch("record_read", json.dumps({"id": bits[0], "offset": offset})))
|
||||
elif name == "graph":
|
||||
if not arg:
|
||||
print(Render.paint("use /graph <mem:id|secret:name|sched:id>", Ansi.RED))
|
||||
else:
|
||||
print(agent.tools.dispatch("graph_query", json.dumps({"node": arg})))
|
||||
elif name == "tools":
|
||||
if arg:
|
||||
match = [schema["function"] for schema in toolschema.TOOL_SCHEMAS if schema["function"]["name"] == arg]
|
||||
if not match:
|
||||
print(Render.paint("unknown tool '%s'" % arg, Ansi.RED))
|
||||
return True
|
||||
print("%s [%s]" % (match[0]["name"], "core" if match[0]["name"] in Consts.CORE_TOOLS else "lazy"))
|
||||
print(match[0]["description"])
|
||||
print("trigger tags: %s" % ", ".join(Consts.TOOL_TAGS.get(match[0]["name"], ())))
|
||||
return True
|
||||
names = [schema["function"]["name"] for schema in toolschema.TOOL_SCHEMAS if ToolRegistry.tool_env_ok(schema["function"]["name"])]
|
||||
cores = sorted(tool for tool in names if tool in Consts.CORE_TOOLS)
|
||||
lazy = sorted(tool for tool in names if tool not in Consts.CORE_TOOLS)
|
||||
print("core, always loaded:\n %s" % "\n ".join(cores))
|
||||
print("lazy, loads when named:\n %s" % "\n ".join(lazy))
|
||||
elif name == "search":
|
||||
if not arg:
|
||||
print(Render.paint("use /search <query>", Ansi.RED))
|
||||
else:
|
||||
print(agent.tools.dispatch("search", json.dumps({"query": arg})))
|
||||
elif name == "install":
|
||||
bits = arg.split()
|
||||
action = bits[0] if bits else "status"
|
||||
if action not in ("status", "install", "upgrade", "reinstall", "uninstall"):
|
||||
print(Render.paint("use /install [status|install|upgrade|reinstall|uninstall] [targets...]", Ansi.RED))
|
||||
else:
|
||||
print(agent.tools.dispatch("install", json.dumps({"action": action, "targets": bits[1:] or list(Consts.INSTALL_TARGETS)})))
|
||||
elif name == "tags":
|
||||
print(agent.tools.dispatch("tags", json.dumps({"prefix": arg} if arg else {})))
|
||||
elif name == "audit":
|
||||
print(agent.tools.dispatch("audit", json.dumps({"path": arg} if arg else {})))
|
||||
elif name == "restore":
|
||||
try:
|
||||
row_id = int(arg or "0")
|
||||
except ValueError:
|
||||
print(Render.paint("use /restore <audit id>", Ansi.RED))
|
||||
return True
|
||||
print(agent.tools.dispatch("restore", json.dumps({"id": row_id})))
|
||||
agent.show_diff()
|
||||
elif name == "release":
|
||||
bits = arg.split(None, 1)
|
||||
if len(bits) != 2 or bits[0] not in ("major", "minor", "patch"):
|
||||
print(Render.paint("use /release <major|minor|patch> <message>", Ansi.RED))
|
||||
else:
|
||||
print(agent.tools.dispatch("release", json.dumps({"part": bits[0], "message": bits[1]})))
|
||||
elif name == "fork":
|
||||
if not arg:
|
||||
print(Render.paint("use /fork <task>", Ansi.RED))
|
||||
else:
|
||||
agent_id = REGISTRY.spawn_agent(arg, agent.profile, 600, agent.config, agent.store.seal, agent.depth + 1, agent.runner_override)
|
||||
print("agent %d started, REPL stays free" % agent_id)
|
||||
elif name == "agents":
|
||||
records = REGISTRY.list_agents(agent.profile)
|
||||
if not records:
|
||||
print("(no agents)")
|
||||
for record in records:
|
||||
print("#%d [%s] %ds %s" % (record["id"], record["status"], record["elapsed"], record["task"][:60]))
|
||||
elif name == "agent":
|
||||
if arg == "clear":
|
||||
print("purged %d finished agents" % REGISTRY.clear_agents(agent.profile))
|
||||
else:
|
||||
try:
|
||||
agent_id = int(arg)
|
||||
except ValueError:
|
||||
print(Render.paint("use /agent <id> or /agent clear", Ansi.RED))
|
||||
return True
|
||||
status, text = REGISTRY.poll_agent(agent_id, 0, agent.profile)
|
||||
print("[%s]\n%s" % (status, text))
|
||||
elif name == "compact":
|
||||
agent.messages, changed = Agent.compact_messages(agent.messages, agent.chat)
|
||||
agent.apply_system()
|
||||
print("compacted" if changed else "nothing to compact")
|
||||
elif name == "clear":
|
||||
agent.reset_history()
|
||||
print("history cleared")
|
||||
else:
|
||||
print(Render.paint("unknown command, try /help", Ansi.RED))
|
||||
return True
|
||||
@@ -0,0 +1,24 @@
|
||||
"""Optional-dependency import guards (readline, fcntl, pty set)."""
|
||||
|
||||
__all__ = ["readline", "fcntl", "pty", "select", "termios", "tty"]
|
||||
|
||||
|
||||
try:
|
||||
import readline
|
||||
except ImportError:
|
||||
readline = None
|
||||
|
||||
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError:
|
||||
fcntl = None
|
||||
|
||||
|
||||
try:
|
||||
import pty
|
||||
import select
|
||||
import termios
|
||||
import tty
|
||||
except ImportError:
|
||||
pty = select = termios = tty = None
|
||||
@@ -0,0 +1,61 @@
|
||||
"""Runtime configuration from environment and arguments."""
|
||||
|
||||
import os
|
||||
import re
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
|
||||
__all__ = ["Config"]
|
||||
|
||||
|
||||
class Config:
|
||||
def __init__(self, args):
|
||||
self.home = os.environ.get("TAI_HOME") or os.path.join(os.path.expanduser("~"), ".tai")
|
||||
self.profiles_dir = os.path.join(self.home, "profiles")
|
||||
self.audio_dir = os.path.join(self.home, "audio")
|
||||
self.db_path = os.path.join(self.home, "memory.db")
|
||||
self.history_path = os.path.join(self.home, "history")
|
||||
self.model = os.environ.get("TAI_MODEL") or Consts.DEFAULT_MODEL
|
||||
self.explicit_model = os.environ.get("TAI_MODEL") or ""
|
||||
self.backends = Config.parse_backends(os.environ.get("TAI_BACKENDS"))
|
||||
self.openrouter_key = os.environ.get("OPENROUTER_API_KEY") or ""
|
||||
self.openrouter_models = Config.split_models(os.environ.get("TAI_OPENROUTER_MODELS"))
|
||||
self.retry_seconds = float(os.environ.get("TAI_RETRY_SECONDS") or 0) if re.match(r"^\d+(\.\d+)?$", (os.environ.get("TAI_RETRY_SECONDS") or "0").strip()) else 0.0
|
||||
self.opencode_enabled = (os.environ.get("TAI_OPENCODE") or "1").strip().lower() not in ("0", "false", "no", "off")
|
||||
self.opencode_models = Config.split_models(os.environ.get("TAI_OPENCODE_MODELS"))
|
||||
self.devplace_key = os.environ.get("DEVPLACE_API_KEY") or ""
|
||||
self.devplace_base = (os.environ.get("TAI_DEVPLACE_BASE") or Consts.DEVPLACE_BASE).rstrip("/")
|
||||
self.devplace_models = Config.split_models(os.environ.get("TAI_DEVPLACE_MODELS"))
|
||||
self.openrouter_enabled = (os.environ.get("TAI_OPENROUTER") or "1").strip().lower() not in ("0", "false", "no", "off")
|
||||
self.pollinations_enabled = (os.environ.get("TAI_POLLINATIONS") or "1").strip().lower() not in ("0", "false", "no", "off")
|
||||
self.pollinations_models = Config.split_models(os.environ.get("TAI_POLLINATIONS_MODELS")) or list(Consts.POLLINATIONS_MODELS)
|
||||
self.zen_enabled = (os.environ.get("TAI_ZEN") or "").strip().lower() in ("1", "true", "yes", "on")
|
||||
self.zen_models = Config.split_models(os.environ.get("TAI_ZEN_MODELS")) or list(Consts.ZEN_FREE_MODELS)
|
||||
self.voice = os.environ.get("TAI_VOICE") or Consts.DEFAULT_VOICE
|
||||
self.auto_approve = bool(args.yes)
|
||||
self.profile = args.profile or "default"
|
||||
os.makedirs(self.profiles_dir, exist_ok=True)
|
||||
os.makedirs(self.audio_dir, exist_ok=True)
|
||||
|
||||
@staticmethod
|
||||
def split_models(spec):
|
||||
return [item.strip() for item in (spec or "").replace("\n", ",").split(",") if item.strip()]
|
||||
|
||||
@staticmethod
|
||||
def parse_backends(spec):
|
||||
found = []
|
||||
for entry in (spec or "").replace("\n", ";").split(";"):
|
||||
fields = [field.strip() for field in entry.split("|")]
|
||||
if len(fields) < 3:
|
||||
continue
|
||||
label, base, key = fields[0], fields[1].rstrip("/"), fields[2]
|
||||
if key.startswith("$"):
|
||||
key = os.environ.get(key[1:], "")
|
||||
model = fields[3].strip() if len(fields) > 3 else ""
|
||||
profile = fields[4].strip().lower() if len(fields) > 4 else ""
|
||||
if profile not in ("", "opencode"):
|
||||
profile = ""
|
||||
if not label or not base:
|
||||
continue
|
||||
found.append({"label": label, "base": base, "key": key, "model": model or "", "profile": profile})
|
||||
return found
|
||||
@@ -0,0 +1,600 @@
|
||||
"""Single source of all literal configuration and shared constants."""
|
||||
|
||||
import os
|
||||
import re
|
||||
|
||||
__all__ = ["Consts"]
|
||||
|
||||
|
||||
class Consts:
|
||||
"""All literal configuration and shared constants."""
|
||||
VERSION = "1.22.0"
|
||||
|
||||
WORKER_STEPS = 12
|
||||
|
||||
CREATE_SKILL_STEPS = 40
|
||||
|
||||
FORK_MAX_DEPTH = 2
|
||||
|
||||
MUSE_LABEL = "muse"
|
||||
|
||||
MUSE_CHAT_URL = "muse://exec"
|
||||
|
||||
MUSE_MODELS_URL = "muse://models"
|
||||
|
||||
GROK_LABEL = "grok"
|
||||
|
||||
GROK_CHAT_URL = "grok://exec"
|
||||
|
||||
GROK_MODELS_URL = "grok://models"
|
||||
|
||||
CLAUDE_LABEL = "claude"
|
||||
|
||||
CLAUDE_CHAT_URL = "claude://exec"
|
||||
|
||||
CLAUDE_MODELS_URL = "claude://models"
|
||||
|
||||
CODEX_LABEL = "codex"
|
||||
|
||||
CODEX_CHAT_URL = "codex://exec"
|
||||
|
||||
CODEX_MODELS_URL = "codex://models"
|
||||
|
||||
GEMINI_LABEL = "gemini"
|
||||
|
||||
GEMINI_CHAT_URL = "gemini://exec"
|
||||
|
||||
GEMINI_MODELS_URL = "gemini://models"
|
||||
|
||||
OLLAMA_LABEL = "ollama"
|
||||
|
||||
OLLAMA_BASE = "http://localhost:11434/v1"
|
||||
|
||||
OLLAMA_LIST_TTL = 600
|
||||
|
||||
CLI_CHAT_URLS = (MUSE_CHAT_URL, GROK_CHAT_URL, CLAUDE_CHAT_URL, CODEX_CHAT_URL, GEMINI_CHAT_URL)
|
||||
|
||||
MUSE_MAX_STEPS = 1
|
||||
|
||||
MUSE_SESSIONS_FILE = "muse_sessions.json"
|
||||
|
||||
MUSE_LOCKS_DIR = "muse_locks"
|
||||
|
||||
MUSE_WORKER_TTL = 86400
|
||||
|
||||
MUSE_REAP_GRACE = 30
|
||||
|
||||
OPENROUTER_LABEL = "openrouter"
|
||||
|
||||
OPENROUTER_BASE = "https://openrouter.ai/api/v1"
|
||||
|
||||
OPENROUTER_MODEL = "deepseek/deepseek-v4.1-flash"
|
||||
|
||||
OPENCODE_LABEL = "opencode"
|
||||
|
||||
OPENCODE_CHAT_URL = "opencode://run"
|
||||
|
||||
OPENCODE_MODELS_URL = "opencode://models"
|
||||
|
||||
OPENCODE_AGENT = "plan"
|
||||
|
||||
OPENCODE_PARSE_MAX_SCANS = 200
|
||||
|
||||
DEVPLACE_LABEL = "devplace"
|
||||
|
||||
DEVPLACE_BASE = "https://devplace.net/openai/v1"
|
||||
|
||||
DEVPLACE_FREE_ROUTES = ("free", "devii")
|
||||
|
||||
ZEN_LABEL = "zen"
|
||||
|
||||
ZEN_BASE = "https://opencode.ai/zen/v1"
|
||||
|
||||
ZEN_FREE_MODELS = ("mimo-v2.6-flash-free", "nemotron-3.5-lightning-free", "nemotron-3-ultra-free", "muse-spark-1.3-contributor-free", "muse-spark-1.2-contributor-free", "jev-1.13-free", "exo-free", "space-bunny-free", "longcat-2.5-preview-free", "ling-3.1-flash-free", "ling-3.0-flash-fin-free", "fledge-alpha-free")
|
||||
|
||||
POLLINATIONS_LABEL = "pollinations"
|
||||
|
||||
POLLINATIONS_BASE = "https://text.pollinations.ai/openai/v1"
|
||||
|
||||
POLLINATIONS_MODELS = ("openai",)
|
||||
|
||||
OPENCODE_USER_AGENT = "opencode/1.18.29 ai-sdk/provider-utils/4.0.23 runtime/bun/1.3.15"
|
||||
|
||||
OPENCODE_CLIENT_NAME = "cli"
|
||||
|
||||
MODEL_HEALTH_FILE = "model_health.json"
|
||||
|
||||
MODEL_NEUTRAL_REWARD = 0.5
|
||||
|
||||
MODEL_WEIGHT_PRIOR = 2.0
|
||||
|
||||
MODEL_SPEED_REF_TPS = 25.0
|
||||
|
||||
MODEL_LATENCY_REF_MS = 8000.0
|
||||
|
||||
MODEL_CIRCUIT_THRESHOLD = 3
|
||||
|
||||
MODEL_CIRCUIT_COOLDOWN = 300.0
|
||||
|
||||
RETRY_BACKOFF = (2, 5, 10, 20, 30, 60)
|
||||
|
||||
MODEL_LIST_TIMEOUT = 8
|
||||
|
||||
OPENROUTER_FREE_CACHE_TTL = 3600
|
||||
|
||||
RSEARCH_BASE = "https://rsearch.app.molodetz.nl"
|
||||
|
||||
FETCH_METHODS = ("GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS")
|
||||
|
||||
CURL_HINT = "hint: use the web_fetch tool for HTTP(S) (methods, headers, bodies, status) instead of curl/wget via shell"
|
||||
|
||||
DEFAULT_MODEL = "muse"
|
||||
|
||||
DEFAULT_VOICE = "en-US-EmmaMultilingualNeural"
|
||||
|
||||
DEFAULT_PASSPHRASE = "tai-default-insecure-change-me"
|
||||
|
||||
BOX_IMAGE = "tai-box:latest"
|
||||
|
||||
BOX_NAME = "tai-box"
|
||||
|
||||
TELEGRAM_API = "https://api.telegram.org/bot"
|
||||
|
||||
TELEGRAM_FILE_API = "https://api.telegram.org/file/bot"
|
||||
|
||||
SKILL_NAME_RE = re.compile(r"^[a-z0-9](?:[a-z0-9-]{0,62}[a-z0-9])?$")
|
||||
|
||||
BASHRC_MARK_BEGIN = "# tai command-not-found hook - start"
|
||||
|
||||
BASHRC_MARK_END = "# tai command-not-found hook - end"
|
||||
|
||||
BASHRC_BLOCK = (
|
||||
BASHRC_MARK_BEGIN + "\n"
|
||||
"alias tai=\"$HOME/.local/bin/tai.py\"\n"
|
||||
"command_not_found_handle() {\n"
|
||||
" \"$HOME/.local/bin/tai.py\" \"$@\"\n"
|
||||
" return $?\n"
|
||||
"}\n" + BASHRC_MARK_END + "\n"
|
||||
)
|
||||
|
||||
ZSH_NOTFOUND_BLOCK = (
|
||||
BASHRC_MARK_BEGIN + "\n"
|
||||
"alias tai=\"$HOME/.local/bin/tai.py\"\n"
|
||||
"command_not_found_handler() {\n"
|
||||
" \"$HOME/.local/bin/tai.py\" \"$@\"\n"
|
||||
" return $?\n"
|
||||
"}\n" + BASHRC_MARK_END + "\n"
|
||||
)
|
||||
|
||||
FISH_NOTFOUND_BLOCK = (
|
||||
BASHRC_MARK_BEGIN + "\n"
|
||||
"function tai\n"
|
||||
" \"$HOME/.local/bin/tai.py\" $argv\n"
|
||||
"end\n"
|
||||
"function fish_command_not_found\n"
|
||||
" \"$HOME/.local/bin/tai.py\" $argv\n"
|
||||
"end\n" + BASHRC_MARK_END + "\n"
|
||||
)
|
||||
|
||||
BOX_CONTAINERFILE = """FROM python:3.12-slim
|
||||
ENV DEBIAN_FRONTEND=noninteractive PIP_NO_CACHE_DIR=1
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends ffmpeg espeak-ng python3-venv && rm -rf /var/lib/apt/lists/*
|
||||
RUN python -m venv /box/venv && /box/venv/bin/pip install --no-cache-dir faster-whisper edge-tts
|
||||
RUN /box/venv/bin/python -c "from faster_whisper import WhisperModel; WhisperModel('tiny', device='cpu', compute_type='int8')"
|
||||
COPY stt.py /box/stt.py
|
||||
COPY tts.py /box/tts.py
|
||||
CMD ["sleep", "infinity"]
|
||||
"""
|
||||
|
||||
BOX_PYTHON = "/box/venv/bin/python"
|
||||
|
||||
BOX_STT = """import sys
|
||||
from faster_whisper import WhisperModel
|
||||
data = sys.stdin.buffer.read()
|
||||
with open("/tmp/in.audio", "wb") as handle:
|
||||
handle.write(data)
|
||||
model = WhisperModel("tiny", device="cpu", compute_type="int8")
|
||||
segments, _info = model.transcribe("/tmp/in.audio")
|
||||
print(" ".join(segment.text for segment in segments).strip(), flush=True)
|
||||
"""
|
||||
|
||||
BOX_TTS = """import asyncio
|
||||
import sys
|
||||
import edge_tts
|
||||
async def speak(text, voice):
|
||||
talk = edge_tts.Communicate(text, voice or "en-US-EmmaMultilingualNeural")
|
||||
await talk.save("/tmp/out.mp3")
|
||||
text = sys.stdin.read().strip()
|
||||
voice = sys.argv[1] if len(sys.argv) > 1 else ""
|
||||
asyncio.run(speak(text, voice))
|
||||
sys.stdout.buffer.write(open("/tmp/out.mp3", "rb").read())
|
||||
"""
|
||||
|
||||
TELEGRAM_UNIT = """[Unit]
|
||||
Description=tai telegram bot
|
||||
After=network-online.target
|
||||
Wants=network-online.target
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
WorkingDirectory=%%h
|
||||
ExecStart=%s --telegram
|
||||
EnvironmentFile=%s
|
||||
Restart=on-failure
|
||||
RestartSec=10
|
||||
|
||||
[Install]
|
||||
WantedBy=default.target
|
||||
"""
|
||||
|
||||
SCHEDULER_UNIT = """[Unit]
|
||||
Description=tai scheduler
|
||||
After=network-online.target
|
||||
Wants=network-online.target
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
WorkingDirectory=%%h
|
||||
ExecStart=%s --scheduler
|
||||
Restart=on-failure
|
||||
RestartSec=10
|
||||
|
||||
[Install]
|
||||
WantedBy=default.target
|
||||
"""
|
||||
|
||||
INSTALL_TARGETS = ("binary", "bash-hook", "zsh-hook", "fish-hook", "plexer-hook", "zsh-plexer-hook", "fish-plexer-hook", "venv", "scheduler-service", "telegram-service", "container")
|
||||
|
||||
CONTEXT_CAP = 32000
|
||||
|
||||
COMPACT_RATIO = 0.8
|
||||
|
||||
KEEP_TURNS = 6
|
||||
|
||||
MAX_STEPS = 25
|
||||
|
||||
TOTAL_STEP_CAP = 500
|
||||
|
||||
LOOP_NUDGE_AT = 3
|
||||
|
||||
LOOP_STOP_AT = 5
|
||||
|
||||
SYSTEM_MAX_CHARS = 12000
|
||||
|
||||
SHELL_TIMEOUT = 120
|
||||
|
||||
SCHEDULER_INTERVAL = 30
|
||||
|
||||
HTTP_TIMEOUT = 60
|
||||
|
||||
STREAM_TIMEOUT = 300
|
||||
|
||||
EDGE_HOST = "speech.platform.bing.com"
|
||||
|
||||
EDGE_WS_PATH = "/consumer/speech/synthesize/readaloud/edge/v1"
|
||||
|
||||
EDGE_TRUSTED_TOKEN = "6A5AA1D4EAFF4E9FB37E23D68491D6F4"
|
||||
|
||||
EDGE_CHROMIUM = "143.0.3650.75"
|
||||
|
||||
PROFILE_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_-]{0,31}$")
|
||||
|
||||
BOT_RE = re.compile(r"^[a-z0-9][a-z0-9_-]{0,31}$")
|
||||
|
||||
BOT_MENTION_RE = re.compile(r"^@([A-Za-z0-9][A-Za-z0-9_-]{0,31})\s+(.*)$", re.DOTALL)
|
||||
|
||||
SECRET_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,63}$")
|
||||
|
||||
LEGACY_PASSWORD_NOTE = "especially passwords (collect and keep every password the user shares), preferences, or standing behavior changes"
|
||||
|
||||
VAULT_MEMORY_NOTE = "preferences, or standing behavior changes. Store passwords, tokens, and secrets with store_secret and reference them by name; never write secret values into memory"
|
||||
|
||||
SAFE_COMMANDS = ("ls", "pwd", "echo", "cat", "head", "tail", "grep", "find", "wc", "sort", "uniq", "diff", "file", "stat", "date", "whoami", "uname", "lsb_release")
|
||||
|
||||
RECORDERS = (("arecord", ("arecord", "-q", "-d", "{seconds}", "-f", "cd", "-t", "wav", "{path}")), ("rec", ("rec", "-q", "{path}", "trim", "0", "{seconds}")), ("ffmpeg", ("ffmpeg", "-y", "-v", "quiet", "-f", "alsa", "-i", "default", "-t", "{seconds}", "{path}")))
|
||||
|
||||
TRANSCRIBERS = ("whisper-cpp", "whisper-cli", "whisper", "faster-whisper")
|
||||
|
||||
SYSINFO_TIMEOUT = 10
|
||||
|
||||
SYSINFO_BINARIES = ("tmux", "git", "ffmpeg", "arecord", "sox", "whisper-cpp", "whisper", "systemctl", "podman", "docker", "muse", "grok", "opencode", "claude", "codex", "gemini", "ollama")
|
||||
|
||||
DEFAULT_SYSTEM = (
|
||||
"You are tai, a professional autonomous assistant. You act with tools, verify results, and report concisely. "
|
||||
"Rules: prefer the smallest change that solves the task; never silently ignore errors; ask when requirements are ambiguous; "
|
||||
"keep answers short and factual. Tools: shell and file tools for the local machine, web_search and web_fetch for the internet, "
|
||||
"web_fetch handles every HTTP method with headers and bodies so never curl or wget through shell, "
|
||||
"speak for voice output, recall to search past session memory. Memory: call remember whenever you learn durable facts, "
|
||||
"preferences, or standing behavior changes. Store passwords, tokens, and secrets with store_secret and reference them by name; "
|
||||
"never write secret values into memory. Name and value suffice for secrets; never pester the user for more. Call remember with a forget instruction to drop outdated knowledge. Delegation: do quick, interactive, or memory-changing work yourself; fork background subagents for independent, long, or context-heavy work and collect them with poll; schedule future work instead of waiting on it. "
|
||||
"Records: save big or durable findings with record_save and tag generously; large tool outputs spill to mem:<hash> ids automatically, page them with record_read, connect related nodes with graph_link. "
|
||||
"Tags: strict rules, lowercase singular only, shortest common word for the subject, one canonical tag per subject, at most 5 per item, always reuse existing tags from the tags tool instead of inventing synonyms. Content words matching known tags attach automatically, and new records link to recent same-tag records. "
|
||||
"Lazy tools: only everyday tools load by default; name any other tool or topic in your reply text and it becomes callable on your next step. "
|
||||
"Bots: one identity holds many bots; a bot is only a system message plus its own history, all vault data is shared. @name routes one turn to another bot and files the exchange in both histories; /bot switches, create_bot makes bots. "
|
||||
"Terminal: the user's question is very likely about what is visible in their terminal, especially right after a terminal capture or when launched with a single question straight from a shell; resolve vague references (this, that, it, the error) against that content, capture it with get_current_terminal_content instead of asking, and answer from what you see. "
|
||||
"Retrieval: use search first for anything stored: one ranked full-text query across records, events, and audit plus graph expansion beats paging and guessing. Chain tool calls until the task is done."
|
||||
)
|
||||
|
||||
MERGE_SYSTEM = (
|
||||
"You maintain an AI assistant system prompt. You receive the current system prompt and one memory instruction. "
|
||||
"Rewrite the system prompt to incorporate the instruction: add new facts, update changed behavior, or remove forgotten items. "
|
||||
"Never persist passwords, tokens, or other secret values; if the instruction holds one, note only that a secret exists and must be stored with store_secret. "
|
||||
"Preserve everything unrelated. Keep it organized with short sections. Output ONLY the rewritten system prompt, no explanation, no code fences."
|
||||
)
|
||||
|
||||
COMPACT_SYSTEM = (
|
||||
"You compress assistant session history into a dense handoff summary. Cover: goal, completed work with outcomes, current state, "
|
||||
"key decisions, files touched, errors and fixes, pending next steps. Output only the summary, no preamble."
|
||||
)
|
||||
|
||||
AUTONOMOUS_NOTE = (
|
||||
"Autonomous mode: never ask the user anything this session. When requirements are unclear, blocked, or need input, "
|
||||
"research deeply with web_search, web_fetch, shell, and subagents, then decide and proceed. State assumptions briefly in your reply."
|
||||
)
|
||||
|
||||
HELP_TEXT = (
|
||||
"commands:\n"
|
||||
" /profile [name] show current profile or switch to it\n"
|
||||
" /profiles list all profiles, current marked with *\n"
|
||||
" /bots list bots, current marked with *\n"
|
||||
" /bot <name> switch bot, history resumes\n"
|
||||
" /env [home|sandbox] show or switch execution environment\n"
|
||||
" /skills list loaded skills with provenance\n"
|
||||
" /skills install ... install an existing skill (copy or symlink)\n"
|
||||
" /secret manage sealed secrets (set|list|delete)\n"
|
||||
" /sysinfo show host environment checks\n"
|
||||
" /models show model roster with speed health\n"
|
||||
" /schedule at|every schedule a prompt for later\n"
|
||||
" /schedules list scheduled prompts\n"
|
||||
" /unschedule <id> delete a scheduled prompt\n"
|
||||
" /records [query] search saved records\n"
|
||||
" /record <id> read one record page\n"
|
||||
" /graph <node> show node neighborhood\n"
|
||||
" /tags [prefix] list tags with usage counts\n"
|
||||
" /tools [name] show core and lazy tools\n"
|
||||
" /install ... install status, install, upgrade, reinstall\n"
|
||||
" /search <query> ranked search over everything\n"
|
||||
" /audit [path] show file audit trail\n"
|
||||
" /restore <id> restore a file from an audit row\n"
|
||||
" /release part msg bump version, back up, log message\n"
|
||||
" /fork <task> spawn a background subagent, REPL stays free\n"
|
||||
" /agents list background subagents\n"
|
||||
" /agent <id|clear> show one result, or purge finished agents\n"
|
||||
" /compact compress history into a summary\n"
|
||||
" /clear drop history, keep system message\n"
|
||||
" /help this text\n"
|
||||
" /quit exit\n"
|
||||
"anything else is sent to the agent."
|
||||
)
|
||||
|
||||
SPILL_LIMIT = 6000
|
||||
|
||||
RECORD_KINDS = ("note", "output", "file", "research", "transcript", "termlog")
|
||||
|
||||
AUDIT_MAX_CHARS = 20000
|
||||
|
||||
FILE_RECORD_MAX = 50000
|
||||
|
||||
SHELL_SNAPSHOT_MAX_FILES = 25
|
||||
|
||||
BACKUP_KEEP = 10
|
||||
|
||||
_UNSET = object()
|
||||
|
||||
SHELL_SEGMENT = (";", "&&", "||", "|", "&", "(", ")")
|
||||
|
||||
SHELL_WRAPPERS = ("sudo", "command", "env", "nice", "nohup", "setsid")
|
||||
|
||||
CHUNK_MAX_LINES = 60
|
||||
|
||||
CHUNK_OVERLAP_LINES = 3
|
||||
|
||||
CHUNK_MAX_CHARS = 4000
|
||||
|
||||
VERSION_RE = re.compile(r"^(?:[ \t]*)VERSION = \"(\d+)\.(\d+)\.(\d+)\"$", re.MULTILINE)
|
||||
|
||||
STREAM_HEIGHT = 4
|
||||
|
||||
DIFF_CONTEXT = 3
|
||||
|
||||
DIFF_MAX_LINES = 120
|
||||
|
||||
BG_ADD = "\033[48;5;28m"
|
||||
|
||||
BG_DEL = "\033[48;5;88m"
|
||||
|
||||
FG_ORANGE = "\033[38;5;214m"
|
||||
|
||||
PY_KEYWORDS = ("False", "None", "True", "and", "as", "assert", "async", "await", "break", "class", "continue", "def", "del", "elif", "else", "except", "finally", "for", "from", "global", "if", "import", "in", "is", "lambda", "nonlocal", "not", "or", "pass", "raise", "return", "try", "while", "with", "yield", "match", "case")
|
||||
|
||||
PY_TOKEN_RE = re.compile(r"(?P<string>\"\"\"(?:\\.|[^\\])*?\"\"\"|'''(?:\\.|[^\\])*?'''|\"(?:\\.|[^\"\\\n])*\"|'(?:\\.|[^'\\\n])*')|(?P<comment>#[^\n]*)|(?P<number>\b\d[\d._]*(?:[eE][+-]?\d+)?[jJ]?\b)|(?P<keyword>\b(?:%s)\b)|(?P<decorator>@[A-Za-z_][\w.]*)|(?P<defname>(?<=def )[A-Za-z_]\w*|(?<=class )[A-Za-z_]\w*)" % "|".join(PY_KEYWORDS))
|
||||
|
||||
ANSI_RE = re.compile(r"\033\[[0-9;]*m")
|
||||
|
||||
MD_BULLETS = ("•", "◦", "▪")
|
||||
|
||||
EXTERNAL_SKILL_DIRNAMES = (os.path.join(".claude", "skills"), os.path.join(".agents", "skills"))
|
||||
|
||||
SKILL_INSTALL_META_DIRNAME = ".installs"
|
||||
|
||||
SKILL_BLUEPRINTS = {
|
||||
"bot-creator": {
|
||||
"hint": "build a dedicated chat bot from a tai.py variant",
|
||||
"scope": "project",
|
||||
"brief": "Deep-research preferred chat-bot features first: transport (Telegram long-poll versus webhook), command routing, per-user allowlists, rate limits, and safe restarts. Also research safe techniques for scripting huge single-file programs: exact-match edits, atomic temp-plus-rename writes, and versioned backups before every reshape. Then copy tai.py to a variant path and reshape the copy into a dedicated bot, editing efficiently and keeping the audit trail intact.",
|
||||
},
|
||||
"api-client": {
|
||||
"hint": "build a resilient REST/OpenAPI client",
|
||||
"scope": "project",
|
||||
"brief": "Deep-research resilient HTTP client design first: retries with backoff, timeouts, pagination, auth header handling, and error classification. Then write a skill that builds small stdlib-only API clients from a base URL and an endpoint list.",
|
||||
},
|
||||
"web-researcher": {
|
||||
"hint": "deep research with triangulated sources",
|
||||
"scope": "project",
|
||||
"brief": "Deep-research professional open-source research methodology first: query expansion, source triangulation, recency checks, and claim grading. Then write a skill that turns a question into a sourced brief with confidence levels.",
|
||||
},
|
||||
"pdf-forms": {
|
||||
"hint": "fill PDF forms and extract field data",
|
||||
"scope": "project",
|
||||
"brief": "Deep-research PDF form handling with freely available tooling first: AcroForm field discovery, filling, flattening, and text extraction. Then write a skill that fills forms and extracts field data from this machine's installed tools.",
|
||||
},
|
||||
"data-wrangler": {
|
||||
"hint": "reshape CSV, JSON, and SQLite data",
|
||||
"scope": "project",
|
||||
"brief": "Deep-research tabular data reshaping first: CSV dialects, JSON normalization, SQLite import ergonomics, and streaming for large files. Then write a skill that reshapes CSV, JSON, and SQLite inputs using stdlib-only scripts.",
|
||||
},
|
||||
"home-sysadmin": {
|
||||
"hint": "routine Linux host care and triage",
|
||||
"scope": "home",
|
||||
"brief": "Deep-research routine single-host Linux care first: disk, memory, service health, log triage, and backup verification. Then write a skill that runs a host checkup and reports findings with suggested fixes.",
|
||||
},
|
||||
}
|
||||
|
||||
BOX_OPENCODE_NODE_VERSION = "20.19.2"
|
||||
|
||||
# Installed on demand, not baked into BOX_CONTAINERFILE: the box exists for
|
||||
# STT/TTS today and most installs already have one running, so this only
|
||||
# pays Node's weight (~800MB, dominated by opencode-ai's own node_modules,
|
||||
# not the runtime itself) on the one machine, the one time, it's actually
|
||||
# needed as a last-resort chat backend (see ChatClient.all_candidates()).
|
||||
# The official nodejs.org tarball, not apt: Debian's nodejs/npm packages
|
||||
# drag in webpack/babel/jest as transitive deps (measured: 1.49GB vs 971MB
|
||||
# for the exact same opencode-ai install) and are not relevant here at all.
|
||||
BOX_OPENCODE_INSTALL_SCRIPT = (
|
||||
"set -e; "
|
||||
"ARCH=\"$(dpkg --print-architecture)\"; "
|
||||
"case \"$ARCH\" in amd64) NODE_ARCH=x64;; arm64) NODE_ARCH=arm64;; armhf) NODE_ARCH=armv7l;; i386) NODE_ARCH=x86;; *) echo \"unsupported arch $ARCH\" >&2; exit 1;; esac; "
|
||||
"apt-get update -qq && apt-get install -y -qq --no-install-recommends curl xz-utils ca-certificates; "
|
||||
"curl -fsSL \"https://nodejs.org/dist/v%(ver)s/node-v%(ver)s-linux-$NODE_ARCH.tar.xz\" -o /tmp/node.tar.xz; "
|
||||
"tar -xJf /tmp/node.tar.xz -C /usr/local --strip-components=1; "
|
||||
"rm /tmp/node.tar.xz; "
|
||||
"apt-get purge -y -qq curl xz-utils; apt-get autoremove -y -qq; rm -rf /var/lib/apt/lists/*; "
|
||||
"npm install -g opencode-ai --no-fund --no-audit"
|
||||
) % {"ver": BOX_OPENCODE_NODE_VERSION}
|
||||
|
||||
MEM_EVENTS_CAP = 10000
|
||||
|
||||
SCHEMA = (
|
||||
("events", ("id INTEGER PRIMARY KEY", "profile TEXT", "ts TEXT", "role TEXT", "kind TEXT", "text TEXT", "tags TEXT"), ()),
|
||||
("secrets", ("profile TEXT", "name TEXT", "value TEXT", "meta TEXT", "updated TEXT"), ("profile", "name")),
|
||||
("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"), ()),
|
||||
("records", ("id TEXT PRIMARY KEY", "kind TEXT", "profile TEXT", "title TEXT", "content TEXT", "size INTEGER", "reads INTEGER", "created TEXT", "updated TEXT"), ()),
|
||||
("tags", ("item TEXT", "tag TEXT", "profile TEXT"), ("item", "tag", "profile")),
|
||||
("edges", ("src TEXT", "dst TEXT", "relation TEXT", "profile TEXT", "created TEXT"), ("src", "dst", "relation", "profile")),
|
||||
("audit", ("id INTEGER PRIMARY KEY", "profile TEXT", "ts TEXT", "actor TEXT", "action TEXT", "path TEXT", "message TEXT", "old_size INTEGER", "new_size INTEGER", "old TEXT", "new TEXT", "tags TEXT"), ()),
|
||||
("chunks", ("id INTEGER PRIMARY KEY", "parent_kind TEXT", "parent_id TEXT", "idx INTEGER", "text TEXT", "profile TEXT", "created TEXT"), ()),
|
||||
)
|
||||
|
||||
SCHEMA_INDEXES = (
|
||||
("idx_events_profile", "events(profile)"),
|
||||
("idx_tags_tag", "tags(tag)"),
|
||||
("idx_edges_src", "edges(src)"),
|
||||
("idx_edges_dst", "edges(dst)"),
|
||||
("idx_audit_path", "audit(path)"),
|
||||
("idx_tags_profile", "tags(profile)"),
|
||||
("idx_edges_profile", "edges(profile)"),
|
||||
("idx_records_profile", "records(profile)"),
|
||||
("idx_audit_profile", "audit(profile)"),
|
||||
("idx_chunks_parent", "chunks(parent_kind, parent_id)"),
|
||||
("idx_chunks_profile", "chunks(profile)"),
|
||||
)
|
||||
|
||||
CORE_TOOLS = ("shell", "read_file", "write_file", "remember", "recall", "load_skill", "fork", "poll")
|
||||
|
||||
LAZY_NAME_STOP = ("get", "set", "list", "add")
|
||||
|
||||
TOOL_TAGS = {
|
||||
"shell": ("run", "execute", "command", "bash", "terminal", "script"),
|
||||
"read_file": ("read", "open", "view", "contents"),
|
||||
"write_file": ("write", "create", "overwrite"),
|
||||
"edit_file": ("edit", "change", "modify", "replace", "patch", "fix", "alter"),
|
||||
"web_search": ("search", "web", "internet", "lookup", "google"),
|
||||
"web_fetch": ("fetch", "download", "url", "link", "page", "curl", "wget"),
|
||||
"speak": ("speak", "say", "voice", "talk", "tts"),
|
||||
"listen": ("listen", "hear", "microphone", "mic", "transcribe", "stt", "dictate"),
|
||||
"remember": ("remember", "memorize", "preference"),
|
||||
"recall": ("recall", "memory", "history", "past"),
|
||||
"load_skill": ("skill", "capability"),
|
||||
"get_current_terminal_content": ("terminal", "tmux", "pane", "scrollback", "screen"),
|
||||
"export_terminal_log": ("terminal", "tmux", "export", "log", "capture", "history", "rag", "index"),
|
||||
"plexer_send": ("plexer", "terminal", "tmux", "keys", "send-keys", "type", "keystroke"),
|
||||
"plexer_read": ("plexer", "terminal", "tmux", "pane", "read", "capture"),
|
||||
"plexer_type": ("plexer", "terminal", "tmux", "type", "animate", "keystroke", "demo"),
|
||||
"plexer_status": ("plexer", "terminal", "tmux", "status", "session"),
|
||||
"fork": ("fork", "subagent", "background", "parallel", "delegate", "spawn"),
|
||||
"poll": ("poll", "collect"),
|
||||
"sysinfo": ("sysinfo", "system", "host", "machine", "hardware", "specs", "installed", "python", "container"),
|
||||
"create_skill": ("skill", "create", "author", "blueprint"),
|
||||
"install_skill": ("skill", "install", "import", "adopt", "copy", "symlink"),
|
||||
"store_secret": ("secret", "password", "token", "credential", "passwd", "vault", "apikey"),
|
||||
"list_secrets": ("secrets", "vault", "credentials"),
|
||||
"delete_secret": ("secret", "remove", "revoke"),
|
||||
"schedule": ("schedule", "cron", "appointment", "reminder", "recurring", "periodic", "later", "daily", "hourly", "interval"),
|
||||
"unschedule": ("unschedule", "cancel"),
|
||||
"schedules": ("schedules", "upcoming", "planned", "agenda"),
|
||||
"record_save": ("record", "note", "document", "memo"),
|
||||
"record_read": ("record", "page"),
|
||||
"record_search": ("record", "find"),
|
||||
"record_delete": ("record", "remove"),
|
||||
"graph_link": ("graph", "link", "connect", "relate", "edge", "relation"),
|
||||
"graph_query": ("graph", "neighbors", "neighborhood", "traverse", "related"),
|
||||
"delete_file": ("delete", "remove", "erase"),
|
||||
"audit": ("audit", "trail", "changes", "blame", "history", "version"),
|
||||
"restore": ("restore", "undo", "revert", "recover", "rollback", "undelete", "previous"),
|
||||
"release": ("release", "version", "bump", "changelog", "publish"),
|
||||
"tags": ("tags", "tag", "label", "categorize", "taxonomy", "vocabulary"),
|
||||
"install": ("install", "uninstall", "reinstall", "upgrade", "setup", "service", "systemd", "venv", "hook", "deploy", "zsh", "fish"),
|
||||
"create_bot": ("bot", "bots", "persona", "character", "mention"),
|
||||
"search": ("search", "find", "lookup", "everything"),
|
||||
}
|
||||
|
||||
TAG_AUTO_SKIP = ("record", "file")
|
||||
|
||||
TAG_AUTO_MAX = 5
|
||||
|
||||
TAG_LINK_MAX = 3
|
||||
|
||||
SINGULAR_KEEP = ("news", "means", "series", "species", "physics")
|
||||
|
||||
PLEXER_DIR_NAME = "plexer"
|
||||
|
||||
PLEXER_PREFIX_KEY = b"\x01" # Ctrl+A, deliberately distinct from tmux's default Ctrl+B
|
||||
|
||||
PLEXER_SOCKET_ENV = "TAI_PLEXER_SOCKET"
|
||||
|
||||
PLEXER_ACTIVE_ENV = "TAI_PLEXER_ACTIVE"
|
||||
|
||||
PLEXER_RECV_MAX = 1 << 20
|
||||
|
||||
PLEXER_BASHRC_MARK_BEGIN = "# tai plexer hook - start"
|
||||
|
||||
PLEXER_BASHRC_MARK_END = "# tai plexer hook - end"
|
||||
|
||||
PLEXER_BASHRC_BLOCK = (
|
||||
PLEXER_BASHRC_MARK_BEGIN + "\n"
|
||||
"if [ -z \"$TAI_PLEXER_ACTIVE\" ] && [ -z \"$TMUX\" ] && [ ! -e \"$HOME/.tai-no-plexer\" ] && case \"$-\" in *i*) true;; *) false;; esac; then\n"
|
||||
" if [ -t 0 ] && [ -t 1 ]; then\n"
|
||||
" exec \"$HOME/.local/bin/tai.py\" --plexer-wrap\n"
|
||||
" fi\n"
|
||||
"fi\n"
|
||||
+ PLEXER_BASHRC_MARK_END + "\n"
|
||||
)
|
||||
|
||||
FISH_PLEXER_BLOCK = (
|
||||
PLEXER_BASHRC_MARK_BEGIN + "\n"
|
||||
"if status is-interactive\n"
|
||||
"and test -z \"$TAI_PLEXER_ACTIVE\"\n"
|
||||
"and test -z \"$TMUX\"\n"
|
||||
"and not test -e \"$HOME/.tai-no-plexer\"\n"
|
||||
"and isatty stdin\n"
|
||||
"and isatty stdout\n"
|
||||
" exec \"$HOME/.local/bin/tai.py\" --plexer-wrap\n"
|
||||
"end\n"
|
||||
+ PLEXER_BASHRC_MARK_END + "\n"
|
||||
)
|
||||
|
||||
COMMANDS = ("help", "profile", "profiles", "bots", "bot", "env", "skills", "sysinfo", "models", "secret", "schedule", "schedules", "unschedule", "records", "record", "graph", "tags", "tools", "install", "search", "audit", "restore", "release", "fork", "agents", "agent", "compact", "clear", "quit", "exit")
|
||||
|
||||
TAI_SCRIPT = os.path.abspath(os.path.join(os.path.dirname(os.path.dirname(__file__)), "tai.py"))
|
||||
|
||||
TAI_CONSTS = os.path.abspath(__file__)
|
||||
+245
@@ -0,0 +1,245 @@
|
||||
"""Edge voice synthesis and local audio capture/playback."""
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import html
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
import shutil
|
||||
import socket
|
||||
import ssl
|
||||
import struct
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import uuid
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.errors import BackendError
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["EdgeTts", "Audio"]
|
||||
|
||||
|
||||
class EdgeTts:
|
||||
"""EdgeTts."""
|
||||
|
||||
@staticmethod
|
||||
def edge_token(skew=0):
|
||||
ticks = int(time.time()) + skew + 11644473600
|
||||
ticks -= ticks % 300
|
||||
ticks *= 10000000
|
||||
return hashlib.sha256(("%d%s" % (ticks, Consts.EDGE_TRUSTED_TOKEN)).encode("ascii")).hexdigest().upper()
|
||||
|
||||
@staticmethod
|
||||
def edge_timestamp():
|
||||
return time.strftime("%a %b %d %Y %H:%M:%S GMT+0000 (Coordinated Universal Time)", time.gmtime())
|
||||
|
||||
@staticmethod
|
||||
def edge_ssml(text, voice):
|
||||
safe = html.escape(text, quote=False).replace("\r", " ").replace("\n", " ")
|
||||
return "<speak version='1.0' xmlns='http://www.w3.org/2001/10/synthesis' xml:lang='en-US'><voice name='%s'><prosody pitch='+0Hz' rate='+0%%' volume='+0%%'>%s</prosody></voice></speak>" % (voice, safe)
|
||||
|
||||
@staticmethod
|
||||
def edge_chunks(text, limit=4000):
|
||||
sentences = re.split(r"(?<=[.!?])\s+", text.strip())
|
||||
parts = []
|
||||
current = ""
|
||||
for sentence in sentences:
|
||||
trial = (current + " " + sentence).strip()
|
||||
if len(trial.encode("utf-8")) > limit and current:
|
||||
parts.append(current)
|
||||
current = sentence
|
||||
else:
|
||||
current = trial
|
||||
if current:
|
||||
parts.append(current)
|
||||
return parts or [text]
|
||||
|
||||
@staticmethod
|
||||
def ws_send(sock, payload):
|
||||
if isinstance(payload, str):
|
||||
payload = payload.encode("utf-8")
|
||||
mask = os.urandom(4)
|
||||
header = bytes([0x81])
|
||||
length = len(payload)
|
||||
if length < 126:
|
||||
header += bytes([0x80 | length])
|
||||
elif length < 65536:
|
||||
header += bytes([0x80 | 126]) + struct.pack(">H", length)
|
||||
else:
|
||||
header += bytes([0x80 | 127]) + struct.pack(">Q", length)
|
||||
masked = bytes(piece ^ mask[pos % 4] for pos, piece in enumerate(payload))
|
||||
sock.sendall(header + mask + masked)
|
||||
|
||||
@staticmethod
|
||||
def ws_recv_exact(sock, count):
|
||||
data = b""
|
||||
while len(data) < count:
|
||||
piece = sock.recv(count - len(data))
|
||||
if not piece:
|
||||
raise BackendError("edge connection closed")
|
||||
data += piece
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def ws_recv(sock):
|
||||
opcode = 0
|
||||
payload = b""
|
||||
first = True
|
||||
while True:
|
||||
head = EdgeTts.ws_recv_exact(sock, 2)
|
||||
done = head[0] & 0x80
|
||||
part = head[0] & 0x0F
|
||||
masked = head[1] & 0x80
|
||||
length = head[1] & 0x7F
|
||||
if length == 126:
|
||||
length = struct.unpack(">H", EdgeTts.ws_recv_exact(sock, 2))[0]
|
||||
elif length == 127:
|
||||
length = struct.unpack(">Q", EdgeTts.ws_recv_exact(sock, 8))[0]
|
||||
key = EdgeTts.ws_recv_exact(sock, 4) if masked else b""
|
||||
body = EdgeTts.ws_recv_exact(sock, length) if length else b""
|
||||
if masked:
|
||||
body = bytes(piece ^ key[pos % 4] for pos, piece in enumerate(body))
|
||||
if part == 0x8:
|
||||
raise BackendError("edge closed stream")
|
||||
if part == 0x9:
|
||||
sock.sendall(bytes([0x8A, 0x00]))
|
||||
continue
|
||||
if part in (0x1, 0x2) and first:
|
||||
opcode = part
|
||||
first = False
|
||||
payload += body
|
||||
if done:
|
||||
return opcode, payload
|
||||
|
||||
@staticmethod
|
||||
def edge_skew(header_text):
|
||||
match = re.search(r"Date:\s*(.+?)\r?$", header_text, re.MULTILINE | re.IGNORECASE)
|
||||
if not match:
|
||||
return 0
|
||||
try:
|
||||
stamp = time.mktime(time.strptime(match.group(1).strip(), "%a, %d %b %Y %H:%M:%S %Z"))
|
||||
return int(stamp - time.time())
|
||||
except (ValueError, OverflowError):
|
||||
return 0
|
||||
|
||||
@staticmethod
|
||||
def edge_connect():
|
||||
skew = 0
|
||||
for attempt in range(2):
|
||||
path = "%s?TrustedClientToken=%s&ConnectionId=%s&Sec-MS-GEC=%s&Sec-MS-GEC-Version=1-%s" % (Consts.EDGE_WS_PATH, Consts.EDGE_TRUSTED_TOKEN, uuid.uuid4().hex, EdgeTts.edge_token(skew), Consts.EDGE_CHROMIUM)
|
||||
raw = socket.create_connection((Consts.EDGE_HOST, 443), timeout=20)
|
||||
sock = ssl.create_default_context().wrap_socket(raw, server_hostname=Consts.EDGE_HOST)
|
||||
sock.settimeout(30)
|
||||
key = base64.b64encode(os.urandom(16)).decode("ascii")
|
||||
muid = "".join(random.choice("0123456789ABCDEF") for _ in range(32))
|
||||
handshake = "\r\n".join([
|
||||
"GET %s HTTP/1.1" % path,
|
||||
"Host: %s" % Consts.EDGE_HOST,
|
||||
"Upgrade: websocket",
|
||||
"Connection: Upgrade",
|
||||
"Sec-WebSocket-Version: 13",
|
||||
"Sec-WebSocket-Key: %s" % key,
|
||||
"Pragma: no-cache",
|
||||
"Cache-Control: no-cache",
|
||||
"Origin: chrome-extension://jdiccldimpdaibmpdkjnbmckianbfold",
|
||||
"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36 Edg/143.0.0.0",
|
||||
"Accept-Encoding: gzip, deflate, br, zstd",
|
||||
"Accept-Language: en-US,en;q=0.9",
|
||||
"Cookie: muid=%s;" % muid,
|
||||
"", "",
|
||||
])
|
||||
sock.sendall(handshake.encode("ascii"))
|
||||
answer = b""
|
||||
while b"\r\n\r\n" not in answer:
|
||||
piece = sock.recv(4096)
|
||||
if not piece:
|
||||
break
|
||||
answer += piece
|
||||
header_text = answer.decode("latin-1", "replace")
|
||||
if " 101 " in header_text.split("\r\n", 1)[0]:
|
||||
return sock
|
||||
sock.close()
|
||||
if " 403 " in header_text and attempt == 0:
|
||||
skew = EdgeTts.edge_skew(header_text)
|
||||
continue
|
||||
raise BackendError("edge handshake failed: %s" % header_text.split("\r\n", 1)[0][:200])
|
||||
raise BackendError("edge handshake failed after retry")
|
||||
|
||||
@staticmethod
|
||||
def edge_synthesize(text, voice):
|
||||
sock = EdgeTts.edge_connect()
|
||||
audio = bytearray()
|
||||
try:
|
||||
config = "X-Timestamp:%s\r\nContent-Type:application/json; charset=utf-8\r\nPath:speech.config\r\n\r\n{\"context\":{\"synthesis\":{\"audio\":{\"metadataoptions\":{\"sentenceBoundaryEnabled\":\"false\",\"wordBoundaryEnabled\":\"true\"},\"outputFormat\":\"audio-24khz-48kbitrate-mono-mp3\"}}}}" % EdgeTts.edge_timestamp()
|
||||
EdgeTts.ws_send(sock, config)
|
||||
for chunk in EdgeTts.edge_chunks(text):
|
||||
message = "X-RequestId:%s\r\nContent-Type:application/ssml+xml\r\nX-Timestamp:%sZ\r\nPath:ssml\r\n\r\n%s" % (uuid.uuid4().hex, EdgeTts.edge_timestamp(), EdgeTts.edge_ssml(chunk, voice))
|
||||
EdgeTts.ws_send(sock, message)
|
||||
while True:
|
||||
opcode, payload = EdgeTts.ws_recv(sock)
|
||||
if opcode == 0x2 and len(payload) > 2:
|
||||
head_len = int.from_bytes(payload[0:2], "big")
|
||||
audio += payload[2 + head_len:]
|
||||
elif opcode == 0x1 and "Path:turn.end" in payload.decode("utf-8", "replace"):
|
||||
break
|
||||
return bytes(audio)
|
||||
finally:
|
||||
try:
|
||||
sock.close()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
class Audio:
|
||||
"""Audio."""
|
||||
|
||||
@staticmethod
|
||||
def play_audio(path):
|
||||
if sys.platform == "darwin" and shutil.which("afplay"):
|
||||
runner = ["afplay", path]
|
||||
elif shutil.which("paplay"):
|
||||
runner = ["paplay", path]
|
||||
elif shutil.which("aplay"):
|
||||
runner = ["aplay", "-q", path]
|
||||
elif shutil.which("pw-play"):
|
||||
runner = ["pw-play", path]
|
||||
elif sys.platform == "win32":
|
||||
runner = ["powershell", "-c", "(New-Object Media.SoundPlayer '%s').PlaySync();" % path]
|
||||
else:
|
||||
return False
|
||||
try:
|
||||
subprocess.run(runner, timeout=300, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=False)
|
||||
return True
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def listen_audio(seconds, workdir):
|
||||
recorder = next((item for item in Consts.RECORDERS if shutil.which(item[0])), None)
|
||||
transcriber = next((name for name in Consts.TRANSCRIBERS if shutil.which(name)), "")
|
||||
if recorder is None or not transcriber:
|
||||
missing = []
|
||||
if recorder is None:
|
||||
missing.append("audio recorder (install arecord, sox, or ffmpeg)")
|
||||
if not transcriber:
|
||||
missing.append("transcriber (install whisper-cpp or whisper)")
|
||||
return "listening unavailable, missing: " + ", ".join(missing)
|
||||
path = os.path.join(workdir, "listen-%d.wav" % int(time.time()))
|
||||
command = [part.replace("{seconds}", str(seconds)).replace("{path}", path) for part in recorder[1]]
|
||||
try:
|
||||
subprocess.run(command, timeout=seconds + 15, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=False)
|
||||
done = subprocess.run([transcriber, path], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=180)
|
||||
except (OSError, subprocess.SubprocessError) as exc:
|
||||
return "listening failed: " + Text.short_error(exc)
|
||||
finally:
|
||||
try:
|
||||
os.remove(path)
|
||||
except OSError:
|
||||
pass
|
||||
if done.returncode != 0:
|
||||
return "transcriber error: " + (done.stderr or "")[:500]
|
||||
text = (done.stdout or "").strip()
|
||||
return Text.truncate(text, 3000, 500) if text else "heard nothing"
|
||||
@@ -0,0 +1,15 @@
|
||||
"""Shared exception types."""
|
||||
|
||||
__all__ = ["BackendError", "Denied"]
|
||||
|
||||
|
||||
class BackendError(Exception):
|
||||
def __init__(self, message, status=0):
|
||||
super().__init__(message)
|
||||
self.status = status
|
||||
|
||||
|
||||
class Denied(Exception):
|
||||
def __init__(self, guidance=""):
|
||||
super().__init__("denied by user")
|
||||
self.guidance = guidance
|
||||
@@ -0,0 +1,71 @@
|
||||
"""Gemini CLI backend."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
|
||||
from tai_pkg.clirun import Cli
|
||||
from tai_pkg.errors import BackendError
|
||||
from tai_pkg.opencode import Opencode
|
||||
|
||||
__all__ = ["Gemini"]
|
||||
|
||||
|
||||
class Gemini:
|
||||
"""Gemini."""
|
||||
|
||||
@staticmethod
|
||||
def gemini_binary():
|
||||
return os.environ.get("GEMINI_BIN") or shutil.which("gemini")
|
||||
|
||||
@staticmethod
|
||||
def gemini_parse_output(stdout):
|
||||
try:
|
||||
data = json.loads(stdout or "")
|
||||
except ValueError:
|
||||
return Opencode.opencode_parse(stdout)
|
||||
if not isinstance(data, dict):
|
||||
return Opencode.opencode_parse(stdout)
|
||||
err = data.get("error")
|
||||
if err:
|
||||
raise BackendError("gemini: %s" % (json.dumps(err) if isinstance(err, dict) else str(err))[:300])
|
||||
response = data.get("response")
|
||||
if isinstance(response, str) and response.strip():
|
||||
return Opencode.opencode_parse(response)
|
||||
return Opencode.opencode_parse(stdout)
|
||||
|
||||
@staticmethod
|
||||
def gemini_run(prompt, timeout):
|
||||
binary = Gemini.gemini_binary()
|
||||
if not binary:
|
||||
raise BackendError("gemini CLI not found (GEMINI_BIN or PATH)")
|
||||
base = [binary, "-p", prompt]
|
||||
variants = [
|
||||
base + ["--output-format", "json", "--approval-mode", "plan"],
|
||||
base + ["--output-format", "json"],
|
||||
base,
|
||||
]
|
||||
last_error = ""
|
||||
for argv in variants:
|
||||
done = Cli._cli_run(argv, None, timeout, "gemini")
|
||||
if done.returncode != 0 and Cli._cli_unknown_flag(done.stderr) and argv is not variants[-1]:
|
||||
last_error = (done.stderr or "").strip()[:200]
|
||||
continue
|
||||
if done.returncode != 0 and not (done.stdout or "").strip():
|
||||
blob = ((done.stderr or "") + "\n" + (done.stdout or "")).strip()
|
||||
if Cli._cli_auth_failure(blob):
|
||||
raise BackendError("gemini not authenticated (run gemini login or set GEMINI_API_KEY): %s" % blob[:200])
|
||||
hint = {42: "input error", 53: "turn limit exceeded"}.get(done.returncode, "")
|
||||
raise BackendError("gemini exit %d%s: %s" % (done.returncode, " (%s)" % hint if hint else "", blob[:300] or last_error or "no output"))
|
||||
if not (done.stdout or "").strip():
|
||||
blob = (done.stderr or "").strip()
|
||||
raise BackendError("gemini returned no output: %s" % (blob[:300] or last_error or "empty stdout"))
|
||||
return done.stdout or ""
|
||||
raise BackendError("gemini failed: %s" % (last_error or "no usable output"))
|
||||
|
||||
@staticmethod
|
||||
def gemini_health():
|
||||
binary = Gemini.gemini_binary()
|
||||
if not binary:
|
||||
return "missing: gemini CLI not found"
|
||||
return Cli._cli_version_only_health(binary, "gemini")
|
||||
@@ -0,0 +1,78 @@
|
||||
"""Grok CLI backend."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
from tai_pkg.clirun import Cli
|
||||
from tai_pkg.errors import BackendError
|
||||
from tai_pkg.opencode import Opencode
|
||||
|
||||
__all__ = ["Grok"]
|
||||
|
||||
|
||||
class Grok:
|
||||
"""Grok."""
|
||||
|
||||
@staticmethod
|
||||
def grok_binary():
|
||||
return os.environ.get("GROK_BIN") or shutil.which("grok") or shutil.which("grokcli")
|
||||
|
||||
@staticmethod
|
||||
def grok_parse_output(stdout):
|
||||
try:
|
||||
data = json.loads(stdout or "")
|
||||
except ValueError:
|
||||
return Opencode.opencode_parse(stdout)
|
||||
if isinstance(data, dict):
|
||||
if data.get("is_error") or data.get("error"):
|
||||
detail = data.get("error") or data.get("result") or data.get("text") or "unknown error"
|
||||
raise BackendError("grok: %s" % str(detail)[:300])
|
||||
if isinstance(data.get("content"), str) or isinstance(data.get("tool_calls"), list):
|
||||
if "content" in data or "tool_calls" in data:
|
||||
return data
|
||||
for key in ("text", "result", "output", "answer", "content"):
|
||||
value = data.get(key)
|
||||
if isinstance(value, str) and value.strip():
|
||||
return Opencode.opencode_parse(value)
|
||||
return Opencode.opencode_parse(stdout)
|
||||
|
||||
@staticmethod
|
||||
def grok_run(prompt, timeout):
|
||||
binary = Grok.grok_binary()
|
||||
if not binary:
|
||||
raise BackendError("grok CLI not found (GROK_BIN or PATH)")
|
||||
with tempfile.TemporaryDirectory(prefix="tai-grok-") as folder:
|
||||
prompt_path = os.path.join(folder, "prompt.txt")
|
||||
with open(prompt_path, "w", encoding="utf-8") as handle:
|
||||
handle.write(prompt)
|
||||
variants = [
|
||||
[binary, "--prompt-file", prompt_path, "--output-format", "json", "--no-auto-update"],
|
||||
[binary, "--prompt-file", prompt_path, "--output-format", "json"],
|
||||
[binary, "-p", prompt, "--output-format", "json"],
|
||||
[binary, "-p", prompt],
|
||||
]
|
||||
last_error = ""
|
||||
for argv in variants:
|
||||
done = Cli._cli_run(argv, None, timeout, "grok")
|
||||
if done.returncode != 0 and Cli._cli_unknown_flag(done.stderr) and argv is not variants[-1]:
|
||||
last_error = (done.stderr or "").strip()[:200]
|
||||
continue
|
||||
if done.returncode != 0 and not (done.stdout or "").strip():
|
||||
blob = ((done.stderr or "") + "\n" + (done.stdout or "")).strip()
|
||||
if Cli._cli_auth_failure(blob):
|
||||
raise BackendError("grok not authenticated (set XAI_API_KEY or run grok login): %s" % blob[:200])
|
||||
raise BackendError("grok exit %d: %s" % (done.returncode, blob[:300] or last_error or "no output"))
|
||||
if not (done.stdout or "").strip():
|
||||
blob = (done.stderr or "").strip()
|
||||
raise BackendError("grok returned no output: %s" % (blob[:300] or last_error or "empty stdout"))
|
||||
return done.stdout or ""
|
||||
raise BackendError("grok failed: %s" % (last_error or "no usable output"))
|
||||
|
||||
@staticmethod
|
||||
def grok_health():
|
||||
binary = Grok.grok_binary()
|
||||
if not binary:
|
||||
return "missing: grok CLI not found"
|
||||
return Cli._cli_version_only_health(binary, "grok")
|
||||
@@ -0,0 +1,676 @@
|
||||
"""Install, hook, service, and container operations."""
|
||||
|
||||
import getpass
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import sqlite3
|
||||
import subprocess
|
||||
import sys
|
||||
import importlib.util
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.errors import BackendError
|
||||
from tai_pkg.sandbox import Sandbox
|
||||
from tai_pkg.telegram import Telegram
|
||||
from tai_pkg.telegram import TelegramBot
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["Install"]
|
||||
|
||||
|
||||
class Install:
|
||||
"""Install."""
|
||||
|
||||
@staticmethod
|
||||
def run_host(argv, timeout=120):
|
||||
try:
|
||||
done = subprocess.run(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=timeout)
|
||||
return done.returncode == 0, (done.stdout or "") + (done.stderr or "")
|
||||
except (OSError, subprocess.SubprocessError) as exc:
|
||||
return False, Text.short_error(exc)
|
||||
|
||||
@staticmethod
|
||||
def venv_available():
|
||||
if importlib.util.find_spec("venv") is None:
|
||||
return False
|
||||
if importlib.util.find_spec("ensurepip") is not None:
|
||||
return True
|
||||
return shutil.which("pip3") is not None or shutil.which("pip") is not None
|
||||
|
||||
@staticmethod
|
||||
def venv_dir(home):
|
||||
return os.path.join(home, "venv")
|
||||
|
||||
@staticmethod
|
||||
def venv_python(home):
|
||||
cand = os.path.join(Install.venv_dir(home), "bin", "python")
|
||||
if os.path.isfile(cand) and os.access(cand, os.X_OK):
|
||||
return cand
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def ensure_venv(home):
|
||||
found = Install.venv_python(home)
|
||||
if found:
|
||||
return True, "kept %s" % found
|
||||
if not Install.venv_available():
|
||||
return False, "python venv module unavailable"
|
||||
try:
|
||||
done = subprocess.run([sys.executable, "-m", "venv", Install.venv_dir(home)], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=300)
|
||||
except (OSError, subprocess.SubprocessError) as exc:
|
||||
return False, Text.short_error(exc)
|
||||
if done.returncode != 0 or not Install.venv_python(home):
|
||||
return False, ((done.stderr or "").strip()[:200] or "venv creation failed")
|
||||
return True, "created %s" % Install.venv_python(home)
|
||||
|
||||
@staticmethod
|
||||
def service_exec(home, script, *args):
|
||||
parts = []
|
||||
python = Install.venv_python(home)
|
||||
if python:
|
||||
parts.append(python)
|
||||
parts.append(script)
|
||||
parts.extend(args)
|
||||
return " ".join(parts)
|
||||
|
||||
@staticmethod
|
||||
def unit_dir(home):
|
||||
return os.path.join(os.path.expanduser("~"), ".config", "systemd", "user")
|
||||
|
||||
@staticmethod
|
||||
def service_state(name):
|
||||
if not shutil.which("systemctl"):
|
||||
return "unknown", "unknown"
|
||||
active, _out = Install.run_host(["systemctl", "--user", "is-active", name], timeout=15)
|
||||
enabled, _out = Install.run_host(["systemctl", "--user", "is-enabled", name], timeout=15)
|
||||
return ("active" if active else "inactive"), ("enabled" if enabled else "disabled")
|
||||
|
||||
@staticmethod
|
||||
def activate_service(name):
|
||||
if not shutil.which("systemctl"):
|
||||
return "systemctl missing, unit written only"
|
||||
Install.run_host(["systemctl", "--user", "daemon-reload"])
|
||||
active, _enabled = Install.service_state(name)
|
||||
if active == "active":
|
||||
ok, out = Install.run_host(["systemctl", "--user", "restart", name])
|
||||
return "restarted, changes effective immediately" if ok else "restart failed: " + out.strip()[:150]
|
||||
ok, out = Install.run_host(["systemctl", "--user", "enable", "--now", name])
|
||||
if ok:
|
||||
user = os.environ.get("USER") or getpass.getuser()
|
||||
Install.run_host(["loginctl", "enable-linger", user])
|
||||
return "enabled and started"
|
||||
return "enable failed: " + out.strip()[:150]
|
||||
|
||||
@staticmethod
|
||||
def remove_service_unit(home, name):
|
||||
path = os.path.join(Install.unit_dir(home), name)
|
||||
if shutil.which("systemctl"):
|
||||
Install.run_host(["systemctl", "--user", "disable", "--now", name])
|
||||
Install.run_host(["systemctl", "--user", "daemon-reload"])
|
||||
if os.path.exists(path):
|
||||
try:
|
||||
os.remove(path)
|
||||
return True, "removed " + path
|
||||
except OSError as exc:
|
||||
return False, Text.short_error(exc)
|
||||
return True, "not present"
|
||||
|
||||
@staticmethod
|
||||
def script_version(path):
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
match = Consts.VERSION_RE.search(handle.read())
|
||||
except OSError:
|
||||
return None
|
||||
return ".".join(match.groups()) if match else None
|
||||
|
||||
@staticmethod
|
||||
def shell_rc_path(shell):
|
||||
"""Where each shell's startup file lives. bash/zsh: a dotfile sourced
|
||||
for interactive shells only. fish: config.fish, which fish sources on
|
||||
every startup (interactive or not, login or not) -- hence the
|
||||
`status is-interactive` guard inside FISH_PLEXER_BLOCK rather than
|
||||
relying on the file only being read at the right time."""
|
||||
home = os.path.expanduser("~")
|
||||
if shell == "fish":
|
||||
config_home = os.environ.get("XDG_CONFIG_HOME") or os.path.join(home, ".config")
|
||||
return os.path.join(config_home, "fish", "config.fish")
|
||||
if shell == "zsh":
|
||||
return os.path.join(home, ".zshrc")
|
||||
return os.path.join(home, ".bashrc")
|
||||
|
||||
@staticmethod
|
||||
def shell_present(shell):
|
||||
return shutil.which(shell) is not None
|
||||
|
||||
@staticmethod
|
||||
def _rc_block_present(path, block):
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
return block in handle.read()
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _rc_block_remove(path, mark_begin, mark_end):
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
raw = handle.read()
|
||||
except OSError:
|
||||
return False
|
||||
stripped = re.sub(re.escape(mark_begin) + r".*?" + re.escape(mark_end) + r"\n?", "", raw, flags=re.DOTALL)
|
||||
if stripped == raw:
|
||||
return False
|
||||
backup = path + ".bak-tai"
|
||||
if raw and not os.path.exists(backup):
|
||||
with open(backup, "w", encoding="utf-8") as handle:
|
||||
handle.write(raw)
|
||||
with open(path, "w", encoding="utf-8") as handle:
|
||||
handle.write(stripped)
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _rc_block_upsert(path, mark_begin, mark_end, block):
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
raw = handle.read()
|
||||
except OSError:
|
||||
raw = ""
|
||||
if block in raw:
|
||||
return False
|
||||
backup = path + ".bak-tai"
|
||||
if raw and not os.path.exists(backup):
|
||||
with open(backup, "w", encoding="utf-8") as handle:
|
||||
handle.write(raw)
|
||||
stripped = re.sub(re.escape(mark_begin) + r".*?" + re.escape(mark_end) + r"\n?", "", raw, flags=re.DOTALL)
|
||||
if stripped and not stripped.endswith("\n"):
|
||||
stripped += "\n"
|
||||
parent = os.path.dirname(path)
|
||||
if parent and not os.path.isdir(parent):
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
with open(path, "w", encoding="utf-8") as handle:
|
||||
handle.write(stripped + block)
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def hook_present(path):
|
||||
return Install._rc_block_present(path, Consts.BASHRC_BLOCK)
|
||||
|
||||
@staticmethod
|
||||
def remove_bashrc_block(path):
|
||||
return Install._rc_block_remove(path, Consts.BASHRC_MARK_BEGIN, Consts.BASHRC_MARK_END)
|
||||
|
||||
@staticmethod
|
||||
def zsh_hook_present(path):
|
||||
return Install._rc_block_present(path, Consts.ZSH_NOTFOUND_BLOCK)
|
||||
|
||||
@staticmethod
|
||||
def remove_zsh_hook_block(path):
|
||||
return Install._rc_block_remove(path, Consts.BASHRC_MARK_BEGIN, Consts.BASHRC_MARK_END)
|
||||
|
||||
@staticmethod
|
||||
def upsert_zsh_hook_block(path):
|
||||
return Install._rc_block_upsert(path, Consts.BASHRC_MARK_BEGIN, Consts.BASHRC_MARK_END, Consts.ZSH_NOTFOUND_BLOCK)
|
||||
|
||||
@staticmethod
|
||||
def fish_hook_present(path):
|
||||
return Install._rc_block_present(path, Consts.FISH_NOTFOUND_BLOCK)
|
||||
|
||||
@staticmethod
|
||||
def remove_fish_hook_block(path):
|
||||
return Install._rc_block_remove(path, Consts.BASHRC_MARK_BEGIN, Consts.BASHRC_MARK_END)
|
||||
|
||||
@staticmethod
|
||||
def upsert_fish_hook_block(path):
|
||||
return Install._rc_block_upsert(path, Consts.BASHRC_MARK_BEGIN, Consts.BASHRC_MARK_END, Consts.FISH_NOTFOUND_BLOCK)
|
||||
|
||||
@staticmethod
|
||||
def plexer_hook_present(path):
|
||||
return Install._rc_block_present(path, Consts.PLEXER_BASHRC_BLOCK)
|
||||
|
||||
@staticmethod
|
||||
def remove_plexer_bashrc_block(path):
|
||||
return Install._rc_block_remove(path, Consts.PLEXER_BASHRC_MARK_BEGIN, Consts.PLEXER_BASHRC_MARK_END)
|
||||
|
||||
@staticmethod
|
||||
def upsert_plexer_bashrc_block(path):
|
||||
return Install._rc_block_upsert(path, Consts.PLEXER_BASHRC_MARK_BEGIN, Consts.PLEXER_BASHRC_MARK_END, Consts.PLEXER_BASHRC_BLOCK)
|
||||
|
||||
@staticmethod
|
||||
# zsh reuses PLEXER_BASHRC_BLOCK byte-for-byte -- zsh is POSIX-enough that
|
||||
# the plain `[`/`case "$-"` bash content works unchanged; only the target
|
||||
# file differs (~/.zshrc, not ~/.bashrc). No ZSH_PLEXER_BLOCK constant exists
|
||||
# on purpose -- there is nothing shell-specific to put in one.
|
||||
def zsh_plexer_hook_present(path):
|
||||
return Install._rc_block_present(path, Consts.PLEXER_BASHRC_BLOCK)
|
||||
|
||||
@staticmethod
|
||||
def remove_zsh_plexer_block(path):
|
||||
return Install._rc_block_remove(path, Consts.PLEXER_BASHRC_MARK_BEGIN, Consts.PLEXER_BASHRC_MARK_END)
|
||||
|
||||
@staticmethod
|
||||
def upsert_zsh_plexer_block(path):
|
||||
return Install._rc_block_upsert(path, Consts.PLEXER_BASHRC_MARK_BEGIN, Consts.PLEXER_BASHRC_MARK_END, Consts.PLEXER_BASHRC_BLOCK)
|
||||
|
||||
@staticmethod
|
||||
def fish_plexer_hook_present(path):
|
||||
return Install._rc_block_present(path, Consts.FISH_PLEXER_BLOCK)
|
||||
|
||||
@staticmethod
|
||||
def remove_fish_plexer_block(path):
|
||||
return Install._rc_block_remove(path, Consts.PLEXER_BASHRC_MARK_BEGIN, Consts.PLEXER_BASHRC_MARK_END)
|
||||
|
||||
@staticmethod
|
||||
def upsert_fish_plexer_block(path):
|
||||
return Install._rc_block_upsert(path, Consts.PLEXER_BASHRC_MARK_BEGIN, Consts.PLEXER_BASHRC_MARK_END, Consts.FISH_PLEXER_BLOCK)
|
||||
|
||||
@staticmethod
|
||||
def box_venv_state(engine):
|
||||
try:
|
||||
probe = Sandbox.box_exec(engine, ["test", "-x", Consts.BOX_PYTHON], timeout=30)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
return "unknown"
|
||||
return "venv ready" if probe.returncode == 0 else "no venv (reinstall container to rebuild)"
|
||||
|
||||
@staticmethod
|
||||
def shell_hook_status_line(label, shell, present_fn):
|
||||
path = Install.shell_rc_path(shell)
|
||||
note = "" if Install.shell_present(shell) else " (%s not found)" % shell
|
||||
return "%s: %s in %s%s" % (label, "installed" if present_fn(path) else "not installed", path, note)
|
||||
|
||||
@staticmethod
|
||||
def status_lines(home, store=None):
|
||||
lines = []
|
||||
binary = os.path.join(os.path.expanduser("~"), ".local", "bin", "tai.py")
|
||||
if os.path.isfile(binary):
|
||||
consts_path = os.path.join(os.path.dirname(binary), "tai_pkg", "consts.py")
|
||||
lines.append("binary: installed %s (%s, running %s)" % (binary, Install.script_version(consts_path) or "unknown", Consts.VERSION))
|
||||
else:
|
||||
lines.append("binary: not installed")
|
||||
lines.append("bash-hook: %s in ~/.bashrc" % ("installed" if Install.hook_present(os.path.join(os.path.expanduser("~"), ".bashrc")) else "not installed"))
|
||||
lines.append(Install.shell_hook_status_line("zsh-hook", "zsh", Install.zsh_hook_present))
|
||||
lines.append(Install.shell_hook_status_line("fish-hook", "fish", Install.fish_hook_present))
|
||||
lines.append("plexer-hook: %s in ~/.bashrc" % ("installed" if Install.plexer_hook_present(os.path.join(os.path.expanduser("~"), ".bashrc")) else "not installed"))
|
||||
lines.append(Install.shell_hook_status_line("zsh-plexer-hook", "zsh", Install.zsh_plexer_hook_present))
|
||||
lines.append(Install.shell_hook_status_line("fish-plexer-hook", "fish", Install.fish_plexer_hook_present))
|
||||
python = Install.venv_python(home)
|
||||
if python:
|
||||
lines.append("venv: ready %s" % python)
|
||||
elif Install.venv_available():
|
||||
lines.append("venv: not created")
|
||||
else:
|
||||
lines.append("venv: unavailable (no venv module)")
|
||||
for name, label in (("tai-scheduler.service", "scheduler-service"), ("tai-telegram.service", "telegram-service")):
|
||||
unit = os.path.join(Install.unit_dir(home), name)
|
||||
if not os.path.isfile(unit):
|
||||
lines.append("%s: not installed" % label)
|
||||
continue
|
||||
active, enabled = Install.service_state(name)
|
||||
lines.append("%s: unit present, %s, %s" % (label, active, enabled))
|
||||
engine = Sandbox.container_engine()
|
||||
if not engine:
|
||||
lines.append("container: no engine (install podman)")
|
||||
else:
|
||||
state = Sandbox.box_state(engine)
|
||||
if state == "running":
|
||||
lines.append("container: running, %s" % Install.box_venv_state(engine))
|
||||
else:
|
||||
lines.append("container: %s" % state)
|
||||
lines.append(Install.vault_line(home, store))
|
||||
return lines
|
||||
|
||||
@staticmethod
|
||||
def vault_line(home, store=None):
|
||||
db_path = os.path.join(home, "memory.db")
|
||||
if not os.path.isfile(db_path):
|
||||
return "vault: empty, install ops never touch data"
|
||||
if store is None:
|
||||
return "vault: kept %s, install ops never touch data" % db_path
|
||||
try:
|
||||
secrets = store.db.execute("SELECT COUNT(*) FROM secrets").fetchone()[0]
|
||||
schedules = store.db.execute("SELECT COUNT(*) FROM schedules").fetchone()[0]
|
||||
records = store.db.execute("SELECT COUNT(*) FROM records").fetchone()[0]
|
||||
except sqlite3.Error:
|
||||
return "vault: kept %s, install ops never touch data" % db_path
|
||||
return "vault: kept (%d secrets, %d schedules, %d records), install ops never touch data" % (secrets, schedules, records)
|
||||
|
||||
@staticmethod
|
||||
def op_binary(mode, home):
|
||||
target = os.path.join(os.path.expanduser("~"), ".local", "bin", "tai.py")
|
||||
pkg_target = os.path.join(os.path.dirname(target), "tai_pkg")
|
||||
consts_target = os.path.join(pkg_target, "consts.py")
|
||||
src = os.path.abspath(Consts.TAI_SCRIPT)
|
||||
src_pkg = os.path.dirname(os.path.abspath(Consts.TAI_CONSTS))
|
||||
current = Install.script_version(consts_target) if os.path.isfile(consts_target) else None
|
||||
if mode == "install" and current is not None:
|
||||
return "binary: already installed %s (running %s)" % (current, Consts.VERSION)
|
||||
if mode == "uninstall":
|
||||
if not os.path.exists(target) and not os.path.exists(pkg_target):
|
||||
return "binary: not present"
|
||||
try:
|
||||
if os.path.exists(target):
|
||||
os.remove(target)
|
||||
if os.path.isdir(pkg_target):
|
||||
shutil.rmtree(pkg_target)
|
||||
return "binary: removed %s" % target
|
||||
except OSError as exc:
|
||||
return "binary: remove failed: %s" % Text.short_error(exc)
|
||||
try:
|
||||
os.makedirs(os.path.dirname(target), exist_ok=True)
|
||||
with open(src, "rb") as handle:
|
||||
blob = handle.read()
|
||||
with open(target, "wb") as handle:
|
||||
handle.write(blob)
|
||||
os.chmod(target, 0o755)
|
||||
if os.path.abspath(src_pkg) != os.path.abspath(pkg_target):
|
||||
if os.path.isdir(pkg_target):
|
||||
shutil.rmtree(pkg_target)
|
||||
shutil.copytree(src_pkg, pkg_target, ignore=shutil.ignore_patterns("__pycache__"))
|
||||
except OSError as exc:
|
||||
return "binary: copy failed: %s" % Text.short_error(exc)
|
||||
if current is None:
|
||||
return "binary: installed %s" % Consts.VERSION
|
||||
return "binary: refreshed %s -> %s" % (current, Consts.VERSION)
|
||||
|
||||
@staticmethod
|
||||
def op_hook(mode, home):
|
||||
path = os.path.join(os.path.expanduser("~"), ".bashrc")
|
||||
if mode == "uninstall":
|
||||
return "bash-hook: %s" % ("removed" if Install.remove_bashrc_block(path) else "not present")
|
||||
if mode == "reinstall":
|
||||
Install.remove_bashrc_block(path)
|
||||
try:
|
||||
changed = Install.upsert_bashrc_block(path)
|
||||
except OSError as exc:
|
||||
return "bash-hook: write failed: %s" % Text.short_error(exc)
|
||||
if changed:
|
||||
return "bash-hook: installed, run: source ~/.bashrc"
|
||||
return "bash-hook: already present"
|
||||
|
||||
@staticmethod
|
||||
def op_plexer_hook(mode, home):
|
||||
path = os.path.join(os.path.expanduser("~"), ".bashrc")
|
||||
if mode == "uninstall":
|
||||
return "plexer-hook: %s" % ("removed" if Install.remove_plexer_bashrc_block(path) else "not present")
|
||||
if mode == "reinstall":
|
||||
Install.remove_plexer_bashrc_block(path)
|
||||
try:
|
||||
changed = Install.upsert_plexer_bashrc_block(path)
|
||||
except OSError as exc:
|
||||
return "plexer-hook: write failed: %s" % Text.short_error(exc)
|
||||
if changed:
|
||||
return "plexer-hook: installed, run: source ~/.bashrc (new interactive shells auto-wrap; existing ones are unaffected until restarted)"
|
||||
return "plexer-hook: already present"
|
||||
|
||||
@staticmethod
|
||||
def op_shell_hook(mode, shell, label, upsert_fn, remove_fn, source_hint):
|
||||
"""Shared body for the zsh/fish hook variants -- mirrors op_hook/op_plexer_hook
|
||||
exactly, plus one addition: install/upgrade skip loudly (never write a dangling
|
||||
hook into a config file for a shell that isn't even on this machine) when the
|
||||
shell binary isn't found. uninstall/reinstall/status still act on the file
|
||||
regardless, so a hook installed before a shell was later removed still cleans up."""
|
||||
path = Install.shell_rc_path(shell)
|
||||
if mode in ("install", "upgrade") and not Install.shell_present(shell):
|
||||
return "%s: %s not found, skipped" % (label, shell)
|
||||
if mode == "uninstall":
|
||||
return "%s: %s" % (label, "removed" if remove_fn(path) else "not present")
|
||||
if mode == "reinstall":
|
||||
remove_fn(path)
|
||||
try:
|
||||
changed = upsert_fn(path)
|
||||
except OSError as exc:
|
||||
return "%s: write failed: %s" % (label, Text.short_error(exc))
|
||||
if changed:
|
||||
return "%s: installed, run: %s" % (label, source_hint)
|
||||
return "%s: already present" % label
|
||||
|
||||
@staticmethod
|
||||
def op_zsh_hook(mode, home):
|
||||
return Install.op_shell_hook(mode, "zsh", "zsh-hook", Install.upsert_zsh_hook_block, Install.remove_zsh_hook_block, "source ~/.zshrc")
|
||||
|
||||
@staticmethod
|
||||
def op_fish_hook(mode, home):
|
||||
return Install.op_shell_hook(mode, "fish", "fish-hook", Install.upsert_fish_hook_block, Install.remove_fish_hook_block, "open a new fish shell")
|
||||
|
||||
@staticmethod
|
||||
def op_zsh_plexer_hook(mode, home):
|
||||
return Install.op_shell_hook(
|
||||
mode, "zsh", "zsh-plexer-hook", Install.upsert_zsh_plexer_block, Install.remove_zsh_plexer_block,
|
||||
"source ~/.zshrc (new interactive shells auto-wrap; existing ones are unaffected until restarted)",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def op_fish_plexer_hook(mode, home):
|
||||
return Install.op_shell_hook(
|
||||
mode, "fish", "fish-plexer-hook", Install.upsert_fish_plexer_block, Install.remove_fish_plexer_block,
|
||||
"open a new fish shell (new interactive shells auto-wrap; existing ones are unaffected until restarted)",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def op_venv(mode, home):
|
||||
if mode == "uninstall":
|
||||
folder = Install.venv_dir(home)
|
||||
if not os.path.isdir(folder):
|
||||
return "venv: not present"
|
||||
try:
|
||||
shutil.rmtree(folder)
|
||||
return "venv: removed %s" % folder
|
||||
except OSError as exc:
|
||||
return "venv: remove failed: %s" % Text.short_error(exc)
|
||||
if mode == "reinstall" and os.path.isdir(Install.venv_dir(home)):
|
||||
try:
|
||||
shutil.rmtree(Install.venv_dir(home))
|
||||
except OSError as exc:
|
||||
return "venv: rebuild failed: %s" % Text.short_error(exc)
|
||||
ok, detail = Install.ensure_venv(home)
|
||||
return "venv: %s" % detail
|
||||
|
||||
@staticmethod
|
||||
def op_service(mode, home, name, label, body):
|
||||
unit = os.path.join(Install.unit_dir(home), name)
|
||||
if mode == "uninstall":
|
||||
_done, detail = Install.remove_service_unit(home, name)
|
||||
return "%s: %s" % (label, detail)
|
||||
existed = os.path.isfile(unit)
|
||||
if mode == "install" and existed:
|
||||
active, enabled = Install.service_state(name)
|
||||
return "%s: already installed (%s, %s)" % (label, active, enabled)
|
||||
if mode == "reinstall" and shutil.which("systemctl"):
|
||||
Install.run_host(["systemctl", "--user", "disable", "--now", name])
|
||||
try:
|
||||
os.makedirs(os.path.dirname(unit), exist_ok=True)
|
||||
with open(unit, "w", encoding="utf-8") as handle:
|
||||
handle.write(body)
|
||||
except OSError as exc:
|
||||
return "%s: unit write failed: %s" % (label, Text.short_error(exc))
|
||||
return "%s: unit %s, %s" % (label, "rewritten" if existed else "written", Install.activate_service(name))
|
||||
|
||||
@staticmethod
|
||||
def op_scheduler_service(mode, home):
|
||||
body = Consts.SCHEDULER_UNIT % Install.service_exec(home, os.path.abspath(Consts.TAI_SCRIPT))
|
||||
return Install.op_service(mode, home, "tai-scheduler.service", "scheduler-service", body)
|
||||
|
||||
@staticmethod
|
||||
def op_telegram_service(mode, home):
|
||||
if mode == "uninstall":
|
||||
_done, detail = Install.remove_service_unit(home, "tai-telegram.service")
|
||||
return "telegram-service: %s" % detail
|
||||
unit = os.path.join(Install.unit_dir(home), "tai-telegram.service")
|
||||
if mode == "install" and os.path.isfile(unit):
|
||||
active, enabled = Install.service_state("tai-telegram.service")
|
||||
return "telegram-service: already installed (%s, %s)" % (active, enabled)
|
||||
token = os.environ.get("TELEGRAM_BOT_TOKEN") or TelegramBot.load_telegram_token(home)
|
||||
if not token:
|
||||
return "telegram-service: needs TELEGRAM_BOT_TOKEN in the environment"
|
||||
try:
|
||||
me = Telegram(token).call("getMe", {})
|
||||
except (OSError, ValueError) as exc:
|
||||
return "telegram-service: token rejected: %s" % Text.short_error(exc)
|
||||
env_path = TelegramBot.telegram_env_path(home)
|
||||
try:
|
||||
with open(env_path, "w", encoding="utf-8") as handle:
|
||||
handle.write("TELEGRAM_BOT_TOKEN=%s\n" % token)
|
||||
os.chmod(env_path, 0o600)
|
||||
except OSError as exc:
|
||||
return "telegram-service: env write failed: %s" % Text.short_error(exc)
|
||||
body = Consts.TELEGRAM_UNIT % (Install.service_exec(home, os.path.abspath(Consts.TAI_SCRIPT)), env_path)
|
||||
if mode == "reinstall" and shutil.which("systemctl"):
|
||||
Install.run_host(["systemctl", "--user", "disable", "--now", "tai-telegram.service"])
|
||||
try:
|
||||
os.makedirs(os.path.dirname(unit), exist_ok=True)
|
||||
with open(unit, "w", encoding="utf-8") as handle:
|
||||
handle.write(body)
|
||||
except OSError as exc:
|
||||
return "telegram-service: unit write failed: %s" % Text.short_error(exc)
|
||||
who = (me.get("result") or {}).get("username", "?")
|
||||
return "telegram-service: @%s, %s" % (who, Install.activate_service("tai-telegram.service"))
|
||||
|
||||
@staticmethod
|
||||
def op_container(mode, _home):
|
||||
engine = Sandbox.container_engine()
|
||||
if not engine:
|
||||
return "container: no engine (install podman)"
|
||||
if mode == "uninstall":
|
||||
Install.run_host([engine, "rm", "-f", Consts.BOX_NAME])
|
||||
ok, _out = Install.run_host([engine, "rmi", "-f", Consts.BOX_IMAGE])
|
||||
return "container: removed%s" % (", image purged" if ok else "")
|
||||
if mode == "reinstall":
|
||||
Install.run_host([engine, "rm", "-f", Consts.BOX_NAME])
|
||||
Install.run_host([engine, "rmi", "-f", Consts.BOX_IMAGE])
|
||||
try:
|
||||
engine = Sandbox.ensure_box()
|
||||
except BackendError as exc:
|
||||
return "container: failed: %s" % exc
|
||||
state = Sandbox.box_state(engine)
|
||||
detail = ", %s" % Install.box_venv_state(engine) if state == "running" else ""
|
||||
if mode == "reinstall":
|
||||
return "container: rebuilt and %s%s" % (state, detail)
|
||||
if mode == "install" and state == "running":
|
||||
return "container: already running%s" % detail
|
||||
return "container: %s%s" % (state, detail)
|
||||
|
||||
@staticmethod
|
||||
def install_report(action, targets, home, store=None):
|
||||
if action == "status":
|
||||
return "\n".join(Install.status_lines(home, store))
|
||||
lines = []
|
||||
for target in targets:
|
||||
if target == "binary":
|
||||
lines.append(Install.op_binary(action, home))
|
||||
elif target == "bash-hook":
|
||||
lines.append(Install.op_hook(action, home))
|
||||
elif target == "zsh-hook":
|
||||
lines.append(Install.op_zsh_hook(action, home))
|
||||
elif target == "fish-hook":
|
||||
lines.append(Install.op_fish_hook(action, home))
|
||||
elif target == "plexer-hook":
|
||||
lines.append(Install.op_plexer_hook(action, home))
|
||||
elif target == "zsh-plexer-hook":
|
||||
lines.append(Install.op_zsh_plexer_hook(action, home))
|
||||
elif target == "fish-plexer-hook":
|
||||
lines.append(Install.op_fish_plexer_hook(action, home))
|
||||
elif target == "venv":
|
||||
lines.append(Install.op_venv(action, home))
|
||||
elif target == "scheduler-service":
|
||||
lines.append(Install.op_scheduler_service(action, home))
|
||||
elif target == "telegram-service":
|
||||
lines.append(Install.op_telegram_service(action, home))
|
||||
elif target == "container":
|
||||
lines.append(Install.op_container(action, home))
|
||||
lines.append(Install.vault_line(home, store))
|
||||
return "\n".join(lines)
|
||||
|
||||
@staticmethod
|
||||
def install_telegram():
|
||||
token = os.environ.get("TELEGRAM_BOT_TOKEN") or ""
|
||||
if token:
|
||||
print("using TELEGRAM_BOT_TOKEN from environment")
|
||||
else:
|
||||
try:
|
||||
token = getpass.getpass("Telegram bot token: ").strip()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
print("\naborted")
|
||||
return 1
|
||||
if not token:
|
||||
print("empty token, aborted")
|
||||
return 1
|
||||
try:
|
||||
me = Telegram(token).call("getMe", {})
|
||||
except (OSError, ValueError) as exc:
|
||||
print("token rejected: %s" % Text.short_error(exc))
|
||||
return 1
|
||||
print("token ok for @%s" % (me.get("result") or {}).get("username", "?"))
|
||||
home = os.environ.get("TAI_HOME") or os.path.join(os.path.expanduser("~"), ".tai")
|
||||
os.makedirs(home, exist_ok=True)
|
||||
env_path = TelegramBot.telegram_env_path(home)
|
||||
with open(env_path, "w", encoding="utf-8") as handle:
|
||||
handle.write("TELEGRAM_BOT_TOKEN=%s\n" % token)
|
||||
try:
|
||||
os.chmod(env_path, 0o600)
|
||||
except OSError:
|
||||
pass
|
||||
print("wrote %s" % env_path)
|
||||
try:
|
||||
Sandbox.ensure_box()
|
||||
except BackendError as exc:
|
||||
print("container failed: %s" % exc)
|
||||
return 1
|
||||
unit_folder = os.path.join(os.path.expanduser("~"), ".config", "systemd", "user")
|
||||
os.makedirs(unit_folder, exist_ok=True)
|
||||
unit_path = os.path.join(unit_folder, "tai-telegram.service")
|
||||
with open(unit_path, "w", encoding="utf-8") as handle:
|
||||
handle.write(Consts.TELEGRAM_UNIT % (Install.service_exec(home, os.path.abspath(Consts.TAI_SCRIPT)), env_path))
|
||||
print("wrote %s" % unit_path)
|
||||
print("service: %s" % Install.activate_service("tai-telegram.service"))
|
||||
return 0
|
||||
|
||||
@staticmethod
|
||||
def uninstall_telegram():
|
||||
unit_path = os.path.join(os.path.expanduser("~"), ".config", "systemd", "user", "tai-telegram.service")
|
||||
if shutil.which("systemctl"):
|
||||
Install.run_host(["systemctl", "--user", "disable", "--now", "tai-telegram.service"])
|
||||
if os.path.exists(unit_path):
|
||||
try:
|
||||
os.remove(unit_path)
|
||||
print("removed %s" % unit_path)
|
||||
except OSError as exc:
|
||||
print("unit remove failed: %s" % Text.short_error(exc))
|
||||
else:
|
||||
print("no service unit found")
|
||||
if shutil.which("systemctl"):
|
||||
Install.run_host(["systemctl", "--user", "daemon-reload"])
|
||||
engine = Sandbox.container_engine()
|
||||
if not engine:
|
||||
print("no container engine found")
|
||||
return 0
|
||||
ok, _out = Install.run_host([engine, "rm", "-f", Consts.BOX_NAME])
|
||||
print("container remove: %s" % ("ok" if ok else "not present"))
|
||||
ok, _out = Install.run_host([engine, "rmi", "-f", Consts.BOX_IMAGE])
|
||||
print("image purge: %s" % ("ok" if ok else "not present"))
|
||||
return 0
|
||||
|
||||
@staticmethod
|
||||
def upsert_bashrc_block(path):
|
||||
return Install._rc_block_upsert(path, Consts.BASHRC_MARK_BEGIN, Consts.BASHRC_MARK_END, Consts.BASHRC_BLOCK)
|
||||
|
||||
@staticmethod
|
||||
def install_self():
|
||||
home = os.environ.get("TAI_HOME") or os.path.join(os.path.expanduser("~"), ".tai")
|
||||
print(Install.op_binary("upgrade", home))
|
||||
print(Install.op_hook("install", home))
|
||||
# bash-hook always goes in, unconditionally, for back-compat -- it's harmless clutter
|
||||
# in ~/.bashrc for a zsh/fish user who never touches bash. $SHELL additionally gets its
|
||||
# own matching hook so the one-shot CLI bootstrap (--install) actually works out of the
|
||||
# box on a non-bash login shell, not just when the user later runs /install zsh-hook.
|
||||
login_shell = os.path.basename(os.environ.get("SHELL") or "")
|
||||
rc = "~/.bashrc"
|
||||
if login_shell == "zsh":
|
||||
print(Install.op_zsh_hook("install", home))
|
||||
rc = "~/.zshrc"
|
||||
elif login_shell == "fish":
|
||||
print(Install.op_fish_hook("install", home))
|
||||
rc = "~/.config/fish/config.fish"
|
||||
print(Install.op_venv("install", home))
|
||||
print("restart your shell or run: source %s" % rc)
|
||||
return 0
|
||||
@@ -0,0 +1,246 @@
|
||||
"""Model health tracking, scoring, and roster probing."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
from tai_pkg.claude import Claude
|
||||
from tai_pkg.codex import Codex
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.gemini import Gemini
|
||||
from tai_pkg.grok import Grok
|
||||
from tai_pkg.muse import Muse
|
||||
from tai_pkg.opencode import Opencode
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["ModelHealth", "ModelScoring"]
|
||||
|
||||
|
||||
class ModelHealth:
|
||||
def __init__(self, home):
|
||||
self.path = os.path.join(home, Consts.MODEL_HEALTH_FILE) if home else ""
|
||||
self.data = {"models": {}, "lists": {}}
|
||||
if self.path:
|
||||
try:
|
||||
with open(self.path, "r", encoding="utf-8") as handle:
|
||||
raw = json.load(handle)
|
||||
if isinstance(raw, dict):
|
||||
if isinstance(raw.get("models"), dict):
|
||||
self.data["models"] = raw["models"]
|
||||
if isinstance(raw.get("lists"), dict):
|
||||
self.data["lists"] = raw["lists"]
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
|
||||
def _save(self):
|
||||
if not self.path:
|
||||
return
|
||||
try:
|
||||
tmp = self.path + ".tmp"
|
||||
with open(tmp, "w", encoding="utf-8") as handle:
|
||||
json.dump(self.data, handle)
|
||||
os.replace(tmp, self.path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def entry(self, provider, model):
|
||||
return self.data["models"].get("%s|%s" % (provider or "default", model or "")) or {}
|
||||
|
||||
def weight(self, provider, model):
|
||||
return ModelScoring.model_weight(self.entry(provider, model))
|
||||
|
||||
def circuit_open(self, provider, model):
|
||||
return (self.entry(provider, model).get("until") or 0.0) > time.time()
|
||||
|
||||
def record(self, provider, model, success, latency_ms=None, tokens_per_second=None):
|
||||
if not model:
|
||||
return
|
||||
key = "%s|%s" % (provider or "default", model or "")
|
||||
entry = self.data["models"].setdefault(key, {"s": 0, "f": 0, "r": 0.0, "lat": 0.0, "tps": 0.0, "cf": 0, "until": 0.0, "used": 0.0, "n": 0})
|
||||
entry["n"] += 1
|
||||
entry["used"] = time.time()
|
||||
if success:
|
||||
reward = Consts.MODEL_NEUTRAL_REWARD if latency_ms is None and tokens_per_second is None else ModelScoring.model_outcome_reward(latency_ms, tokens_per_second)
|
||||
entry["s"] += 1
|
||||
entry["cf"] = 0
|
||||
entry["until"] = 0.0
|
||||
entry["r"] += reward
|
||||
if latency_ms is not None:
|
||||
entry["lat"] += latency_ms
|
||||
if tokens_per_second is not None:
|
||||
entry["tps"] += tokens_per_second
|
||||
else:
|
||||
entry["f"] += 1
|
||||
entry["cf"] += 1
|
||||
if entry["cf"] >= Consts.MODEL_CIRCUIT_THRESHOLD:
|
||||
entry["until"] = time.time() + Consts.MODEL_CIRCUIT_COOLDOWN
|
||||
self._save()
|
||||
|
||||
def cached_list(self, name, ttl=Consts.OPENROUTER_FREE_CACHE_TTL):
|
||||
row = self.data["lists"].get(name)
|
||||
if not isinstance(row, dict):
|
||||
return None
|
||||
models = row.get("models")
|
||||
if not models or (time.time() - (row.get("at") or 0)) > ttl:
|
||||
return None
|
||||
return [str(item) for item in models if str(item).strip()]
|
||||
|
||||
def store_list(self, name, models):
|
||||
cleaned = [str(item).strip() for item in models if str(item).strip()]
|
||||
if not cleaned:
|
||||
return
|
||||
self.data["lists"][name] = {"at": time.time(), "models": cleaned}
|
||||
self._save()
|
||||
|
||||
def snapshot(self):
|
||||
rows = []
|
||||
for key, entry in self.data["models"].items():
|
||||
provider, _, model = key.partition("|")
|
||||
ok = entry.get("s") or 0
|
||||
rows.append({
|
||||
"provider": provider,
|
||||
"model": model,
|
||||
"weight": round(ModelScoring.model_weight(entry), 4),
|
||||
"ok": ok,
|
||||
"fail": entry.get("f") or 0,
|
||||
"avg_ms": round((entry.get("lat") or 0.0) / ok) if ok and entry.get("lat") else None,
|
||||
"avg_tps": round((entry.get("tps") or 0.0) / ok, 1) if ok and entry.get("tps") else None,
|
||||
"circuit_open": (entry.get("until") or 0.0) > time.time(),
|
||||
"used": entry.get("used") or 0.0,
|
||||
})
|
||||
rows.sort(key=lambda row: (-row["weight"], row["provider"], row["model"]))
|
||||
return rows
|
||||
|
||||
|
||||
class ModelScoring:
|
||||
"""ModelScoring."""
|
||||
|
||||
@staticmethod
|
||||
def model_speed_reward(tokens_per_second):
|
||||
if not tokens_per_second or tokens_per_second <= 0:
|
||||
return Consts.MODEL_NEUTRAL_REWARD
|
||||
return tokens_per_second / (tokens_per_second + Consts.MODEL_SPEED_REF_TPS)
|
||||
|
||||
@staticmethod
|
||||
def model_latency_reward(latency_ms):
|
||||
if not latency_ms or latency_ms <= 0:
|
||||
return 1.0
|
||||
return 1.0 / (1.0 + latency_ms / Consts.MODEL_LATENCY_REF_MS)
|
||||
|
||||
@staticmethod
|
||||
def model_outcome_reward(latency_ms=None, tokens_per_second=None):
|
||||
return ModelScoring.model_speed_reward(tokens_per_second) * ModelScoring.model_latency_reward(latency_ms)
|
||||
|
||||
@staticmethod
|
||||
def model_weight(entry):
|
||||
entry = entry or {}
|
||||
attempts = (entry.get("s") or 0) + (entry.get("f") or 0)
|
||||
return ((entry.get("r") or 0.0) + Consts.MODEL_NEUTRAL_REWARD * Consts.MODEL_WEIGHT_PRIOR) / (attempts + Consts.MODEL_WEIGHT_PRIOR)
|
||||
|
||||
@staticmethod
|
||||
def fetch_model_ids(models_url, key="", extra_headers=None, timeout=Consts.MODEL_LIST_TIMEOUT):
|
||||
headers = dict(extra_headers or {})
|
||||
if key:
|
||||
headers["Authorization"] = "Bearer " + key
|
||||
try:
|
||||
request = urllib.request.Request(models_url, headers=headers, method="GET")
|
||||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||||
payload = json.loads(response.read().decode("utf-8", "replace"))
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
rows = payload.get("data") if isinstance(payload, dict) else None
|
||||
if not isinstance(rows, list):
|
||||
return None
|
||||
ids = [str(row.get("id")).strip() for row in rows if isinstance(row, dict) and str(row.get("id") or "").strip()]
|
||||
return ids or None
|
||||
|
||||
@staticmethod
|
||||
def embedded_failure_message(data):
|
||||
if not isinstance(data, dict):
|
||||
return ""
|
||||
error = data.get("error")
|
||||
if isinstance(error, dict):
|
||||
message = error.get("message") or error.get("type") or ""
|
||||
return str(message).strip()
|
||||
if isinstance(error, str) and error.strip():
|
||||
return error.strip()
|
||||
if data.get("type") == "error":
|
||||
return str(data.get("message") or "upstream error").strip()
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def prefer_muse(models):
|
||||
return sorted(models, key=lambda item: 0 if Consts.MUSE_LABEL in item.lower() else 1)
|
||||
|
||||
@staticmethod
|
||||
def openrouter_free_tier(cand):
|
||||
model = (cand.get("model") or "").lower()
|
||||
muse = 0 if Consts.MUSE_LABEL in model else 1
|
||||
if cand.get("provider") == Consts.OPENCODE_LABEL:
|
||||
return muse
|
||||
if cand.get("provider") == Consts.DEVPLACE_LABEL:
|
||||
return 2 + muse
|
||||
if cand.get("provider") == Consts.OPENROUTER_LABEL and Text.is_free_model(model):
|
||||
return 4 + muse
|
||||
return 6
|
||||
|
||||
@staticmethod
|
||||
def models_report(chat):
|
||||
lines = ["label/model | weight ok/fail avg-ms avg-tps status"]
|
||||
seen = set()
|
||||
for cand in chat.ordered_candidates():
|
||||
key = (cand["provider"], cand["model"])
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
entry = chat.health.entry(cand["provider"], cand["model"])
|
||||
ok = entry.get("s") or 0
|
||||
fail = entry.get("f") or 0
|
||||
avg_ms = str(round((entry.get("lat") or 0.0) / ok)) if ok and entry.get("lat") else "-"
|
||||
avg_tps = str(round((entry.get("tps") or 0.0) / ok, 1)) if ok and entry.get("tps") else "-"
|
||||
if cand["chat_url"] == Consts.MUSE_CHAT_URL:
|
||||
status = Muse.muse_health()
|
||||
elif chat.health.circuit_open(cand["provider"], cand["model"]):
|
||||
status = "circuit open"
|
||||
elif not ok and not fail:
|
||||
status = "untested"
|
||||
else:
|
||||
status = "ready"
|
||||
lines.append("%s/%s | %.4f %d/%d %s %s %s" % (cand["provider"], cand["model"], ModelScoring.model_weight(entry), ok, fail, avg_ms, avg_tps, status))
|
||||
return "\n".join(lines)
|
||||
|
||||
@staticmethod
|
||||
def probe_backend(models_url, key):
|
||||
if models_url == Consts.MUSE_MODELS_URL:
|
||||
return Muse.muse_health()
|
||||
if models_url == Consts.GROK_MODELS_URL:
|
||||
return Grok.grok_health()
|
||||
if models_url == Consts.CLAUDE_MODELS_URL:
|
||||
return Claude.claude_health()
|
||||
if models_url == Consts.CODEX_MODELS_URL:
|
||||
return Codex.codex_health()
|
||||
if models_url == Consts.GEMINI_MODELS_URL:
|
||||
return Gemini.gemini_health()
|
||||
if models_url == Consts.OPENCODE_MODELS_URL:
|
||||
return "ok, %s" % Opencode.opencode_binary() if Opencode.opencode_binary() else "opencode CLI not found"
|
||||
request = urllib.request.Request(models_url, headers={"Authorization": "Bearer " + key})
|
||||
started = time.time()
|
||||
try:
|
||||
with urllib.request.urlopen(request, timeout=6) as response:
|
||||
try:
|
||||
payload = json.loads(response.read().decode("utf-8", "replace"))
|
||||
except ValueError:
|
||||
return "bad response"
|
||||
data = payload.get("data", []) if isinstance(payload, dict) else []
|
||||
count = len(data) if isinstance(data, list) else 0
|
||||
return "ok, %d models, %dms" % (count, (time.time() - started) * 1000)
|
||||
except urllib.error.HTTPError as exc:
|
||||
if exc.code == 401:
|
||||
return "needs key"
|
||||
return "http %d" % exc.code
|
||||
except OSError as exc:
|
||||
return "fail: %s" % Text.short_error(exc)[:60]
|
||||
except Exception as exc:
|
||||
return "fail: %s" % Text.short_error(exc)[:60]
|
||||
+342
@@ -0,0 +1,342 @@
|
||||
"""Muse CLI backend and persistent session map."""
|
||||
|
||||
import glob
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
|
||||
from tai_pkg.clirun import Cli
|
||||
from tai_pkg.compat import fcntl
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.errors import BackendError
|
||||
from tai_pkg.release import Release
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["Muse", "_MUSE_THREAD_LOCKS", "_MUSE_THREAD_GUARD"]
|
||||
|
||||
|
||||
class Muse:
|
||||
"""Muse."""
|
||||
|
||||
@staticmethod
|
||||
def muse_binary():
|
||||
return os.environ.get("MUSE_BIN") or shutil.which("muse")
|
||||
|
||||
@staticmethod
|
||||
def muse_output_schema():
|
||||
call = {"type": "object", "additionalProperties": False, "properties": {"name": {"type": "string"}, "arguments": {"type": "string"}}, "required": ["name", "arguments"]}
|
||||
return {"type": "object", "additionalProperties": False, "properties": {"content": {"type": "string"}, "tool_calls": {"type": "array", "items": call}}, "required": ["content", "tool_calls"]}
|
||||
|
||||
@staticmethod
|
||||
def muse_prompt(messages, tools):
|
||||
parts = []
|
||||
for message in messages or []:
|
||||
role = (message.get("role") or "user").upper()
|
||||
body = message.get("content") or ""
|
||||
if isinstance(body, list):
|
||||
body = " ".join(str(part.get("text") or "") for part in body if isinstance(part, dict))
|
||||
if message.get("tool_call_id"):
|
||||
parts.append("%s %s:\n%s" % (role, message["tool_call_id"], body))
|
||||
else:
|
||||
parts.append("%s:\n%s" % (role, body))
|
||||
for call in message.get("tool_calls") or []:
|
||||
func = call.get("function") or {}
|
||||
parts.append("TOOL CALL %s:\n%s %s" % (call.get("id") or "", func.get("name") or "", func.get("arguments") or ""))
|
||||
if tools:
|
||||
parts.append("TOOLS:\n%s" % json.dumps(tools))
|
||||
parts.append("Answer with content text plus tool_calls for every tool to invoke now. Empty tool_calls ends the turn.")
|
||||
return "\n\n".join(part for part in parts if part.strip())
|
||||
|
||||
@staticmethod
|
||||
def muse_argv(config, schema_path, prompt_path, session_id=None):
|
||||
argv = [Muse.muse_binary(), "exec", "--json", "--max-model-steps", str(Consts.MUSE_MAX_STEPS), "--disable-web-tools", "--disable-reminders", "--allow-workspace-switch", "--output-schema", schema_path, "--prompt-file", prompt_path]
|
||||
if session_id is None:
|
||||
argv.append("--no-session-log")
|
||||
else:
|
||||
argv.extend(["--session-id", session_id])
|
||||
if config.explicit_model:
|
||||
argv.extend(["--model", config.explicit_model])
|
||||
return argv
|
||||
|
||||
@staticmethod
|
||||
def muse_write_file(path, text):
|
||||
with open(path, "w", encoding="utf-8") as handle:
|
||||
handle.write(text)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
|
||||
@staticmethod
|
||||
def muse_message_hash(message):
|
||||
blob = json.dumps(message, sort_keys=True, separators=(",", ":"), default=str)
|
||||
return hashlib.sha256(blob.encode("utf-8")).hexdigest()
|
||||
|
||||
@staticmethod
|
||||
def muse_map_path(home):
|
||||
return os.path.join(home, Consts.MUSE_SESSIONS_FILE)
|
||||
|
||||
@staticmethod
|
||||
def muse_load_map(home):
|
||||
path = Muse.muse_map_path(home)
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
data = json.load(handle)
|
||||
except FileNotFoundError:
|
||||
return {}
|
||||
except (OSError, ValueError):
|
||||
Release.quarantine(path)
|
||||
return {}
|
||||
if not isinstance(data, dict):
|
||||
Release.quarantine(path)
|
||||
return {}
|
||||
now = time.time()
|
||||
pruned = False
|
||||
for key in list(data):
|
||||
entry = data[key]
|
||||
if not isinstance(entry, dict) or not isinstance(entry.get("id"), str) or not isinstance(entry.get("hashes"), list):
|
||||
del data[key]
|
||||
pruned = True
|
||||
elif key.startswith("worker/"):
|
||||
try:
|
||||
stamp = Text._fromisoformat(entry.get("updated") or "").timestamp()
|
||||
except ValueError:
|
||||
stamp = 0
|
||||
if now - stamp > Consts.MUSE_WORKER_TTL:
|
||||
del data[key]
|
||||
pruned = True
|
||||
if pruned:
|
||||
Muse.muse_save_map(home, data)
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def muse_save_map(home, data):
|
||||
Muse.muse_write_file(Muse.muse_map_path(home), json.dumps(data, sort_keys=True) + "\n")
|
||||
|
||||
@staticmethod
|
||||
def muse_acquire(home, name):
|
||||
with _MUSE_THREAD_GUARD:
|
||||
thread_lock = _MUSE_THREAD_LOCKS.setdefault(name, threading.Lock())
|
||||
thread_lock.acquire()
|
||||
handle = None
|
||||
if fcntl is not None:
|
||||
try:
|
||||
folder = os.path.join(home, Consts.MUSE_LOCKS_DIR)
|
||||
os.makedirs(folder, exist_ok=True)
|
||||
handle = open(os.path.join(folder, hashlib.sha256(name.encode("utf-8")).hexdigest() + ".lock"), "w")
|
||||
fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
|
||||
except OSError:
|
||||
if handle is not None:
|
||||
try:
|
||||
handle.close()
|
||||
except OSError:
|
||||
pass
|
||||
handle = None
|
||||
|
||||
def release():
|
||||
try:
|
||||
if handle is not None:
|
||||
if fcntl is not None:
|
||||
try:
|
||||
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
handle.close()
|
||||
except OSError:
|
||||
pass
|
||||
finally:
|
||||
thread_lock.release()
|
||||
|
||||
return release
|
||||
|
||||
@staticmethod
|
||||
def muse_transact(home, func):
|
||||
release = Muse.muse_acquire(home, "map")
|
||||
try:
|
||||
data = Muse.muse_load_map(home)
|
||||
result = func(data)
|
||||
Muse.muse_save_map(home, data)
|
||||
return result
|
||||
finally:
|
||||
release()
|
||||
|
||||
@staticmethod
|
||||
def muse_session_root():
|
||||
base = os.environ.get("XDG_DATA_HOME") or os.path.join(os.path.expanduser("~"), ".local", "share")
|
||||
return os.path.join(base, "muse", "sessions")
|
||||
|
||||
@staticmethod
|
||||
def muse_session_alive(session_id):
|
||||
try:
|
||||
pattern = os.path.join(Muse.muse_session_root(), "*", "*", "*", session_id, "session.jsonl")
|
||||
return bool(glob.glob(pattern))
|
||||
except (OSError, ValueError):
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def muse_session_split(messages, hashes):
|
||||
if len(hashes) > len(messages):
|
||||
return None
|
||||
for pos, digest in enumerate(hashes):
|
||||
if Muse.muse_message_hash(messages[pos]) != digest:
|
||||
return None
|
||||
return list(messages[len(hashes):])
|
||||
|
||||
@staticmethod
|
||||
def muse_valid_answer(text):
|
||||
try:
|
||||
data = json.loads(text)
|
||||
except ValueError:
|
||||
return None
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
if not isinstance(data.get("content"), str) or not isinstance(data.get("tool_calls"), list):
|
||||
return None
|
||||
for call in data["tool_calls"]:
|
||||
if not isinstance(call, dict) or not isinstance(call.get("name"), str) or not isinstance(call.get("arguments"), str):
|
||||
return None
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def muse_gate(func):
|
||||
def gated(ok):
|
||||
func(ok)
|
||||
|
||||
gated.armed = threading.Event()
|
||||
return gated
|
||||
|
||||
@staticmethod
|
||||
def muse_reap(proc, stderr_done, on_reap):
|
||||
def target():
|
||||
ok = False
|
||||
try:
|
||||
try:
|
||||
for _line in proc.stdout:
|
||||
pass
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
try:
|
||||
proc.stdout.close()
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
try:
|
||||
ok = proc.wait(timeout=Consts.MUSE_REAP_GRACE) == 0
|
||||
except subprocess.TimeoutExpired:
|
||||
try:
|
||||
proc.kill()
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
proc.wait(timeout=10)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
pass
|
||||
try:
|
||||
stderr_done.wait(timeout=10)
|
||||
except (OSError, RuntimeError):
|
||||
pass
|
||||
finally:
|
||||
if on_reap is not None:
|
||||
try:
|
||||
on_reap(ok)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
thread = threading.Thread(target=target, daemon=True)
|
||||
thread.start()
|
||||
return thread
|
||||
|
||||
@staticmethod
|
||||
def muse_run(argv, cwd, stream_sink, timeout, on_reap=None):
|
||||
try:
|
||||
proc = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, cwd=cwd, bufsize=1)
|
||||
except OSError as exc:
|
||||
raise BackendError("muse exec failed: %s" % Text.short_error(exc))
|
||||
gate = getattr(on_reap, "armed", None)
|
||||
if gate is not None:
|
||||
gate.set()
|
||||
stderr_parts = []
|
||||
stderr_done = threading.Event()
|
||||
|
||||
def collect_stderr():
|
||||
try:
|
||||
for chunk in proc.stderr:
|
||||
stderr_parts.append(chunk)
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
finally:
|
||||
stderr_done.set()
|
||||
|
||||
threading.Thread(target=collect_stderr, daemon=True).start()
|
||||
finished = threading.Event()
|
||||
expired = threading.Event()
|
||||
|
||||
def watchdog():
|
||||
if not finished.wait(timeout):
|
||||
expired.set()
|
||||
try:
|
||||
proc.kill()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
threading.Thread(target=watchdog, daemon=True).start()
|
||||
answer = None
|
||||
terminal = None
|
||||
buffer = []
|
||||
try:
|
||||
for line in proc.stdout:
|
||||
text = line.strip()
|
||||
if not text.startswith("{"):
|
||||
continue
|
||||
try:
|
||||
event = json.loads(text)
|
||||
except ValueError:
|
||||
continue
|
||||
kind = event.get("payload_type") or ""
|
||||
payload = event.get("payload") or {}
|
||||
if kind == "run.output.delta" and payload.get("text"):
|
||||
piece = payload["text"]
|
||||
if stream_sink is not None:
|
||||
stream_sink(piece)
|
||||
buffer.append(piece)
|
||||
if answer is None:
|
||||
answer = Muse.muse_valid_answer("".join(buffer))
|
||||
if answer is not None:
|
||||
break
|
||||
elif kind.startswith("run.terminal."):
|
||||
terminal = payload
|
||||
break
|
||||
finally:
|
||||
finished.set()
|
||||
Muse.muse_reap(proc, stderr_done, on_reap)
|
||||
if terminal is not None:
|
||||
if terminal.get("terminal") != "completed":
|
||||
raise BackendError("muse %s: %s" % (terminal.get("terminal") or "failed", (terminal.get("reason") or "unknown")[:300]))
|
||||
parsed = Muse.muse_valid_answer(terminal.get("text") or "")
|
||||
if parsed is None:
|
||||
raise BackendError("muse returned invalid structured answer")
|
||||
return parsed
|
||||
if answer is not None:
|
||||
return answer
|
||||
if expired.is_set():
|
||||
raise BackendError("muse timed out after %ds" % timeout)
|
||||
try:
|
||||
code = proc.returncode
|
||||
except (OSError, ValueError):
|
||||
code = "?"
|
||||
raise BackendError("muse exit %s: %s" % (code, "".join(stderr_parts).strip()[-300:] or "no terminal event"))
|
||||
|
||||
@staticmethod
|
||||
def muse_health():
|
||||
binary = Muse.muse_binary()
|
||||
if not binary:
|
||||
return "missing: muse CLI not found"
|
||||
return Cli._cli_version_only_health(binary, "muse")
|
||||
|
||||
|
||||
_MUSE_THREAD_LOCKS = {}
|
||||
|
||||
|
||||
_MUSE_THREAD_GUARD = threading.Lock()
|
||||
@@ -0,0 +1,122 @@
|
||||
"""Opencode CLI backend, prompts, and response parsing."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.errors import BackendError
|
||||
from tai_pkg.muse import Muse
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["Opencode", "_OPENCODE_ID_ALPHABET", "_OPENCODE_SESSION"]
|
||||
|
||||
|
||||
class Opencode:
|
||||
"""Opencode."""
|
||||
|
||||
@staticmethod
|
||||
def opencode_binary():
|
||||
return os.environ.get("OPENCODE_BIN") or shutil.which("opencode")
|
||||
|
||||
@staticmethod
|
||||
def opencode_prompt(messages, tools):
|
||||
schema = json.dumps(Muse.muse_output_schema())
|
||||
rules = (
|
||||
"You are the model behind another agent. The transcript above is the whole conversation; lines starting with TOOL <id> are results of tool calls that already ran. "
|
||||
"Do not use any of your own tools and do not repeat a tool call whose result is already shown. "
|
||||
"If the results answer the user, put the final answer in content with an empty tool_calls list; otherwise request the next tool calls from TOOLS. "
|
||||
"Reply with one JSON object only, no code fence, matching this schema: "
|
||||
)
|
||||
return Muse.muse_prompt(messages, tools) + "\n\n" + rules + schema
|
||||
|
||||
@staticmethod
|
||||
def opencode_parse(text):
|
||||
text = (text or "").strip()
|
||||
fenced = re.match(r"^```(?:json)?\s*(.*?)\s*```$", text, re.S)
|
||||
if fenced:
|
||||
text = fenced.group(1)
|
||||
# Scan every "{" last-first and decode one value at each spot: free
|
||||
# models often echo the prompt schema or tool JSON before the real
|
||||
# answer, and a first-"{"-to-last-"}" slice spans both into garbage.
|
||||
# Last-first means the model's own trailing answer wins over echoes.
|
||||
decoder = json.JSONDecoder()
|
||||
starts = [pos for pos, char in enumerate(text) if char == "{"]
|
||||
for pos in reversed(starts[-Consts.OPENCODE_PARSE_MAX_SCANS:]):
|
||||
try:
|
||||
answer, _end = decoder.raw_decode(text, pos)
|
||||
except ValueError:
|
||||
continue
|
||||
if isinstance(answer, dict) and ("content" in answer or "tool_calls" in answer):
|
||||
return answer
|
||||
return {"content": text, "tool_calls": []}
|
||||
|
||||
@staticmethod
|
||||
def opencode_parse_events(stdout_text, returncode, stderr_text):
|
||||
"""Shared by opencode_run() (local CLI) and box_opencode_run() (same CLI,
|
||||
executed inside the sandbox box) -- the JSON-lines event stream is
|
||||
identical either way, only how the process gets launched differs."""
|
||||
texts = []
|
||||
for line in stdout_text.splitlines():
|
||||
try:
|
||||
event = json.loads(line)
|
||||
except ValueError:
|
||||
continue
|
||||
if not isinstance(event, dict):
|
||||
continue
|
||||
if event.get("type") == "error":
|
||||
data = (event.get("error") or {}).get("data") or {}
|
||||
raise BackendError("opencode: %s" % (data.get("message") or json.dumps(event.get("error"))[:300]), data.get("statusCode") or 0)
|
||||
part = event.get("part") or {}
|
||||
if event.get("type") == "text" and part.get("text"):
|
||||
texts.append(part["text"])
|
||||
if not texts:
|
||||
raise BackendError("opencode returned no text (exit %d): %s" % (returncode, (stderr_text or stdout_text)[-300:]))
|
||||
return "".join(texts)
|
||||
|
||||
@staticmethod
|
||||
def opencode_run(model, prompt, timeout):
|
||||
binary = Opencode.opencode_binary()
|
||||
if not binary:
|
||||
raise BackendError("opencode CLI not found (OPENCODE_BIN or PATH)")
|
||||
argv = [binary, "run", "--pure", "--agent", Consts.OPENCODE_AGENT, "--format", "json", "-m", "opencode/" + model]
|
||||
with tempfile.TemporaryDirectory(prefix="tai-opencode-") as folder:
|
||||
try:
|
||||
done = subprocess.run(argv, input=prompt, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, cwd=folder, timeout=timeout)
|
||||
except subprocess.TimeoutExpired:
|
||||
raise BackendError("opencode run timed out after %ds" % timeout)
|
||||
except OSError as exc:
|
||||
raise BackendError("opencode run failed: %s" % Text.short_error(exc))
|
||||
return Opencode.opencode_parse_events(done.stdout, done.returncode, done.stderr)
|
||||
|
||||
@staticmethod
|
||||
def opencode_free_models():
|
||||
binary = Opencode.opencode_binary()
|
||||
if not binary:
|
||||
return None
|
||||
try:
|
||||
done = subprocess.run([binary, "models", "opencode"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=Consts.MODEL_LIST_TIMEOUT * 4)
|
||||
except (OSError, subprocess.TimeoutExpired):
|
||||
return None
|
||||
models = [line.strip().split("/", 1)[1] for line in done.stdout.splitlines() if line.strip().startswith("opencode/")]
|
||||
return [item for item in models if Text.is_free_model(item)] or None
|
||||
|
||||
@staticmethod
|
||||
def opencode_headers():
|
||||
return {
|
||||
"User-Agent": os.environ.get("OPENCODE_CLIENT_USER_AGENT") or Consts.OPENCODE_USER_AGENT,
|
||||
"x-opencode-client": os.environ.get("OPENCODE_CLIENT_NAME") or Consts.OPENCODE_CLIENT_NAME,
|
||||
"x-opencode-project": "global",
|
||||
"x-opencode-session": _OPENCODE_SESSION,
|
||||
"x-opencode-request": "msg_" + "".join(random.choice(_OPENCODE_ID_ALPHABET) for _ in range(26)),
|
||||
}
|
||||
|
||||
|
||||
_OPENCODE_ID_ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
||||
|
||||
|
||||
_OPENCODE_SESSION = "ses_" + "".join(random.choice(_OPENCODE_ID_ALPHABET) for _ in range(26))
|
||||
@@ -0,0 +1,570 @@
|
||||
"""Terminal multiplexer controllers and wrapping setup."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
import struct
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from tai_pkg.compat import fcntl, pty, select, termios, tty
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["PlexerError", "PlexerController", "NativePlexerController", "TmuxPlexerController", "PlexerControllerFactory", "PlexerBackend", "NativePlexerBackend", "PlexerSetup"]
|
||||
|
||||
|
||||
class PlexerError(Exception):
|
||||
"""Raised by a PlexerController when the backend can't be reached or refuses a command."""
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Controller side: what Tools.run_plexer_* actually calls. One small
|
||||
# interface, two backends, chosen by which environment variable is set —
|
||||
# TAI_PLEXER_SOCKET (our own NativePlexerBackend already wrapping this shell)
|
||||
# or TMUX_PANE (tmux is already the active manager of this shell). Neither
|
||||
# set means no plexer session is active and the plexer_* tools stay hidden
|
||||
# (see TOOL_ENV_GATES below) rather than erroring.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
class PlexerController:
|
||||
kind = "abstract"
|
||||
|
||||
def send_keys(self, text, enter=False):
|
||||
raise NotImplementedError
|
||||
|
||||
def read_pane(self, lines=200):
|
||||
raise NotImplementedError
|
||||
|
||||
def type_text(self, text, delay=0.03, enter=True):
|
||||
for ch in text:
|
||||
self.send_keys(ch)
|
||||
time.sleep(delay)
|
||||
if enter:
|
||||
self.send_keys("", enter=True)
|
||||
|
||||
def status(self):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class NativePlexerController(PlexerController):
|
||||
kind = "native"
|
||||
|
||||
def __init__(self, socket_path):
|
||||
self.socket_path = socket_path
|
||||
|
||||
def _call(self, payload):
|
||||
try:
|
||||
conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
conn.settimeout(10)
|
||||
conn.connect(self.socket_path)
|
||||
except OSError as exc:
|
||||
raise PlexerError("cannot reach native plexer backend: " + Text.short_error(exc))
|
||||
try:
|
||||
conn.sendall((json.dumps(payload) + "\n").encode("utf-8"))
|
||||
chunks = []
|
||||
while True:
|
||||
chunk = conn.recv(65536)
|
||||
if not chunk:
|
||||
break
|
||||
chunks.append(chunk)
|
||||
if chunk.endswith(b"\n"):
|
||||
break
|
||||
except OSError as exc:
|
||||
raise PlexerError("native plexer backend read/write failed: " + Text.short_error(exc))
|
||||
finally:
|
||||
conn.close()
|
||||
raw = b"".join(chunks).decode("utf-8", "replace").strip()
|
||||
try:
|
||||
reply = json.loads(raw) if raw else {}
|
||||
except ValueError:
|
||||
raise PlexerError("native plexer backend sent an unparseable reply")
|
||||
if not reply.get("ok", False):
|
||||
raise PlexerError(reply.get("error") or "native plexer backend refused the command")
|
||||
return reply.get("data")
|
||||
|
||||
def send_keys(self, text, enter=False):
|
||||
return self._call({"cmd": "send", "text": text, "enter": bool(enter)})
|
||||
|
||||
def read_pane(self, lines=200):
|
||||
return self._call({"cmd": "read", "lines": int(lines)})
|
||||
|
||||
def status(self):
|
||||
return self._call({"cmd": "status"})
|
||||
|
||||
|
||||
class TmuxPlexerController(PlexerController):
|
||||
kind = "tmux"
|
||||
|
||||
def __init__(self, pane):
|
||||
self.pane = pane
|
||||
|
||||
def _tmux(self, *args):
|
||||
try:
|
||||
done = subprocess.run(["tmux"] + list(args), stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=15)
|
||||
except (OSError, subprocess.SubprocessError) as exc:
|
||||
raise PlexerError("tmux call failed: " + Text.short_error(exc))
|
||||
if done.returncode != 0:
|
||||
raise PlexerError("tmux error: " + (done.stderr or "").strip()[:300])
|
||||
return done.stdout
|
||||
|
||||
def send_keys(self, text, enter=False):
|
||||
if text:
|
||||
self._tmux("send-keys", "-t", self.pane, "-l", text)
|
||||
if enter:
|
||||
self._tmux("send-keys", "-t", self.pane, "Enter")
|
||||
return True
|
||||
|
||||
def read_pane(self, lines=200):
|
||||
return self._tmux("capture-pane", "-p", "-t", self.pane, "-S", "-%d" % max(1, int(lines)))
|
||||
|
||||
def status(self):
|
||||
# list-panes, not display-message -p: display-message -p silently
|
||||
# falls back to tmux's default/current pane when -t names a target
|
||||
# that doesn't exist (confirmed empirically -- it still exits 0 with
|
||||
# blank format substitutions), while list-panes -t actually validates
|
||||
# the target and errors like send-keys/capture-pane already do.
|
||||
info = self._tmux("list-panes", "-t", self.pane, "-F", "#{session_name}:#{window_index}.#{pane_index} #{pane_current_command}")
|
||||
return {"backend": "tmux", "pane": self.pane, "info": info.strip()}
|
||||
|
||||
|
||||
class PlexerControllerFactory:
|
||||
"""The whole detection flow lives in one place: which env var is set
|
||||
decides which backend is already managing this shell. Nothing here ever
|
||||
spawns a backend — spawning only happens via `tai.py --plexer-wrap`,
|
||||
invoked by the bashrc hook, and only when neither var is already set."""
|
||||
|
||||
@staticmethod
|
||||
def detect():
|
||||
socket_path = os.environ.get(Consts.PLEXER_SOCKET_ENV)
|
||||
if socket_path:
|
||||
return NativePlexerController(socket_path)
|
||||
pane = os.environ.get("TMUX_PANE")
|
||||
if pane:
|
||||
return TmuxPlexerController(pane)
|
||||
return None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Backend side: only NativePlexerBackend actually exists, because tmux
|
||||
# already IS a backend when it's active — there is nothing for tai to wrap
|
||||
# in that case (see PlexerControllerFactory). PlexerBackend is kept as the
|
||||
# extension point for a future alternative (e.g. a screen-based backend)
|
||||
# without touching the controller side or the detection flow above.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
class PlexerBackend:
|
||||
def run(self):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class NativePlexerBackend(PlexerBackend):
|
||||
"""Owns the pty directly: forks the user's real shell onto one, relays
|
||||
bytes between the real terminal and that pty in both directions, forwards
|
||||
SIGWINCH as a resize, tees everything to a log file, and serves control
|
||||
commands (send/read/status) over a Unix socket so a separate tai process
|
||||
(the agent the user is actually talking to) can drive this session.
|
||||
|
||||
One prefix key (Ctrl+A, PLEXER_PREFIX_KEY) is reserved for tai's own
|
||||
shortcuts in the raw keystroke stream. There is deliberately no
|
||||
tmux-compatible keymap here: when tmux is the active backend, tai never
|
||||
owns the raw keystroke stream at all (see TmuxPlexerController), so there
|
||||
is nothing to avoid colliding with.
|
||||
|
||||
A characteristic worth knowing, not a bug: this relay loop is single-
|
||||
threaded and synchronous, so its own os.write(stdout_fd, ...) applies
|
||||
normal pty backpressure -- if whatever is reading the OTHER end of the
|
||||
outer pty (normally a real terminal emulator, which always reads
|
||||
continuously) ever stops draining for a stretch, a large enough burst of
|
||||
child output can block that write and stall the whole loop until the
|
||||
reader catches up. This is standard for any synchronous pty relay
|
||||
(tmux, ssh, `script` all have the same property) and is why writes
|
||||
headed INTO the child are queued and drained cooperatively rather than
|
||||
written directly (see _pending_to_child) -- that fix was for a
|
||||
self-inflicted deadlock on the inner side; this is normal, external
|
||||
backpressure on the outer side, not something to engineer around.
|
||||
"""
|
||||
|
||||
def __init__(self, home, profile, shell_cmd=None):
|
||||
self.home = home
|
||||
self.profile = profile
|
||||
self.shell_cmd = shell_cmd or [os.environ.get("SHELL") or "/bin/bash"]
|
||||
self.session_id = "%s-%d" % (datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S"), os.getpid())
|
||||
self.socket_path = os.path.join(PlexerSetup.plexer_dir(home), self.session_id + ".sock")
|
||||
self.log_path = os.path.join(PlexerSetup.plexer_dir(home), self.session_id + ".log")
|
||||
self.pid = None
|
||||
self.master_fd = None
|
||||
self._log_fh = None
|
||||
self._server = None
|
||||
self._saved_termios = None
|
||||
self._resize_pending = False
|
||||
self._running = False
|
||||
self._pending_to_child = bytearray()
|
||||
|
||||
def _set_raw(self):
|
||||
if termios is None:
|
||||
return
|
||||
fd = sys.stdin.fileno()
|
||||
try:
|
||||
self._saved_termios = termios.tcgetattr(fd)
|
||||
tty.setraw(fd)
|
||||
except (termios.error, OSError, ValueError):
|
||||
self._saved_termios = None
|
||||
|
||||
def _restore_tty(self):
|
||||
if termios is None or self._saved_termios is None:
|
||||
return
|
||||
try:
|
||||
termios.tcsetattr(sys.stdin.fileno(), termios.TCSADRAIN, self._saved_termios)
|
||||
except (termios.error, OSError):
|
||||
pass
|
||||
|
||||
def _get_winsize(self, fd):
|
||||
try:
|
||||
packed = fcntl.ioctl(fd, termios.TIOCGWINSZ, b"\0" * 8)
|
||||
except OSError:
|
||||
return None
|
||||
rows, cols = struct.unpack("HH", packed[:4])
|
||||
return rows, cols
|
||||
|
||||
def _apply_winsize(self):
|
||||
size = self._get_winsize(sys.stdin.fileno())
|
||||
if size is None or self.master_fd is None:
|
||||
return
|
||||
rows, cols = size
|
||||
try:
|
||||
fcntl.ioctl(self.master_fd, termios.TIOCSWINSZ, struct.pack("HHHH", rows, cols, 0, 0))
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def _on_sigwinch(self, _signum, _frame):
|
||||
self._resize_pending = True
|
||||
|
||||
def _spawn_child(self):
|
||||
pid, master_fd = pty.fork()
|
||||
if pid == 0:
|
||||
env = dict(os.environ)
|
||||
env[Consts.PLEXER_ACTIVE_ENV] = "1"
|
||||
env[Consts.PLEXER_SOCKET_ENV] = self.socket_path
|
||||
try:
|
||||
os.execvpe(self.shell_cmd[0], self.shell_cmd, env)
|
||||
except OSError as exc:
|
||||
sys.stderr.write("tai plexer: failed to exec %s: %s\n" % (self.shell_cmd[0], exc))
|
||||
os._exit(127)
|
||||
return pid, master_fd
|
||||
|
||||
def _open_server(self):
|
||||
try:
|
||||
os.unlink(self.socket_path)
|
||||
except OSError:
|
||||
pass
|
||||
server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
server.bind(self.socket_path)
|
||||
# Same reasoning as the log file: a plain AF_UNIX bind() creates the
|
||||
# socket file with umask-default permissions, often world-connectable
|
||||
# on a multi-user machine -- which would let another local user send
|
||||
# keystrokes into this terminal session. Owner-only, no exceptions.
|
||||
try:
|
||||
os.chmod(self.socket_path, 0o600)
|
||||
except OSError:
|
||||
pass
|
||||
server.listen(8)
|
||||
server.setblocking(False)
|
||||
return server
|
||||
|
||||
def _recv_request(self, conn):
|
||||
"""Mirrors NativePlexerController._call()'s own framing exactly: loop
|
||||
until a newline-terminated chunk arrives, not a single recv() call.
|
||||
json.dumps always escapes a real newline inside a string as the two
|
||||
characters '\\n', so the one literal 0x0A byte in the whole message
|
||||
is always the trailing terminator added below -- a single recv() would
|
||||
silently truncate any request too large for one read (e.g. a long
|
||||
plexer_type payload), which is exactly the kind of bug that only shows
|
||||
up under real load, not in a quick manual check."""
|
||||
chunks = []
|
||||
total = 0
|
||||
while True:
|
||||
chunk = conn.recv(65536)
|
||||
if not chunk:
|
||||
break
|
||||
chunks.append(chunk)
|
||||
total += len(chunk)
|
||||
if chunk.endswith(b"\n") or total >= Consts.PLEXER_RECV_MAX:
|
||||
break
|
||||
return b"".join(chunks)
|
||||
|
||||
def _handle_control(self, conn):
|
||||
try:
|
||||
conn.settimeout(2)
|
||||
raw = self._recv_request(conn)
|
||||
request = json.loads(raw.decode("utf-8", "replace")) if raw.strip() else {}
|
||||
if not isinstance(request, dict):
|
||||
raise ValueError("request must be a JSON object")
|
||||
reply = self._dispatch(request)
|
||||
except Exception as exc:
|
||||
# Deliberately broad: a malformed or unexpected request must never
|
||||
# be able to crash this relay loop and take the user's whole live
|
||||
# terminal session down with it. This is the one place in the
|
||||
# whole backend where "catch everything" is the correct choice,
|
||||
# not laziness -- everything else in this class keeps narrow,
|
||||
# specific except clauses on purpose.
|
||||
reply = {"ok": False, "error": Text.short_error(exc)}
|
||||
try:
|
||||
conn.sendall((json.dumps(reply) + "\n").encode("utf-8"))
|
||||
except OSError:
|
||||
pass
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _dispatch(self, request):
|
||||
cmd = request.get("cmd")
|
||||
if cmd == "send":
|
||||
text = str(request.get("text") or "")
|
||||
if text:
|
||||
self._pending_to_child += text.encode("utf-8", "replace")
|
||||
if request.get("enter"):
|
||||
self._pending_to_child += b"\r"
|
||||
return {"ok": True, "data": True}
|
||||
if cmd == "read":
|
||||
limit = max(1, min(int(request.get("lines") or 200), 5000))
|
||||
return {"ok": True, "data": self._tail_log(limit)}
|
||||
if cmd == "status":
|
||||
return {"ok": True, "data": {"backend": "native", "pid": self.pid, "log": self.log_path, "socket": self.socket_path}}
|
||||
return {"ok": False, "error": "unknown command %r" % cmd}
|
||||
|
||||
def _tail_log(self, lines, chunk_size=65536, max_bytes=1 << 24):
|
||||
"""Reads only the tail, in growing byte chunks from the end, instead
|
||||
of loading the whole file -- a session log is exactly the kind of
|
||||
file this feature encourages to grow large over a long-running
|
||||
terminal session, and read_pane()/plexer_read get called repeatedly,
|
||||
so re-reading the entire file on every call would make the backend
|
||||
slower the longer the very session it's meant to support stays open."""
|
||||
try:
|
||||
size = os.path.getsize(self.log_path)
|
||||
except OSError:
|
||||
return ""
|
||||
with open(self.log_path, "rb") as handle:
|
||||
read_from = size
|
||||
data = b""
|
||||
while True:
|
||||
read_from = max(0, read_from - chunk_size)
|
||||
handle.seek(read_from)
|
||||
data = handle.read(size - read_from)
|
||||
if data.count(b"\n") >= lines or read_from == 0 or len(data) >= max_bytes:
|
||||
break
|
||||
text = data.decode("utf-8", "replace")
|
||||
return "\n".join(text.splitlines()[-lines:])
|
||||
|
||||
def run(self):
|
||||
if pty is None or termios is None or select is None:
|
||||
return "error: pty/termios/select not available on this platform"
|
||||
self.pid, self.master_fd = self._spawn_child()
|
||||
# The log captures the entire terminal session verbatim -- anything typed,
|
||||
# including passwords and tokens -- so it must never be created world- or
|
||||
# group-readable. os.open's mode is applied atomically at creation (no
|
||||
# window where a looser default permission briefly exists), unlike a
|
||||
# plain open() + a later os.chmod.
|
||||
log_fd = os.open(self.log_path, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o600)
|
||||
self._log_fh = os.fdopen(log_fd, "ab", buffering=0)
|
||||
self._server = self._open_server()
|
||||
self._set_raw()
|
||||
self._apply_winsize()
|
||||
try:
|
||||
old_handler = signal.signal(signal.SIGWINCH, self._on_sigwinch) if hasattr(signal, "SIGWINCH") else None
|
||||
except (ValueError, OSError):
|
||||
old_handler = None
|
||||
self._running = True
|
||||
crash = None
|
||||
try:
|
||||
self._relay_loop()
|
||||
except Exception as exc: # noqa: BLE001 -- see comment below
|
||||
# Deliberately broad, and deliberately not re-raised bare: letting
|
||||
# an unexpected exception propagate straight out of run() would
|
||||
# skip the waitpid() below entirely (a bare `finally` only runs
|
||||
# cleanup, it does not stop the exception from continuing past
|
||||
# it), orphaning the wrapped shell with no supervisor left to
|
||||
# reap it. The user's actual terminal child process must not be
|
||||
# allowed to leak just because something in the relay loop broke.
|
||||
crash = exc
|
||||
finally:
|
||||
self._running = False
|
||||
if old_handler is not None:
|
||||
try:
|
||||
signal.signal(signal.SIGWINCH, old_handler)
|
||||
except (ValueError, OSError):
|
||||
pass
|
||||
self._restore_tty()
|
||||
try:
|
||||
self._log_fh.close()
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
self._server.close()
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
os.unlink(self.socket_path)
|
||||
except OSError:
|
||||
pass
|
||||
if crash is not None:
|
||||
try:
|
||||
os.killpg(self.pid, signal.SIGKILL)
|
||||
except OSError:
|
||||
try:
|
||||
os.kill(self.pid, signal.SIGKILL)
|
||||
except OSError:
|
||||
pass
|
||||
sys.stderr.write("tai plexer: wrapper crashed (%s), terminating the session\n" % Text.short_error(crash))
|
||||
try:
|
||||
_, status = os.waitpid(self.pid, 0)
|
||||
return os.waitstatus_to_exitcode(status) if hasattr(os, "waitstatus_to_exitcode") else (1 if crash else 0)
|
||||
except OSError:
|
||||
return 1 if crash else 0
|
||||
|
||||
def _relay_loop(self):
|
||||
"""Every byte headed INTO the child goes through self._pending_to_child,
|
||||
drained here in bounded chunks only when select() says master_fd is
|
||||
writable -- never a direct, unbounded os.write(self.master_fd, ...)
|
||||
anywhere, including from the keystroke path below. A real deadlock
|
||||
was found and fixed here: a large plexer_send payload written
|
||||
directly and synchronously into the pty (echo-enabled, so the write
|
||||
itself generates output) filled the pty's buffer while the only
|
||||
thing that could ever drain that output -- this very loop -- was
|
||||
itself blocked inside that same write call. Queuing and draining
|
||||
cooperatively, alongside reading master_fd's output on the same
|
||||
select() pass, is what actually avoids that, not a bigger buffer."""
|
||||
stdin_fd = sys.stdin.fileno()
|
||||
stdout_fd = sys.stdout.fileno()
|
||||
pending_prefix = False
|
||||
while True:
|
||||
if self._resize_pending:
|
||||
self._resize_pending = False
|
||||
self._apply_winsize()
|
||||
watch_read = [stdin_fd, self.master_fd, self._server.fileno()]
|
||||
watch_write = [self.master_fd] if self._pending_to_child else []
|
||||
try:
|
||||
readable, writable, _ = select.select(watch_read, watch_write, [], 1.0)
|
||||
except InterruptedError:
|
||||
continue
|
||||
if not readable and not writable and not self._process_alive():
|
||||
break
|
||||
if self._server.fileno() in readable:
|
||||
try:
|
||||
conn, _addr = self._server.accept()
|
||||
self._handle_control(conn)
|
||||
except OSError:
|
||||
pass
|
||||
if self.master_fd in writable and self._pending_to_child:
|
||||
try:
|
||||
written = os.write(self.master_fd, bytes(self._pending_to_child[:4096]))
|
||||
del self._pending_to_child[:written]
|
||||
except OSError:
|
||||
self._pending_to_child = bytearray()
|
||||
if self.master_fd in readable:
|
||||
try:
|
||||
data = os.read(self.master_fd, 65536)
|
||||
except OSError:
|
||||
data = b""
|
||||
if not data:
|
||||
break
|
||||
self._log_fh.write(data)
|
||||
os.write(stdout_fd, data)
|
||||
if stdin_fd in readable:
|
||||
try:
|
||||
data = os.read(stdin_fd, 4096)
|
||||
except OSError:
|
||||
data = b""
|
||||
if not data:
|
||||
break
|
||||
data, pending_prefix = self._strip_prefix_commands(data, pending_prefix)
|
||||
if data:
|
||||
self._pending_to_child += data
|
||||
if not self._process_alive():
|
||||
break
|
||||
|
||||
def _strip_prefix_commands(self, data, pending_prefix):
|
||||
"""Intercepts tai's own Ctrl+A shortcuts; everything else passes through untouched.
|
||||
|
||||
'q' quits the wrapped session. Verified empirically (see
|
||||
tests/plexer.py) that in this environment an interactive bash running
|
||||
as a pty.fork() session leader does not die from an external SIGHUP
|
||||
OR SIGTERM -- confirmed independent of this class with a bare
|
||||
pty.fork()+bash+os.kill script, so this is not something tai is
|
||||
doing wrong, it is this environment's actual bash behavior. SIGKILL
|
||||
is the one signal no process can ignore, so that's what's used here.
|
||||
This is a clean kill, not a true tmux-style detach-and-reattach --
|
||||
there is no session persistence to reattach to yet, so it is named
|
||||
and documented as "quit," not "detach", to not promise what isn't built."""
|
||||
out = bytearray()
|
||||
i = 0
|
||||
while i < len(data):
|
||||
byte = data[i:i + 1]
|
||||
if pending_prefix:
|
||||
pending_prefix = False
|
||||
if byte == b"q":
|
||||
# killpg, not kill: pty.fork()'s child is its own session
|
||||
# and process-group leader (pid == pgid), so this reaches
|
||||
# a plain foreground child the bare shell pid alone would
|
||||
# miss. It does NOT reach every descendant, and verified
|
||||
# empirically rather than assumed: an interactive bash's
|
||||
# own job control puts each backgrounded job (`sleep &`)
|
||||
# in a NEW process group of its own, specifically so
|
||||
# signals to the shell's group -- Ctrl+C, this killpg,
|
||||
# closing a real terminal -- don't also take down its
|
||||
# background jobs. A backgrounded job surviving `quit` is
|
||||
# therefore expected, standard shell behavior, identical
|
||||
# to what closing a real terminal does, not a gap to
|
||||
# close here.
|
||||
try:
|
||||
os.killpg(self.pid, signal.SIGKILL)
|
||||
except OSError:
|
||||
os.kill(self.pid, signal.SIGKILL)
|
||||
elif byte == Consts.PLEXER_PREFIX_KEY:
|
||||
out += Consts.PLEXER_PREFIX_KEY
|
||||
i += 1
|
||||
continue
|
||||
if byte == Consts.PLEXER_PREFIX_KEY:
|
||||
pending_prefix = True
|
||||
i += 1
|
||||
continue
|
||||
out += byte
|
||||
i += 1
|
||||
return bytes(out), pending_prefix
|
||||
|
||||
def _process_alive(self):
|
||||
try:
|
||||
pid, _status = os.waitpid(self.pid, os.WNOHANG)
|
||||
except OSError:
|
||||
return False
|
||||
return pid == 0
|
||||
|
||||
|
||||
class PlexerSetup:
|
||||
"""PlexerSetup."""
|
||||
|
||||
@staticmethod
|
||||
def plexer_dir(home):
|
||||
path = os.path.join(home, Consts.PLEXER_DIR_NAME)
|
||||
os.makedirs(path, exist_ok=True)
|
||||
try:
|
||||
os.chmod(path, 0o700)
|
||||
except OSError:
|
||||
pass
|
||||
return path
|
||||
|
||||
@staticmethod
|
||||
def plexer_should_wrap(env):
|
||||
"""The Python-side half of the detection flow (see research/shell-multiplexer-auto-attach/).
|
||||
The cheap, fast checks ($- interactive flag, isatty) already happened in the bashrc
|
||||
block itself before tai.py was even invoked -- this only covers what needs Python:
|
||||
never wrap if a plexer is already active, and never wrap if tmux already is."""
|
||||
if env.get(Consts.PLEXER_ACTIVE_ENV):
|
||||
return False
|
||||
if env.get("TMUX"):
|
||||
return False
|
||||
if os.path.exists(os.path.join(os.path.expanduser("~"), ".tai-no-plexer")):
|
||||
return False
|
||||
return True
|
||||
@@ -0,0 +1,133 @@
|
||||
"""Version bumps, self-backups, and release bookkeeping."""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import sqlite3
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["Release"]
|
||||
|
||||
|
||||
class Release:
|
||||
"""Release."""
|
||||
|
||||
@staticmethod
|
||||
def next_version(current, part):
|
||||
major, minor, patch = (int(item) for item in current.split("."))
|
||||
if part == "major":
|
||||
return "%d.0.0" % (major + 1)
|
||||
if part == "minor":
|
||||
return "%d.%d.0" % (major, minor + 1)
|
||||
if part == "patch":
|
||||
return "%d.%d.%d" % (major, minor, patch + 1)
|
||||
raise ValueError("part must be major, minor, or patch")
|
||||
|
||||
@staticmethod
|
||||
def atomic_write(path, content):
|
||||
tmp = "%s.tmp-%d" % (path, os.getpid())
|
||||
with open(tmp, "w", encoding="utf-8") as handle:
|
||||
handle.write(content)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
os.replace(tmp, path)
|
||||
|
||||
@staticmethod
|
||||
def sealed_write(path, content):
|
||||
Release.atomic_write(path, content)
|
||||
try:
|
||||
os.chmod(path, 0o600)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def quarantine(path):
|
||||
try:
|
||||
if os.path.exists(path):
|
||||
os.replace(path, "%s.corrupt-%d" % (path, int(time.time())))
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def release_bump(script_file, part):
|
||||
with open(script_file, "r", encoding="utf-8") as handle:
|
||||
text = handle.read()
|
||||
match = Consts.VERSION_RE.search(text)
|
||||
if match is None:
|
||||
raise ValueError("no VERSION line in %s" % script_file)
|
||||
old = ".".join(match.groups())
|
||||
new = Release.next_version(old, part)
|
||||
matched = text[match.start():match.end()]
|
||||
indent = matched[:len(matched) - len(matched.lstrip())]
|
||||
Release.atomic_write(script_file, text[:match.start()] + indent + "VERSION = \"%s\"" % new + text[match.end():])
|
||||
return old, new
|
||||
|
||||
@staticmethod
|
||||
def backups_dir(home):
|
||||
return os.path.join(home, "backups")
|
||||
|
||||
@staticmethod
|
||||
def prune_backups(folder):
|
||||
try:
|
||||
names = sorted((entry for entry in os.listdir(folder) if entry.startswith("tai-") and entry.endswith(".py")), key=lambda entry: os.path.getmtime(os.path.join(folder, entry)))
|
||||
except OSError:
|
||||
return 0
|
||||
pruned = 0
|
||||
while len(names) > Consts.BACKUP_KEEP:
|
||||
try:
|
||||
os.remove(os.path.join(folder, names.pop(0)))
|
||||
pruned += 1
|
||||
except OSError:
|
||||
break
|
||||
return pruned
|
||||
|
||||
@staticmethod
|
||||
def snapshot_self(script_file, folder, version):
|
||||
with open(script_file, "rb") as handle:
|
||||
data = handle.read()
|
||||
digest = hashlib.sha256(data).hexdigest()[:8]
|
||||
stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
|
||||
name = "tai-%s-%s-%s.py" % (version, stamp, digest)
|
||||
os.makedirs(folder, exist_ok=True)
|
||||
dest = os.path.join(folder, name)
|
||||
with open(dest, "wb") as handle:
|
||||
handle.write(data)
|
||||
Release.prune_backups(folder)
|
||||
return dest, digest
|
||||
|
||||
@staticmethod
|
||||
def backup_has_digest(folder, digest):
|
||||
try:
|
||||
return any(entry.endswith("-%s.py" % digest) for entry in os.listdir(folder))
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def ensure_self_backup(config, store, script_file=None, actor="main"):
|
||||
script_file = script_file or os.path.abspath(Consts.TAI_SCRIPT)
|
||||
folder = Release.backups_dir(config.home)
|
||||
with open(script_file, "rb") as handle:
|
||||
digest = hashlib.sha256(handle.read()).hexdigest()[:8]
|
||||
if Release.backup_has_digest(folder, digest):
|
||||
return None
|
||||
dest, _digest = Release.snapshot_self(script_file, folder, Consts.VERSION)
|
||||
try:
|
||||
store.audit_event(actor, "snapshot", script_file, "backed up as %s" % os.path.basename(dest), tags=["snapshot", "v" + Consts.VERSION])
|
||||
text, override = Text.read_capped(script_file, Consts.FILE_RECORD_MAX)
|
||||
if text is not None:
|
||||
store.upsert_file_record(script_file, text, override)
|
||||
except (sqlite3.Error, OSError):
|
||||
pass
|
||||
return dest
|
||||
|
||||
@staticmethod
|
||||
def do_release(store, script_file, folder, part, message, actor="main"):
|
||||
if not message or not message.strip():
|
||||
raise ValueError("release message is required")
|
||||
old, new = Release.release_bump(script_file, part)
|
||||
dest, _digest = Release.snapshot_self(script_file, folder, new)
|
||||
row_id = store.audit_event(actor, "release", script_file, message.strip()[:500], old=old, new=new, tags=["release", "v" + new])
|
||||
return old, new, dest, row_id
|
||||
@@ -0,0 +1,578 @@
|
||||
"""Terminal output: color, markdown, diffs, spinners, live views."""
|
||||
|
||||
import collections
|
||||
import difflib
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["Ansi", "Spinner", "LiveStream", "Render"]
|
||||
|
||||
|
||||
class Ansi:
|
||||
RESET = "\033[0m"
|
||||
BOLD = "\033[1m"
|
||||
DIM = "\033[2m"
|
||||
UNDERLINE = "\033[4m"
|
||||
STRIKE = "\033[9m"
|
||||
RED = "\033[31m"
|
||||
GREEN = "\033[32m"
|
||||
YELLOW = "\033[33m"
|
||||
BLUE = "\033[34m"
|
||||
MAGENTA = "\033[35m"
|
||||
CYAN = "\033[36m"
|
||||
GRAY = "\033[90m"
|
||||
|
||||
|
||||
class Spinner:
|
||||
FRAMES = ("⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏")
|
||||
|
||||
def __init__(self, label):
|
||||
self.label = label
|
||||
self.stop_flag = threading.Event()
|
||||
self.worker = None
|
||||
|
||||
def start(self):
|
||||
if not sys.stderr.isatty():
|
||||
return
|
||||
self.worker = threading.Thread(target=self.spin, daemon=True)
|
||||
self.worker.start()
|
||||
|
||||
def spin(self):
|
||||
pos = 0
|
||||
while not self.stop_flag.is_set():
|
||||
frame = self.FRAMES[pos % len(self.FRAMES)]
|
||||
sys.stderr.write("\r%s %s" % (frame, self.label))
|
||||
sys.stderr.flush()
|
||||
pos += 1
|
||||
time.sleep(0.08)
|
||||
sys.stderr.write("\r%s\r" % (" " * (len(self.label) + 2)))
|
||||
sys.stderr.flush()
|
||||
|
||||
def stop(self):
|
||||
self.stop_flag.set()
|
||||
if self.worker is not None:
|
||||
self.worker.join(timeout=1)
|
||||
|
||||
|
||||
class LiveStream:
|
||||
def __init__(self, height=Consts.STREAM_HEIGHT, width=0, file=None, enabled=True):
|
||||
self.height = max(1, height)
|
||||
self.width = width or shutil.get_terminal_size().columns
|
||||
self.file = file if file is not None else sys.stdout
|
||||
self.enabled = enabled
|
||||
self.lines = collections.deque(maxlen=self.height)
|
||||
self.shown = 0
|
||||
|
||||
def feed(self, line):
|
||||
if not self.enabled:
|
||||
return
|
||||
cleaned = Render.truncate_ansi(line.split("\r")[-1].rstrip("\n").expandtabs(8), max(8, self.width - 2))
|
||||
self.lines.append(cleaned)
|
||||
if self.shown < self.height:
|
||||
self.file.write("│ %s\n" % cleaned)
|
||||
self.shown += 1
|
||||
else:
|
||||
self.file.write("\033[%dA" % self.height)
|
||||
for kept in self.lines:
|
||||
self.file.write("\r\033[K│ %s\n" % kept)
|
||||
self.file.flush()
|
||||
|
||||
def close(self, status):
|
||||
if not self.enabled:
|
||||
return
|
||||
if self.shown == 0:
|
||||
self.file.write("│ (no output)\n")
|
||||
self.file.write("│ %s\n" % Render.truncate_ansi(status, max(8, self.width - 2)))
|
||||
self.file.flush()
|
||||
|
||||
|
||||
class Render:
|
||||
"""Render."""
|
||||
|
||||
@staticmethod
|
||||
def paint(text, *codes):
|
||||
if os.environ.get("NO_COLOR") or not sys.stdout.isatty():
|
||||
return text
|
||||
return "".join(codes) + text + Ansi.RESET
|
||||
|
||||
@staticmethod
|
||||
def stream_status(code, output, elapsed, timed_out):
|
||||
detail = "%d lines · %s" % (len(output.splitlines()), Text.human_size(len(output.encode("utf-8", "replace"))))
|
||||
if timed_out:
|
||||
return "timed out after %ds (killed) · %s" % (Consts.SHELL_TIMEOUT, detail)
|
||||
return "exit %d · %.1fs · %s" % (code, elapsed, detail)
|
||||
|
||||
@staticmethod
|
||||
def truncate_ansi(text, width):
|
||||
plain = Consts.ANSI_RE.sub("", text)
|
||||
if len(plain) <= width:
|
||||
return text
|
||||
out = []
|
||||
count = 0
|
||||
pos = 0
|
||||
for match in Consts.ANSI_RE.finditer(text):
|
||||
chunk = text[pos:match.start()]
|
||||
take = min(len(chunk), width - count)
|
||||
out.append(chunk[:take])
|
||||
count += take
|
||||
if count >= width:
|
||||
pos = match.start()
|
||||
break
|
||||
out.append(match.group(0))
|
||||
pos = match.end()
|
||||
if count < width:
|
||||
out.append(Consts.ANSI_RE.sub("", text[pos:])[:width - count])
|
||||
return "".join(out) + Ansi.RESET
|
||||
|
||||
@staticmethod
|
||||
def run_live(argv, timeout=Consts.SHELL_TIMEOUT, shell=False, cwd=None, env=None, stream=None):
|
||||
started = time.time()
|
||||
proc = subprocess.Popen(argv, shell=shell, cwd=cwd, env=env, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True, errors="replace", bufsize=1)
|
||||
chunks = []
|
||||
pending = queue.Queue()
|
||||
|
||||
def pump():
|
||||
try:
|
||||
for line in proc.stdout:
|
||||
chunks.append(line)
|
||||
pending.put(line)
|
||||
finally:
|
||||
pending.put(None)
|
||||
|
||||
reader = threading.Thread(target=pump, daemon=True)
|
||||
reader.start()
|
||||
timed_out = False
|
||||
try:
|
||||
deadline = started + timeout
|
||||
while True:
|
||||
remaining = deadline - time.time()
|
||||
if remaining <= 0:
|
||||
timed_out = True
|
||||
break
|
||||
try:
|
||||
item = pending.get(timeout=min(0.05, remaining))
|
||||
except queue.Empty:
|
||||
continue
|
||||
if item is None:
|
||||
break
|
||||
if stream is not None:
|
||||
stream.feed(item)
|
||||
except BaseException:
|
||||
proc.kill()
|
||||
proc.wait()
|
||||
raise
|
||||
finally:
|
||||
if timed_out:
|
||||
proc.kill()
|
||||
try:
|
||||
proc.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
proc.wait()
|
||||
reader.join(timeout=10)
|
||||
if not timed_out:
|
||||
while not pending.empty():
|
||||
item = pending.get_nowait()
|
||||
if item is None:
|
||||
break
|
||||
if stream is not None:
|
||||
stream.feed(item)
|
||||
return proc.returncode, "".join(chunks), time.time() - started, timed_out
|
||||
|
||||
@staticmethod
|
||||
def color_enabled():
|
||||
return not os.environ.get("NO_COLOR") and sys.stdout.isatty()
|
||||
|
||||
@staticmethod
|
||||
def highlight_python(line):
|
||||
parts = []
|
||||
pos = 0
|
||||
for match in Consts.PY_TOKEN_RE.finditer(line):
|
||||
parts.append(line[pos:match.start()])
|
||||
kind = match.lastgroup
|
||||
text = match.group(0)
|
||||
if kind == "string":
|
||||
parts.append(Ansi.YELLOW + text + Ansi.RESET)
|
||||
elif kind == "comment":
|
||||
parts.append(Ansi.GRAY + text + Ansi.RESET)
|
||||
elif kind == "number":
|
||||
parts.append(Consts.FG_ORANGE + text + Ansi.RESET)
|
||||
elif kind == "keyword":
|
||||
parts.append(Ansi.MAGENTA + text + Ansi.RESET)
|
||||
else:
|
||||
parts.append(Ansi.CYAN + text + Ansi.RESET)
|
||||
pos = match.end()
|
||||
parts.append(line[pos:])
|
||||
return "".join(parts)
|
||||
|
||||
@staticmethod
|
||||
def split_diff_lines(text):
|
||||
if not text:
|
||||
return []
|
||||
lines = text.split("\n")
|
||||
if lines and lines[-1] == "":
|
||||
lines.pop()
|
||||
return lines
|
||||
|
||||
@staticmethod
|
||||
def render_diff(path, old, new, width=0, colors=True):
|
||||
if new is not None and (old or "") == new:
|
||||
return ""
|
||||
old_lines = Render.split_diff_lines(old or "")
|
||||
new_lines = Render.split_diff_lines(new) if new is not None else []
|
||||
if not width:
|
||||
try:
|
||||
width = shutil.get_terminal_size().columns
|
||||
except OSError:
|
||||
width = 80
|
||||
matcher = difflib.SequenceMatcher(None, old_lines, new_lines, autojunk=False)
|
||||
rows = []
|
||||
for group in matcher.get_grouped_opcodes(Consts.DIFF_CONTEXT):
|
||||
if rows:
|
||||
rows.append(("gap", 0, 0, ""))
|
||||
for tag, i1, i2, j1, j2 in group:
|
||||
if tag == "equal":
|
||||
for offset in range(i2 - i1):
|
||||
rows.append(("ctx", i1 + offset + 1, j1 + offset + 1, old_lines[i1 + offset]))
|
||||
elif tag == "delete":
|
||||
for offset in range(i2 - i1):
|
||||
rows.append(("del", i1 + offset + 1, 0, old_lines[i1 + offset]))
|
||||
elif tag == "insert":
|
||||
for offset in range(j2 - j1):
|
||||
rows.append(("add", 0, j1 + offset + 1, new_lines[j1 + offset]))
|
||||
else:
|
||||
for offset in range(i2 - i1):
|
||||
rows.append(("del", i1 + offset + 1, 0, old_lines[i1 + offset]))
|
||||
for offset in range(j2 - j1):
|
||||
rows.append(("add", 0, j1 + offset + 1, new_lines[j1 + offset]))
|
||||
added = sum(1 for kind, _a, _b, _t in rows if kind == "add")
|
||||
removed = sum(1 for kind, _a, _b, _t in rows if kind == "del")
|
||||
if new is None:
|
||||
head = "── %s (deleted, %d lines)" % (path, len(old_lines))
|
||||
elif not old_lines:
|
||||
head = "── %s (new file, %d lines)" % (path, len(new_lines))
|
||||
else:
|
||||
head = "── %s (+%d -%d)" % (path, added, removed)
|
||||
python = str(path).lower().endswith(".py")
|
||||
num_width = max(4, len(str(max(len(old_lines), len(new_lines), 1))))
|
||||
code_width = max(20, width - num_width - 4)
|
||||
fence = None
|
||||
out = []
|
||||
out.append(Ansi.DIM + "│ " + head + Ansi.RESET if colors else "│ " + head)
|
||||
for kind, old_no, new_no, text in rows[:Consts.DIFF_MAX_LINES]:
|
||||
if kind == "gap":
|
||||
out.append(Ansi.DIM + "···" + Ansi.RESET if colors else "···")
|
||||
continue
|
||||
code = Render.strip_ansi(text).expandtabs(8)
|
||||
if fence is not None:
|
||||
if python and colors:
|
||||
code = Ansi.YELLOW + code + Ansi.RESET
|
||||
if fence in text:
|
||||
fence = None
|
||||
elif python and colors:
|
||||
for mark in ('"""', "'''"):
|
||||
if text.count(mark) % 2 == 1:
|
||||
fence = mark
|
||||
code = Ansi.YELLOW + code + Ansi.RESET
|
||||
break
|
||||
else:
|
||||
code = Render.highlight_python(code)
|
||||
number = old_no if kind in ("ctx", "del") else new_no
|
||||
gutter = str(number).rjust(num_width)
|
||||
if not colors:
|
||||
sign = {"ctx": " ", "del": "-", "add": "+"}[kind]
|
||||
out.append("%s %s %s" % (gutter, sign, code[:code_width]))
|
||||
continue
|
||||
dimmed = Ansi.DIM + gutter + Ansi.RESET
|
||||
if kind == "ctx":
|
||||
out.append("%s %s" % (dimmed, Render.truncate_ansi(code, code_width)))
|
||||
elif kind == "del":
|
||||
row = "%s %s %s" % (dimmed, Ansi.RED + Ansi.BOLD + "-" + Ansi.RESET, Render.truncate_ansi(code, code_width))
|
||||
out.append(Render.pad_tinted(Render.style_span(row, Consts.BG_DEL), width))
|
||||
else:
|
||||
row = "%s %s %s" % (dimmed, Ansi.GREEN + Ansi.BOLD + "+" + Ansi.RESET, Render.truncate_ansi(code, code_width))
|
||||
out.append(Render.pad_tinted(Render.style_span(row, Consts.BG_ADD), width))
|
||||
if len(rows) > Consts.DIFF_MAX_LINES:
|
||||
note = "··· %d more lines hidden" % (len(rows) - Consts.DIFF_MAX_LINES)
|
||||
out.append(Ansi.DIM + note + Ansi.RESET if colors else note)
|
||||
return "\n".join(out)
|
||||
|
||||
@staticmethod
|
||||
def pad_tinted(line, width):
|
||||
plain = len(Consts.ANSI_RE.sub("", line))
|
||||
if plain < width:
|
||||
fill = " " * (width - plain)
|
||||
if line.endswith(Ansi.RESET):
|
||||
return line[:-len(Ansi.RESET)] + fill + Ansi.RESET
|
||||
return line + fill
|
||||
return line
|
||||
|
||||
@staticmethod
|
||||
def strip_ansi(text):
|
||||
return Consts.ANSI_RE.sub("", text)
|
||||
|
||||
@staticmethod
|
||||
def style_text(text, *codes):
|
||||
return "".join(codes) + text + Ansi.RESET
|
||||
|
||||
@staticmethod
|
||||
def style_span(text, *codes):
|
||||
prefix = "".join(codes)
|
||||
return prefix + text.replace(Ansi.RESET, Ansi.RESET + prefix) + Ansi.RESET
|
||||
|
||||
@staticmethod
|
||||
def wrap_ansi(text, width, first_prefix="", next_prefix=""):
|
||||
parts = re.findall(r"\033\[[0-9;]*m|[^\s\033]+|\s+", text)
|
||||
lines = []
|
||||
current = first_prefix
|
||||
length = len(Render.strip_ansi(first_prefix))
|
||||
base = length
|
||||
active = ""
|
||||
for part in parts:
|
||||
if not part:
|
||||
continue
|
||||
if part.startswith("\033"):
|
||||
current += part
|
||||
if part == Ansi.RESET:
|
||||
active = ""
|
||||
else:
|
||||
active += part
|
||||
continue
|
||||
if part.isspace():
|
||||
if length < width and length > base:
|
||||
current += " "
|
||||
length += 1
|
||||
continue
|
||||
if length + len(part) > width and length > base:
|
||||
if active:
|
||||
current += Ansi.RESET
|
||||
lines.append(current.rstrip())
|
||||
current = next_prefix + active
|
||||
length = len(Render.strip_ansi(next_prefix))
|
||||
base = length
|
||||
current += part
|
||||
length += len(part)
|
||||
if active:
|
||||
current += Ansi.RESET
|
||||
if current.strip():
|
||||
lines.append(current.rstrip())
|
||||
return lines
|
||||
|
||||
@staticmethod
|
||||
def md_inline(text):
|
||||
stashed = []
|
||||
|
||||
def stash(value):
|
||||
stashed.append(value)
|
||||
return "\x00%d\x00" % (len(stashed) - 1)
|
||||
|
||||
text = re.sub(r"`([^`\n]+)`", lambda match: stash(Render.style_text(match.group(1), Ansi.CYAN)), text)
|
||||
text = re.sub(r"\\([!\"#$%&'()*+,\-./:;<=>?@\[\\\]^_`{|}~])", lambda match: stash(match.group(1)), text)
|
||||
text = re.sub(r"!\[([^\]\n]*)\]\(([^)\s\n]+)(?:\s+\"[^\"]*\")?\)", lambda match: stash("[image: %s]" % (match.group(1) or match.group(2))), text)
|
||||
text = re.sub(r"\[([^\]\n]+)\]\(([^)\s\n]+)(?:\s+\"[^\"]*\")?\)", lambda match: Render.style_span(Render.md_inline(match.group(1)), Ansi.BLUE, Ansi.UNDERLINE) + Render.style_text(" (%s)" % match.group(2), Ansi.DIM), text)
|
||||
text = re.sub(r"\*\*([^*\n]+)\*\*", lambda match: Render.style_text(match.group(1), Ansi.BOLD), text)
|
||||
text = re.sub(r"(?<!\w)__([^_\n]+)__(?!\w)", lambda match: Render.style_text(match.group(1), Ansi.BOLD), text)
|
||||
text = re.sub(r"~~([^~\n]+)~~", lambda match: Render.style_text(match.group(1), Ansi.STRIKE), text)
|
||||
text = re.sub(r"\*([^*\n]+)\*", lambda match: Render.style_text(match.group(1), Ansi.DIM), text)
|
||||
text = re.sub(r"(?<!\w)_([^_\n]+)_(?!\w)", lambda match: Render.style_text(match.group(1), Ansi.DIM), text)
|
||||
for pos, value in enumerate(stashed):
|
||||
text = text.replace("\x00%d\x00" % pos, value)
|
||||
return text
|
||||
|
||||
@staticmethod
|
||||
def md_split_row(line):
|
||||
line = line.strip()
|
||||
if line.startswith("|"):
|
||||
line = line[1:]
|
||||
if line.endswith("|") and not line.endswith("\\|"):
|
||||
line = line[:-1]
|
||||
return [cell.replace("\\|", "|").strip() for cell in re.split(r"(?<!\\)\|", line)]
|
||||
|
||||
@staticmethod
|
||||
def md_is_separator(line):
|
||||
cells = Render.md_split_row(line)
|
||||
return bool(cells) and all(re.fullmatch(r":?-{1,}:?", cell) for cell in cells)
|
||||
|
||||
@staticmethod
|
||||
def md_align_cell(cell, width, align):
|
||||
pad = max(0, width - len(Render.strip_ansi(cell)))
|
||||
if align == "right":
|
||||
return " " * pad + cell
|
||||
if align == "center":
|
||||
left = pad // 2
|
||||
return " " * left + cell + " " * (pad - left)
|
||||
return cell + " " * pad
|
||||
|
||||
@staticmethod
|
||||
def md_table_block(header, aligns, rows):
|
||||
span = max([len(header)] + [len(row) for row in rows])
|
||||
header = (header + [""] * span)[:span]
|
||||
rows = [(row + [""] * span)[:span] for row in rows]
|
||||
head = [Render.md_inline(cell) for cell in header]
|
||||
body = [[Render.md_inline(cell) for cell in row] for row in rows]
|
||||
widths = [0] * span
|
||||
for pos in range(span):
|
||||
widths[pos] = max([len(Render.strip_ansi(head[pos]))] + [len(Render.strip_ansi(row[pos])) for row in body])
|
||||
bar = Render.style_text("│", Ansi.DIM)
|
||||
lines = [Render.style_text("┌" + "┬".join("─" * (width + 2) for width in widths) + "┐", Ansi.DIM)]
|
||||
lines.append(bar + bar.join(" " + Render.style_span(Render.md_align_cell(cell, widths[pos], aligns[pos]), Ansi.BOLD) + " " for pos, cell in enumerate(head)) + bar)
|
||||
lines.append(Render.style_text("├" + "┼".join("─" * (width + 2) for width in widths) + "┤", Ansi.DIM))
|
||||
for row in body:
|
||||
lines.append(bar + bar.join(" " + Render.md_align_cell(cell, widths[pos], aligns[pos]) + " " for pos, cell in enumerate(row)) + bar)
|
||||
lines.append(Render.style_text("└" + "┴".join("─" * (width + 2) for width in widths) + "┘", Ansi.DIM))
|
||||
return lines
|
||||
|
||||
@staticmethod
|
||||
def render_markdown(text, width=None, color=None):
|
||||
if color is None:
|
||||
color = not os.environ.get("NO_COLOR") and sys.stdout.isatty()
|
||||
if not color:
|
||||
return text
|
||||
if width is None:
|
||||
try:
|
||||
width = shutil.get_terminal_size().columns
|
||||
except OSError:
|
||||
width = 80
|
||||
text = (text or "").replace("\r\n", "\n").expandtabs(8)
|
||||
text = re.sub(r"<br\s*/?>", " \n", text, flags=re.IGNORECASE)
|
||||
raw = text.split("\n")
|
||||
out = []
|
||||
pos = 0
|
||||
total = len(raw)
|
||||
while pos < total:
|
||||
line = raw[pos]
|
||||
stripped = line.strip()
|
||||
if not stripped:
|
||||
pos += 1
|
||||
continue
|
||||
fence = re.match(r"^\s*(`{3,}|~{3,})\s*\S*\s*$", line)
|
||||
if fence:
|
||||
char = fence.group(1)[0]
|
||||
pos += 1
|
||||
while pos < total:
|
||||
closer = raw[pos].strip()
|
||||
if closer and set(closer) == {char} and len(closer) >= 3:
|
||||
pos += 1
|
||||
break
|
||||
out.append(" " + Render.style_text(raw[pos], Ansi.DIM) if raw[pos].strip() else "")
|
||||
pos += 1
|
||||
out.append("")
|
||||
continue
|
||||
if "|" in line and pos + 1 < total and Render.md_is_separator(raw[pos + 1]):
|
||||
header = Render.md_split_row(line)
|
||||
aligns = []
|
||||
for cell in Render.md_split_row(raw[pos + 1]):
|
||||
if cell.startswith(":") and cell.endswith(":") and len(cell) > 2:
|
||||
aligns.append("center")
|
||||
elif cell.endswith(":"):
|
||||
aligns.append("right")
|
||||
else:
|
||||
aligns.append("left")
|
||||
pos += 2
|
||||
rows = []
|
||||
while pos < total and "|" in raw[pos] and raw[pos].strip():
|
||||
rows.append(Render.md_split_row(raw[pos]))
|
||||
pos += 1
|
||||
while len(aligns) < max([len(header)] + [len(row) for row in rows]):
|
||||
aligns.append("left")
|
||||
out.extend(Render.md_table_block(header, aligns, rows))
|
||||
out.append("")
|
||||
continue
|
||||
heading = re.match(r"^(#{1,6})\s+(.*\S)\s*$", stripped)
|
||||
if heading:
|
||||
level = len(heading.group(1))
|
||||
body = Render.md_inline(heading.group(2))
|
||||
if level == 1:
|
||||
out.append(Render.style_span(body, Ansi.BOLD, Ansi.CYAN))
|
||||
elif level == 2:
|
||||
out.append(Render.style_span(body, Ansi.BOLD))
|
||||
else:
|
||||
out.append(Render.style_span(body, Ansi.BOLD, Ansi.DIM))
|
||||
out.append("")
|
||||
pos += 1
|
||||
continue
|
||||
if re.fullmatch(r"\*{3,}|-{3,}|_{3,}", stripped.replace(" ", "")):
|
||||
out.append(Render.style_text("─" * width, Ansi.DIM))
|
||||
out.append("")
|
||||
pos += 1
|
||||
continue
|
||||
if stripped.startswith(">"):
|
||||
paras = []
|
||||
current = []
|
||||
while pos < total and raw[pos].strip().startswith(">"):
|
||||
inner = re.sub(r"^\s*(?:>\s?)+", "", raw[pos])
|
||||
if inner.strip():
|
||||
current.append(inner.strip())
|
||||
elif current:
|
||||
paras.append(current)
|
||||
current = []
|
||||
pos += 1
|
||||
if current:
|
||||
paras.append(current)
|
||||
for index, para in enumerate(paras):
|
||||
if index:
|
||||
out.append(Render.style_text("│", Ansi.DIM))
|
||||
prefix = Render.style_text("│ ", Ansi.DIM)
|
||||
for wrapped in Render.wrap_ansi(Render.md_inline(" ".join(para)), width, prefix, prefix):
|
||||
out.append(wrapped)
|
||||
out.append("")
|
||||
continue
|
||||
bullet = re.match(r"^(\s*)[-*+]\s+(.*\S)\s*$", line)
|
||||
numbered = re.match(r"^(\s*)\d+[.)]\s+(.*\S)\s*$", line)
|
||||
if bullet or numbered:
|
||||
counters = {}
|
||||
while pos < total:
|
||||
again_bullet = re.match(r"^(\s*)[-*+]\s+(.*\S)\s*$", raw[pos])
|
||||
again_numbered = re.match(r"^(\s*)\d+[.)]\s+(.*\S)\s*$", raw[pos])
|
||||
if not again_bullet and not again_numbered:
|
||||
break
|
||||
if again_bullet:
|
||||
level = len(again_bullet.group(1)) // 2
|
||||
marker = Consts.MD_BULLETS[level % len(Consts.MD_BULLETS)]
|
||||
body = again_bullet.group(2)
|
||||
check = re.match(r"^\[([ xX])\]\s+(.*\S)\s*$", body)
|
||||
if check:
|
||||
marker = "☑" if check.group(1).lower() == "x" else "☐"
|
||||
body = check.group(2)
|
||||
else:
|
||||
level = len(again_numbered.group(1)) // 2
|
||||
counters[level] = counters.get(level, 0) + 1
|
||||
for deeper in [key for key in counters if key > level]:
|
||||
del counters[deeper]
|
||||
marker = "%d." % counters[level]
|
||||
body = again_numbered.group(2)
|
||||
pad = " " * level
|
||||
for wrapped in Render.wrap_ansi(Render.md_inline(body), width, pad + marker + " ", pad + " " * (len(marker) + 1)):
|
||||
out.append(wrapped)
|
||||
pos += 1
|
||||
out.append("")
|
||||
continue
|
||||
para = []
|
||||
while pos < total and raw[pos].strip():
|
||||
para.append(raw[pos])
|
||||
pos += 1
|
||||
segments = []
|
||||
current = []
|
||||
for chunk in para:
|
||||
current.append(chunk.strip())
|
||||
if chunk.endswith(" "):
|
||||
segments.append(" ".join(current))
|
||||
current = []
|
||||
if current:
|
||||
segments.append(" ".join(current))
|
||||
for segment in segments:
|
||||
for wrapped in Render.wrap_ansi(Render.md_inline(segment), width):
|
||||
out.append(wrapped)
|
||||
out.append("")
|
||||
while out and not out[-1]:
|
||||
out.pop()
|
||||
return "\n".join(out)
|
||||
@@ -0,0 +1,149 @@
|
||||
"""Container sandbox box and in-box command execution."""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.errors import BackendError
|
||||
from tai_pkg.opencode import Opencode
|
||||
from tai_pkg.render import Ansi
|
||||
from tai_pkg.render import Render
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["Sandbox", "_BOX_PYTHON_CACHE", "_BOX_OPENCODE_CACHE"]
|
||||
|
||||
|
||||
class Sandbox:
|
||||
"""Sandbox."""
|
||||
|
||||
@staticmethod
|
||||
def container_engine():
|
||||
if shutil.which("podman"):
|
||||
return "podman"
|
||||
if shutil.which("docker"):
|
||||
return "docker"
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def box_state(engine):
|
||||
try:
|
||||
done = subprocess.run([engine, "inspect", "-f", "{{.State.Running}}", Consts.BOX_NAME], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=15)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
return "missing"
|
||||
if done.returncode != 0:
|
||||
return "missing"
|
||||
return "running" if done.stdout.strip() == "true" else "stopped"
|
||||
|
||||
@staticmethod
|
||||
def ensure_box():
|
||||
engine = Sandbox.container_engine()
|
||||
if not engine:
|
||||
raise BackendError("no container engine found, install podman")
|
||||
state = Sandbox.box_state(engine)
|
||||
if state == "running":
|
||||
return engine
|
||||
if state == "stopped":
|
||||
done = subprocess.run([engine, "start", Consts.BOX_NAME], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=60)
|
||||
if done.returncode != 0:
|
||||
raise BackendError("cannot start box: " + (done.stderr or "").strip()[:200])
|
||||
return engine
|
||||
print("building sandbox image %s (one-time, takes minutes)..." % Consts.BOX_IMAGE)
|
||||
with tempfile.TemporaryDirectory() as build:
|
||||
for filename, content in (("Containerfile", Consts.BOX_CONTAINERFILE), ("stt.py", Consts.BOX_STT), ("tts.py", Consts.BOX_TTS)):
|
||||
with open(os.path.join(build, filename), "w", encoding="utf-8") as handle:
|
||||
handle.write(content)
|
||||
done = subprocess.run([engine, "build", "-t", Consts.BOX_IMAGE, build], timeout=1800)
|
||||
if done.returncode != 0:
|
||||
raise BackendError("box image build failed")
|
||||
done = subprocess.run([engine, "run", "-d", "--name", Consts.BOX_NAME, "--restart", "unless-stopped", Consts.BOX_IMAGE], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=120)
|
||||
if done.returncode != 0:
|
||||
raise BackendError("cannot start box: " + (done.stderr or "").strip()[:200])
|
||||
return engine
|
||||
|
||||
@staticmethod
|
||||
def box_exec(engine, argv, extra=(), input_data=None, timeout=180):
|
||||
return subprocess.run([engine, "exec", *extra, "-i", Consts.BOX_NAME, *argv], input=input_data, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=timeout)
|
||||
|
||||
@staticmethod
|
||||
def box_python(engine):
|
||||
if engine not in _BOX_PYTHON_CACHE:
|
||||
try:
|
||||
probe = Sandbox.box_exec(engine, ["test", "-x", Consts.BOX_PYTHON], timeout=30)
|
||||
_BOX_PYTHON_CACHE[engine] = Consts.BOX_PYTHON if probe.returncode == 0 else "python3"
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
_BOX_PYTHON_CACHE[engine] = "python3"
|
||||
return _BOX_PYTHON_CACHE[engine]
|
||||
|
||||
@staticmethod
|
||||
def box_transcribe(engine, audio, timeout=300):
|
||||
done = Sandbox.box_exec(engine, [Sandbox.box_python(engine), "/box/stt.py"], input_data=audio, timeout=timeout)
|
||||
if done.returncode != 0:
|
||||
raise BackendError("box stt failed: " + done.stderr.decode("utf-8", "replace")[:300])
|
||||
return done.stdout.decode("utf-8", "replace").strip()
|
||||
|
||||
@staticmethod
|
||||
def box_speak(engine, text, voice=""):
|
||||
argv = [Sandbox.box_python(engine), "/box/tts.py"] + ([voice] if voice else [])
|
||||
done = Sandbox.box_exec(engine, argv, input_data=text.encode("utf-8"), timeout=180)
|
||||
if done.returncode != 0 or not done.stdout:
|
||||
raise BackendError("box tts failed: " + done.stderr.decode("utf-8", "replace")[:300])
|
||||
return done.stdout
|
||||
|
||||
@staticmethod
|
||||
def box_opencode_binary(engine):
|
||||
"""True once the opencode CLI is confirmed present inside the sandbox
|
||||
box, installing Node.js + `npm install -g opencode-ai` into the box
|
||||
first if it isn't there yet. Only ever called as a last resort (see
|
||||
ChatClient.all_candidates()) and only against a box that is ALREADY
|
||||
running -- this never creates the box itself; that stays the explicit,
|
||||
approval-gated action it already is (ensure_box(), via --env sandbox or
|
||||
the install tool). Cached per engine for the process lifetime so normal
|
||||
turns that succeed through a real backend never pay this cost twice."""
|
||||
if engine in _BOX_OPENCODE_CACHE:
|
||||
return _BOX_OPENCODE_CACHE[engine]
|
||||
try:
|
||||
probe = Sandbox.box_exec(engine, ["which", "opencode"], timeout=10)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
_BOX_OPENCODE_CACHE[engine] = False
|
||||
return False
|
||||
if probe.returncode == 0:
|
||||
_BOX_OPENCODE_CACHE[engine] = True
|
||||
return True
|
||||
print(Render.paint(
|
||||
"opencode: no local CLI on this machine; installing it into the existing sandbox "
|
||||
"box (container '%s') instead, as a last-resort backend. Nothing outside that "
|
||||
"container is touched. Streaming the install below:" % Consts.BOX_NAME, Ansi.DIM
|
||||
), file=sys.stderr)
|
||||
try:
|
||||
# Deliberately not box_exec(): this one streams straight to the
|
||||
# real stdout/stderr instead of capturing, so the person watching
|
||||
# sees the actual apt/curl/npm output live, not a silent pause.
|
||||
done = subprocess.run([engine, "exec", Consts.BOX_NAME, "sh", "-c", Consts.BOX_OPENCODE_INSTALL_SCRIPT], timeout=600)
|
||||
except (OSError, subprocess.SubprocessError) as exc:
|
||||
print(Render.paint("opencode box install failed: %s" % Text.short_error(exc), Ansi.DIM), file=sys.stderr)
|
||||
_BOX_OPENCODE_CACHE[engine] = False
|
||||
return False
|
||||
ok = done.returncode == 0
|
||||
print(Render.paint("opencode box install %s" % ("succeeded" if ok else "failed (exit %d)" % done.returncode), Ansi.DIM), file=sys.stderr)
|
||||
_BOX_OPENCODE_CACHE[engine] = ok
|
||||
return ok
|
||||
|
||||
@staticmethod
|
||||
def box_opencode_run(engine, model, prompt, timeout):
|
||||
argv = ["opencode", "run", "--pure", "--agent", Consts.OPENCODE_AGENT, "--format", "json", "-m", "opencode/" + model]
|
||||
try:
|
||||
done = Sandbox.box_exec(engine, argv, input_data=prompt.encode("utf-8"), timeout=timeout)
|
||||
except subprocess.TimeoutExpired:
|
||||
raise BackendError("opencode (box) run timed out after %ds" % timeout)
|
||||
except (OSError, subprocess.SubprocessError) as exc:
|
||||
raise BackendError("opencode (box) run failed: %s" % Text.short_error(exc))
|
||||
return Opencode.opencode_parse_events(done.stdout.decode("utf-8", "replace"), done.returncode, done.stderr.decode("utf-8", "replace"))
|
||||
|
||||
|
||||
_BOX_PYTHON_CACHE = {}
|
||||
|
||||
|
||||
_BOX_OPENCODE_CACHE = {}
|
||||
@@ -0,0 +1,149 @@
|
||||
"""Schedule parsing, display, and the scheduler loop."""
|
||||
|
||||
import re
|
||||
import sqlite3
|
||||
import threading
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
from tai_pkg.agents import REGISTRY
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.seal import Seal
|
||||
from tai_pkg.seal import SealError
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["Scheduler"]
|
||||
|
||||
|
||||
class Scheduler:
|
||||
"""Scheduler."""
|
||||
|
||||
@staticmethod
|
||||
def parse_schedule_at(raw):
|
||||
text = str(raw or "").strip()
|
||||
if not text:
|
||||
raise ValueError("empty datetime")
|
||||
if text[-1:] in ("Z", "z"):
|
||||
text = text[:-1] + "+00:00"
|
||||
try:
|
||||
moment = Text._fromisoformat(text)
|
||||
except ValueError:
|
||||
raise ValueError("invalid datetime, use ISO like 2026-10-08T09:00")
|
||||
return moment.astimezone(timezone.utc).isoformat()
|
||||
|
||||
@staticmethod
|
||||
def parse_duration(raw):
|
||||
match = re.fullmatch(r"(\d+)\s*([smhdw])?", str(raw or "").strip().lower())
|
||||
if not match:
|
||||
return None
|
||||
return int(match.group(1)) * {"s": 1, "m": 60, "h": 3600, "d": 86400, "w": 604800}[match.group(2) or "s"]
|
||||
|
||||
@staticmethod
|
||||
def format_delay(seconds):
|
||||
seconds = int(seconds)
|
||||
if seconds < 0:
|
||||
return "overdue"
|
||||
if seconds < 60:
|
||||
return "%ds" % seconds
|
||||
if seconds < 3600:
|
||||
return "%dm" % (seconds // 60)
|
||||
if seconds < 86400:
|
||||
return "%dh" % (seconds // 3600)
|
||||
return "%dd" % (seconds // 86400)
|
||||
|
||||
@staticmethod
|
||||
def local_display(utc_iso):
|
||||
try:
|
||||
return Text._fromisoformat(utc_iso).astimezone().strftime("%Y-%m-%d %H:%M")
|
||||
except ValueError:
|
||||
return utc_iso
|
||||
|
||||
@staticmethod
|
||||
def format_schedule_lines(items, now):
|
||||
if not items:
|
||||
return "(no schedules)"
|
||||
lines = []
|
||||
for item in items:
|
||||
try:
|
||||
due_in = (Text._fromisoformat(item["next_run"]) - now).total_seconds()
|
||||
except ValueError:
|
||||
due_in = -1
|
||||
when = "every %s" % Scheduler.format_delay(item["every"]) if item["every"] else "at %s" % Scheduler.local_display(item["next_run"])
|
||||
marked = " [%s]" % " ".join(item["tags"]) if item["tags"] else ""
|
||||
lines.append("#%d [%s] %s (%s) %s%s :: %s" % (item["id"], item["status"], when, Scheduler.format_delay(due_in), item["name"] or "(no name)", marked, item["prompt"][:80]))
|
||||
if item["last"]:
|
||||
extra = " :: %s" % item["result"][:120] if item["result"] else ""
|
||||
lines.append(" last: %s%s" % (item["last"], extra))
|
||||
return "\n".join(lines)
|
||||
|
||||
@staticmethod
|
||||
def note_schedule_result(db_path, row_id, status, result=""):
|
||||
try:
|
||||
db = sqlite3.connect(db_path, timeout=30)
|
||||
db.execute("PRAGMA journal_mode=WAL")
|
||||
db.execute("PRAGMA busy_timeout=30000")
|
||||
try:
|
||||
db.execute("UPDATE schedules SET last_status = ?, last_result = ?, updated = ? WHERE id = ?", (status, result[:2000], datetime.now(timezone.utc).isoformat(), row_id))
|
||||
db.commit()
|
||||
finally:
|
||||
db.close()
|
||||
except sqlite3.Error:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def watch_schedule(db_path, row_id, agent_id, timeout):
|
||||
status, text = REGISTRY.poll_agent(agent_id, wait=timeout + 30)
|
||||
Scheduler.note_schedule_result(db_path, row_id, status, text)
|
||||
|
||||
@staticmethod
|
||||
def scheduler_tick(config, seal, runner=None):
|
||||
fired = 0
|
||||
now = datetime.now(timezone.utc)
|
||||
db = sqlite3.connect(config.db_path, timeout=30)
|
||||
db.execute("PRAGMA journal_mode=WAL")
|
||||
db.execute("PRAGMA busy_timeout=30000")
|
||||
try:
|
||||
rows = db.execute("SELECT id, prompt, profile, every_sec, next_run, timeout FROM schedules WHERE status = 'pending' AND next_run <= ? ORDER BY next_run", (now.isoformat(),)).fetchall()
|
||||
for row_id, prompt, profile, every_sec, next_run, timeout in rows:
|
||||
stamp = datetime.now(timezone.utc).isoformat()
|
||||
if every_sec:
|
||||
try:
|
||||
claimed_next = Text._fromisoformat(next_run)
|
||||
except ValueError:
|
||||
continue
|
||||
while claimed_next <= datetime.now(timezone.utc):
|
||||
claimed_next += timedelta(seconds=every_sec)
|
||||
done = db.execute("UPDATE schedules SET next_run = ?, updated = ? WHERE id = ? AND status = 'pending' AND next_run = ?", (claimed_next.isoformat(), stamp, row_id, next_run))
|
||||
else:
|
||||
done = db.execute("UPDATE schedules SET status = 'done', updated = ? WHERE id = ? AND status = 'pending'", (stamp, row_id))
|
||||
db.commit()
|
||||
if done.rowcount != 1:
|
||||
continue
|
||||
if prompt.startswith(Seal.PREFIX):
|
||||
try:
|
||||
prompt = seal.unlock(prompt)
|
||||
except SealError as exc:
|
||||
Scheduler.note_schedule_result(config.db_path, row_id, "error: " + Text.short_error(exc))
|
||||
continue
|
||||
agent_id = REGISTRY.spawn_agent(prompt, profile, timeout, config, seal, 0, runner)
|
||||
watcher = threading.Thread(target=Scheduler.watch_schedule, args=(config.db_path, row_id, agent_id, timeout), daemon=True)
|
||||
watcher.start()
|
||||
fired += 1
|
||||
finally:
|
||||
db.close()
|
||||
return fired
|
||||
|
||||
@staticmethod
|
||||
def run_scheduler_loop(config, seal, stop_flag):
|
||||
while not stop_flag.is_set():
|
||||
try:
|
||||
Scheduler.scheduler_tick(config, seal)
|
||||
except Exception as exc:
|
||||
print("scheduler error: %s" % Text.short_error(exc))
|
||||
stop_flag.wait(Consts.SCHEDULER_INTERVAL)
|
||||
|
||||
@staticmethod
|
||||
def start_scheduler(config, seal):
|
||||
stop_flag = threading.Event()
|
||||
worker = threading.Thread(target=Scheduler.run_scheduler_loop, args=(config, seal, stop_flag), daemon=True)
|
||||
worker.start()
|
||||
return stop_flag
|
||||
+173
@@ -0,0 +1,173 @@
|
||||
"""Vault encryption seal and passphrase management."""
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import struct
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
|
||||
__all__ = ["SealError", "Seal", "SealManager"]
|
||||
|
||||
|
||||
class SealError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class Seal:
|
||||
PREFIX = "tai1$"
|
||||
ROUNDS = 200000
|
||||
|
||||
def __init__(self, home, passphrase):
|
||||
self.enabled = bool(passphrase)
|
||||
self.default_key = False
|
||||
self.master = b""
|
||||
self.enc_key = b""
|
||||
self.mac_key = b""
|
||||
if self.enabled:
|
||||
self.master = self.load_master(home, passphrase)
|
||||
self.enc_key = hmac.new(self.master, b"tai-enc-1", hashlib.sha256).digest()
|
||||
self.mac_key = hmac.new(self.master, b"tai-mac-1", hashlib.sha256).digest()
|
||||
|
||||
def load_master(self, home, passphrase):
|
||||
path = os.path.join(home, ".seal")
|
||||
raw = ""
|
||||
if os.path.exists(path):
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
raw = handle.read()
|
||||
saved = {}
|
||||
if raw:
|
||||
try:
|
||||
saved = json.loads(raw)
|
||||
salt = base64.b64decode(saved["salt"])
|
||||
except (ValueError, KeyError):
|
||||
raise SealError("seal file is corrupt")
|
||||
else:
|
||||
salt = os.urandom(16)
|
||||
master = hashlib.pbkdf2_hmac("sha256", passphrase.encode("utf-8"), salt, self.ROUNDS, 32)
|
||||
check = hmac.new(master, b"tai-check-1", hashlib.sha256).hexdigest()
|
||||
if raw:
|
||||
if not hmac.compare_digest(saved.get("check", ""), check):
|
||||
raise SealError("wrong TAI_PASSPHRASE")
|
||||
return master
|
||||
with open(path, "w", encoding="utf-8") as handle:
|
||||
handle.write(json.dumps({"salt": base64.b64encode(salt).decode("ascii"), "check": check}))
|
||||
try:
|
||||
os.chmod(path, 0o600)
|
||||
except OSError:
|
||||
pass
|
||||
return master
|
||||
|
||||
def stream(self, nonce, length):
|
||||
out = bytearray()
|
||||
counter = 0
|
||||
while len(out) < length:
|
||||
out += hmac.new(self.enc_key, nonce + struct.pack(">I", counter), hashlib.sha256).digest()
|
||||
counter += 1
|
||||
return bytes(out[:length])
|
||||
|
||||
def lock(self, text):
|
||||
if not self.enabled or text.startswith(self.PREFIX):
|
||||
return text
|
||||
nonce = os.urandom(16)
|
||||
body = text.encode("utf-8")
|
||||
keystream = self.stream(nonce, len(body))
|
||||
cipher = bytes(piece ^ keystream[pos] for pos, piece in enumerate(body))
|
||||
mac = hmac.new(self.mac_key, nonce + cipher, hashlib.sha256).digest()
|
||||
parts = [base64.b64encode(blob).decode("ascii") for blob in (nonce, cipher, mac)]
|
||||
return self.PREFIX + "$".join(parts)
|
||||
|
||||
def unlock(self, text):
|
||||
if not text.startswith(self.PREFIX):
|
||||
raise SealError("not sealed")
|
||||
try:
|
||||
_tag, nonce64, cipher64, mac64 = text.split("$", 3)
|
||||
nonce = base64.b64decode(nonce64)
|
||||
cipher = base64.b64decode(cipher64)
|
||||
mac = base64.b64decode(mac64)
|
||||
except ValueError:
|
||||
raise SealError("sealed value is corrupt")
|
||||
expect = hmac.new(self.mac_key, nonce + cipher, hashlib.sha256).digest()
|
||||
if not hmac.compare_digest(mac, expect):
|
||||
raise SealError("seal check failed")
|
||||
keystream = self.stream(nonce, len(cipher))
|
||||
plain = bytes(piece ^ keystream[pos] for pos, piece in enumerate(cipher))
|
||||
return plain.decode("utf-8")
|
||||
|
||||
def safe_unlock(self, text):
|
||||
try:
|
||||
return self.unlock(text)
|
||||
except SealError:
|
||||
return "[sealed: decrypt failed]"
|
||||
|
||||
|
||||
class SealManager:
|
||||
"""SealManager."""
|
||||
|
||||
@staticmethod
|
||||
def resolve_passphrase():
|
||||
raw = os.environ.get("TAI_PASSPHRASE")
|
||||
if raw is None:
|
||||
return Consts.DEFAULT_PASSPHRASE, True
|
||||
return raw, False
|
||||
|
||||
@staticmethod
|
||||
def rotate_seal(config, old_seal, new_passphrase):
|
||||
db = sqlite3.connect(config.db_path)
|
||||
rows = db.execute("SELECT id, text FROM events").fetchall()
|
||||
tables = [row[0] for row in db.execute("SELECT name FROM sqlite_master WHERE type = 'table'").fetchall()]
|
||||
secret_cols = [row[1] for row in db.execute("PRAGMA table_info(secrets)").fetchall()] if "secrets" in tables else []
|
||||
has_profile = "profile" in secret_cols
|
||||
secret_rows = db.execute("SELECT profile, name, value FROM secrets").fetchall() if has_profile else [(None, name, value) for name, value in db.execute("SELECT name, value FROM secrets").fetchall()] if "secrets" in tables else []
|
||||
meta_rows = db.execute("SELECT profile, name, meta FROM secrets WHERE meta IS NOT NULL").fetchall() if has_profile and "meta" in secret_cols else [(None, name, meta) for name, meta in db.execute("SELECT name, meta FROM secrets WHERE meta IS NOT NULL").fetchall()] if "meta" in secret_cols else []
|
||||
planned_rows = db.execute("SELECT id, prompt FROM schedules").fetchall() if "schedules" in tables else []
|
||||
db.close()
|
||||
texts = {}
|
||||
for row_id, text in rows:
|
||||
texts[row_id] = old_seal.unlock(text) if text.startswith(Seal.PREFIX) else text
|
||||
secret_texts = {}
|
||||
for profile, name, value in secret_rows:
|
||||
secret_texts[(profile, name)] = old_seal.unlock(value) if value.startswith(Seal.PREFIX) else value
|
||||
secret_metas = {}
|
||||
for profile, name, meta in meta_rows:
|
||||
secret_metas[(profile, name)] = old_seal.unlock(meta) if meta.startswith(Seal.PREFIX) else meta
|
||||
planned_texts = {}
|
||||
for row_id, prompt in planned_rows:
|
||||
planned_texts[row_id] = old_seal.unlock(prompt) if prompt.startswith(Seal.PREFIX) else prompt
|
||||
files = {}
|
||||
for entry in os.listdir(config.profiles_dir):
|
||||
if not entry.endswith((".sys.md", ".session.json", ".bots.json")):
|
||||
continue
|
||||
path = os.path.join(config.profiles_dir, entry)
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
raw = handle.read()
|
||||
files[path] = old_seal.unlock(raw) if raw.startswith(Seal.PREFIX) else raw
|
||||
try:
|
||||
os.remove(os.path.join(config.home, ".seal"))
|
||||
except OSError:
|
||||
pass
|
||||
fresh = Seal(config.home, new_passphrase)
|
||||
db = sqlite3.connect(config.db_path)
|
||||
for row_id, plain in texts.items():
|
||||
db.execute("UPDATE events SET text = ? WHERE id = ?", (fresh.lock(plain), row_id))
|
||||
for (profile, name), plain in secret_texts.items():
|
||||
if profile is None:
|
||||
db.execute("UPDATE secrets SET value = ? WHERE name = ?", (fresh.lock(plain), name))
|
||||
else:
|
||||
db.execute("UPDATE secrets SET value = ? WHERE profile = ? AND name = ?", (fresh.lock(plain), profile, name))
|
||||
for (profile, name), plain in secret_metas.items():
|
||||
if profile is None:
|
||||
db.execute("UPDATE secrets SET meta = ? WHERE name = ?", (fresh.lock(plain), name))
|
||||
else:
|
||||
db.execute("UPDATE secrets SET meta = ? WHERE profile = ? AND name = ?", (fresh.lock(plain), profile, name))
|
||||
for row_id, plain in planned_texts.items():
|
||||
db.execute("UPDATE schedules SET prompt = ? WHERE id = ?", (fresh.lock(plain), row_id))
|
||||
db.commit()
|
||||
db.close()
|
||||
for path, plain in files.items():
|
||||
with open(path, "w", encoding="utf-8") as handle:
|
||||
handle.write(fresh.lock(plain))
|
||||
return fresh
|
||||
@@ -0,0 +1,130 @@
|
||||
"""Shell side-effect snapshots and output spill helpers."""
|
||||
|
||||
import glob
|
||||
import os
|
||||
import re
|
||||
import shlex
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["ShellSnap"]
|
||||
|
||||
|
||||
class ShellSnap:
|
||||
"""ShellSnap."""
|
||||
|
||||
@staticmethod
|
||||
def shell_target_paths(command, workdir):
|
||||
spaced = re.sub(r"([;&|()<>])", r" \1 ", command)
|
||||
try:
|
||||
tokens = shlex.split(spaced, posix=True)
|
||||
except ValueError:
|
||||
return []
|
||||
found = []
|
||||
segment = []
|
||||
|
||||
def flush():
|
||||
found.extend(ShellSnap.segment_targets(segment))
|
||||
|
||||
for token in tokens:
|
||||
if token in Consts.SHELL_SEGMENT:
|
||||
flush()
|
||||
segment = []
|
||||
else:
|
||||
segment.append(token)
|
||||
flush()
|
||||
resolved = []
|
||||
for item in found:
|
||||
if len(resolved) >= Consts.SHELL_SNAPSHOT_MAX_FILES:
|
||||
break
|
||||
base = item if os.path.isabs(item) else os.path.join(workdir, item)
|
||||
if any(mark in item for mark in ("*", "?", "[")):
|
||||
try:
|
||||
hits = sorted(glob.glob(base, recursive=True))
|
||||
except (OSError, ValueError):
|
||||
continue
|
||||
else:
|
||||
hits = [base]
|
||||
for hit in hits:
|
||||
if len(resolved) >= Consts.SHELL_SNAPSHOT_MAX_FILES:
|
||||
break
|
||||
if os.path.isfile(hit):
|
||||
resolved.append(os.path.abspath(hit))
|
||||
elif os.path.isdir(hit):
|
||||
for root, _dirs, files in os.walk(hit):
|
||||
for entry in sorted(files):
|
||||
if len(resolved) >= Consts.SHELL_SNAPSHOT_MAX_FILES:
|
||||
break
|
||||
full = os.path.join(root, entry)
|
||||
if os.path.isfile(full):
|
||||
resolved.append(os.path.abspath(full))
|
||||
if len(resolved) >= Consts.SHELL_SNAPSHOT_MAX_FILES:
|
||||
break
|
||||
seen = []
|
||||
for item in resolved:
|
||||
if item not in seen:
|
||||
seen.append(item)
|
||||
return seen
|
||||
|
||||
@staticmethod
|
||||
def segment_targets(tokens):
|
||||
targets = []
|
||||
pos = 0
|
||||
while pos < len(tokens):
|
||||
token = tokens[pos]
|
||||
if "=" in token and not token.startswith((">", "<", "-", "/")):
|
||||
pos += 1
|
||||
elif token in Consts.SHELL_WRAPPERS:
|
||||
pos += 1
|
||||
elif token == "timeout" and pos + 1 < len(tokens) and re.fullmatch(r"[0-9.]+[smhd]?", tokens[pos + 1] or ""):
|
||||
pos += 2
|
||||
else:
|
||||
break
|
||||
rest = tokens[pos:]
|
||||
for index, token in enumerate(rest):
|
||||
if token in (">", ">>") and index + 1 < len(rest):
|
||||
cand = rest[index + 1]
|
||||
if cand and not cand.startswith("&") and not cand.startswith("/dev/"):
|
||||
targets.append(cand)
|
||||
if not rest:
|
||||
return targets
|
||||
name = os.path.basename(rest[0])
|
||||
operands = [item for item in rest[1:] if not item.startswith("-") or item in ("-", "--")]
|
||||
operands = [item for item in operands if item not in (">", ">>", "<", "--")]
|
||||
if name == "rm":
|
||||
targets.extend(operands)
|
||||
elif name in ("mv", "cp") and operands:
|
||||
targets.append(operands[-1])
|
||||
elif name == "tee":
|
||||
targets.extend(operands)
|
||||
elif name == "truncate":
|
||||
for item in operands:
|
||||
if re.fullmatch(r"[0-9]+[KMGTPE]?[iB]?", item or ""):
|
||||
continue
|
||||
targets.append(item)
|
||||
elif name == "dd":
|
||||
for item in operands:
|
||||
if item.startswith("of="):
|
||||
targets.append(item[3:])
|
||||
elif name == "shred":
|
||||
targets.extend(operands)
|
||||
return [item for item in targets if item and item != "-"]
|
||||
|
||||
@staticmethod
|
||||
def spill_output(store, kind, title, text, tags=(), profile=None):
|
||||
profile = profile if profile is not None else store.profile
|
||||
record_id = store.add_record(kind, title, text, tags, profile)
|
||||
store.add_chunks("record", record_id, Text.chunk_lines(text), profile)
|
||||
return record_id
|
||||
|
||||
@staticmethod
|
||||
def spilled_view(text, record_id, size, head=1500, tail=500):
|
||||
return text[:head] + "\n...[%d chars spilled to %s, record_read pages the rest]...\n" % (size, record_id) + text[-tail:]
|
||||
|
||||
@staticmethod
|
||||
def spill_or_truncate(store, kind, title, text, tags=(), profile=None):
|
||||
if store is None or len(text) <= Consts.SPILL_LIMIT:
|
||||
return Text.truncate(text)
|
||||
record_id = ShellSnap.spill_output(store, kind, title, text, tags, profile)
|
||||
return ShellSnap.spilled_view(text, record_id, len(text))
|
||||
@@ -0,0 +1,166 @@
|
||||
"""SKILL.md discovery, install metadata, and creation prompts."""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
|
||||
__all__ = ["Skills"]
|
||||
|
||||
|
||||
class Skills:
|
||||
"""Skills."""
|
||||
|
||||
@staticmethod
|
||||
def parse_skill_front(lines):
|
||||
front = {}
|
||||
pos = 0
|
||||
while pos < len(lines):
|
||||
line = lines[pos]
|
||||
if ":" in line and not line.startswith((" ", "\t")):
|
||||
key, value = line.split(":", 1)
|
||||
value = value.strip()
|
||||
if value in (">", "|"):
|
||||
joined = []
|
||||
pos += 1
|
||||
while pos < len(lines) and lines[pos].startswith((" ", "\t")):
|
||||
joined.append(lines[pos].strip())
|
||||
pos += 1
|
||||
front[key.strip()] = (" " if value == ">" else "\n").join(joined)
|
||||
continue
|
||||
front[key.strip()] = value.strip("'\"")
|
||||
pos += 1
|
||||
return front
|
||||
|
||||
@staticmethod
|
||||
def parse_skill_file(path):
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
raw = handle.read()
|
||||
if not raw.startswith("---"):
|
||||
return None
|
||||
end = raw.find("\n---", 3)
|
||||
if end < 0:
|
||||
return None
|
||||
front = Skills.parse_skill_front(raw[3:end].strip().splitlines())
|
||||
name = front.get("name", "")
|
||||
description = front.get("description", "")
|
||||
if not name or not description or not Consts.SKILL_NAME_RE.match(name):
|
||||
return None
|
||||
return {"name": name, "description": description, "body": raw[end + 4:].lstrip("\n")}
|
||||
|
||||
@staticmethod
|
||||
def skill_scope_dir(scope, home_dir, project_dir, name=""):
|
||||
"""Shared by create_skill and install_skill so both land skills in the same place."""
|
||||
base = os.path.join(project_dir, ".tai", "skills") if scope == "project" else os.path.join(home_dir, "skills")
|
||||
return os.path.join(base, name) if name else base
|
||||
|
||||
@staticmethod
|
||||
def default_skill_scope(project_dir):
|
||||
for marker in (".git", ".tai") + Consts.EXTERNAL_SKILL_DIRNAMES:
|
||||
if os.path.exists(os.path.join(project_dir, marker)):
|
||||
return "project"
|
||||
return "home"
|
||||
|
||||
@staticmethod
|
||||
def skill_extra_files(root):
|
||||
extras = []
|
||||
for sub in ("scripts", "references", "assets"):
|
||||
folder = os.path.join(root, sub)
|
||||
if os.path.isdir(folder):
|
||||
for entry in sorted(os.listdir(folder)):
|
||||
extras.append(os.path.join(root, sub, entry))
|
||||
return extras
|
||||
|
||||
@staticmethod
|
||||
def skill_install_meta_path(scope_dir, entry):
|
||||
return os.path.join(scope_dir, Consts.SKILL_INSTALL_META_DIRNAME, entry + ".json")
|
||||
|
||||
@staticmethod
|
||||
def load_skill_install_meta(scope_dir, entry):
|
||||
try:
|
||||
with open(Skills.skill_install_meta_path(scope_dir, entry), "r", encoding="utf-8") as handle:
|
||||
return json.load(handle)
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def save_skill_install_meta(scope_dir, entry, meta):
|
||||
path = Skills.skill_install_meta_path(scope_dir, entry)
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
with open(path, "w", encoding="utf-8") as handle:
|
||||
json.dump(meta, handle)
|
||||
|
||||
@staticmethod
|
||||
def is_external_skill_root(path):
|
||||
return os.path.dirname(os.path.normpath(path)).endswith(Consts.EXTERNAL_SKILL_DIRNAMES)
|
||||
|
||||
@staticmethod
|
||||
def skill_provenance_note(skill):
|
||||
install = skill.get("install")
|
||||
if install:
|
||||
return "[installed %s, %s scope, from %s]" % (install.get("mode", "copy"), install.get("scope", "?"), install.get("source", "?"))
|
||||
if not skill.get("managed", True):
|
||||
return "[external, not managed by tai: %s]" % skill["root"]
|
||||
return "[native]"
|
||||
|
||||
@staticmethod
|
||||
def discover_skills(home_dir, project_dir):
|
||||
found = {}
|
||||
tai_roots = (os.path.join(home_dir, "skills"), os.path.join(project_dir, ".tai", "skills"))
|
||||
external_roots = (os.path.join(os.path.expanduser("~"), ".claude", "skills"),) + tuple(os.path.join(project_dir, dirname) for dirname in Consts.EXTERNAL_SKILL_DIRNAMES)
|
||||
# External roots scan first so tai's own dirs always win on a name collision (dict assignment below overwrites).
|
||||
for base in external_roots + tai_roots:
|
||||
if not os.path.isdir(base):
|
||||
continue
|
||||
managed = base in tai_roots
|
||||
for entry in sorted(os.listdir(base)):
|
||||
path = os.path.join(base, entry, "SKILL.md")
|
||||
if not os.path.isfile(path):
|
||||
continue
|
||||
try:
|
||||
skill = Skills.parse_skill_file(path)
|
||||
except OSError:
|
||||
continue
|
||||
if skill is None:
|
||||
continue
|
||||
skill["root"] = os.path.join(base, entry)
|
||||
skill["managed"] = managed
|
||||
skill["install"] = Skills.load_skill_install_meta(base, entry) if managed else None
|
||||
found[skill["name"]] = skill
|
||||
return found
|
||||
|
||||
@staticmethod
|
||||
def skill_catalog(skills):
|
||||
lines = []
|
||||
if skills:
|
||||
lines += ["", "", "## Available skills"]
|
||||
for name in sorted(skills):
|
||||
lines.append("- %s: %s" % (name, skills[name]["description"][:300]))
|
||||
wanted = sorted(name for name in Consts.SKILL_BLUEPRINTS if name not in (skills or {}))
|
||||
if wanted:
|
||||
lines += ["", "", "## Buildable skill blueprints (load_skill builds one on demand)"]
|
||||
for name in wanted:
|
||||
lines.append("- %s: %s" % (name, Consts.SKILL_BLUEPRINTS[name]["hint"]))
|
||||
return "\n".join(lines)
|
||||
lines.append("Load full instructions with load_skill before using one.")
|
||||
return "\n".join(lines)
|
||||
|
||||
@staticmethod
|
||||
def create_skill_prompt(name, brief, skill_dir):
|
||||
return (
|
||||
"Create a new agent skill named '%s' in %s.\n"
|
||||
"Brief: %s\n\n"
|
||||
"Follow this process exactly:\n"
|
||||
"1. Call sysinfo first so the skill matches this machine (os, installed tools, root, container, venv).\n"
|
||||
"2. Research the subject deeply with web_search, web_fetch, and shell. Keep researching until you literally "
|
||||
"have enough verified information to write complete instructions. Verify factual claims against at least two "
|
||||
"independent sources. Never invent commands, paths, versions, or API details.\n"
|
||||
"3. Write SKILL.md in the skill directory with this exact frontmatter:\n"
|
||||
"---\nname: %s\ndescription: <one sentence saying what the skill does and when to use it>\n---\n"
|
||||
"followed by clear markdown instructions grounded in your research.\n"
|
||||
"4. Add scripts/, references/, or assets/ files only when they carry real weight; keep every file focused.\n"
|
||||
"5. Read back every file you wrote, fix mistakes, and confirm the frontmatter holds a name plus a description.\n"
|
||||
"6. Reply with a short summary: what the skill covers, which sources verified it, and which files you wrote.\n"
|
||||
"Do not stop early: an incomplete or unverified skill is a failure. Do not ask questions; decide and record assumptions in the skill body."
|
||||
% (name, skill_dir, brief, name)
|
||||
)
|
||||
@@ -0,0 +1,968 @@
|
||||
"""SQLite vault storage, sessions, FTS, and migrations."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.release import Release
|
||||
from tai_pkg.seal import Seal
|
||||
from tai_pkg.seal import SealError
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["Store"]
|
||||
|
||||
|
||||
class Store:
|
||||
def __init__(self, config, seal):
|
||||
self.config = config
|
||||
self.seal = seal
|
||||
self.db = sqlite3.connect(config.db_path, timeout=30)
|
||||
self.db.execute("PRAGMA journal_mode=WAL")
|
||||
self.db.execute("PRAGMA synchronous=NORMAL")
|
||||
self.db.execute("PRAGMA busy_timeout=30000")
|
||||
self.schema_notes = Store.ensure_schema(self.db)
|
||||
Store.migrate_profiles(self.db)
|
||||
self.db.commit()
|
||||
self.profile = config.profile
|
||||
self.secret_cache = {}
|
||||
self.fts_ok = Store.setup_fts(self.db)
|
||||
if seal.enabled:
|
||||
self.db.create_function("tai_enc", 1, seal.lock)
|
||||
self.db.create_function("tai_dec", 1, seal.safe_unlock)
|
||||
self.db.execute("UPDATE events SET text = tai_enc(text) WHERE text NOT LIKE 'tai1$%'")
|
||||
self.db.execute("UPDATE secrets SET value = tai_enc(value) WHERE value NOT LIKE 'tai1$%'")
|
||||
self.db.execute("UPDATE secrets SET meta = tai_enc(COALESCE(meta, '{}')) WHERE meta IS NULL OR meta NOT LIKE 'tai1$%'")
|
||||
self.db.execute("UPDATE schedules SET prompt = tai_enc(prompt) WHERE prompt NOT LIKE 'tai1$%'")
|
||||
self.db.commit()
|
||||
self.seal_files()
|
||||
else:
|
||||
sealed = self.db.execute("SELECT COUNT(*) FROM events WHERE text LIKE 'tai1$%'").fetchone()[0]
|
||||
if sealed:
|
||||
raise SealError("memory is sealed, set TAI_PASSPHRASE")
|
||||
locked = self.db.execute("SELECT COUNT(*) FROM secrets WHERE value LIKE 'tai1$%' OR meta LIKE 'tai1$%'").fetchone()[0]
|
||||
if locked:
|
||||
raise SealError("secrets are sealed, set TAI_PASSPHRASE")
|
||||
planned = self.db.execute("SELECT COUNT(*) FROM schedules WHERE prompt LIKE 'tai1$%'").fetchone()[0]
|
||||
if planned:
|
||||
raise SealError("schedules are sealed, set TAI_PASSPHRASE")
|
||||
self.memdb = None
|
||||
self.mem_events_max = 0
|
||||
if seal.enabled and self.fts_ok:
|
||||
try:
|
||||
self.memdb = Store.open_mem_events()
|
||||
floor = self.db.execute("SELECT COALESCE(MAX(id), 0) FROM events").fetchone()[0] - Consts.MEM_EVENTS_CAP
|
||||
self.mem_events_max = max(0, floor)
|
||||
self._sync_mem_events()
|
||||
except sqlite3.Error:
|
||||
self.memdb = None
|
||||
|
||||
def _sync_mem_events(self):
|
||||
if self.memdb is None:
|
||||
return
|
||||
try:
|
||||
rows = self.db.execute("SELECT id, profile, ts, role, kind, tai_dec(text) FROM events WHERE id > ? ORDER BY id ASC", (self.mem_events_max,)).fetchall()
|
||||
except sqlite3.Error:
|
||||
return
|
||||
if not rows:
|
||||
return
|
||||
try:
|
||||
self.memdb.executemany("INSERT INTO mem_events(item, kind, sub, title, body, profile) VALUES (?, 'event', ? || '/' || ?, ? || ' ' || ? || '/' || ?, ?, ?)", [(row[0], row[3], row[4], row[2], row[3], row[4], row[5], row[1]) for row in rows])
|
||||
self.memdb.commit()
|
||||
except sqlite3.Error:
|
||||
return
|
||||
self.mem_events_max = rows[-1][0]
|
||||
try:
|
||||
self.memdb.execute("DELETE FROM mem_events WHERE CAST(item AS INTEGER) <= ?", (self.mem_events_max - Consts.MEM_EVENTS_CAP,))
|
||||
self.memdb.commit()
|
||||
except sqlite3.Error:
|
||||
pass
|
||||
|
||||
def seal_files(self):
|
||||
for entry in os.listdir(self.config.profiles_dir):
|
||||
if not entry.endswith((".sys.md", ".session.json", ".bots.json")):
|
||||
continue
|
||||
path = os.path.join(self.config.profiles_dir, entry)
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
raw = handle.read()
|
||||
except OSError:
|
||||
continue
|
||||
if raw.startswith(Seal.PREFIX):
|
||||
continue
|
||||
try:
|
||||
with open(path, "w", encoding="utf-8") as handle:
|
||||
handle.write(self.seal.lock(raw))
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def log_event(self, profile, role, kind, text, tags=()):
|
||||
value = (text or "")[:2000]
|
||||
stamped = " " + " ".join([role, kind] + Text.normalize_tags(tags)) + " "
|
||||
try:
|
||||
if self.seal.enabled:
|
||||
self.db.execute("INSERT INTO events (profile, ts, role, kind, text, tags) VALUES (?, ?, ?, ?, tai_enc(?), ?)", (profile, Text.now_iso(), role, kind, value, stamped))
|
||||
else:
|
||||
self.db.execute("INSERT INTO events (profile, ts, role, kind, text, tags) VALUES (?, ?, ?, ?, ?, ?)", (profile, Text.now_iso(), role, kind, value, stamped))
|
||||
self.db.commit()
|
||||
self._sync_mem_events()
|
||||
except sqlite3.Error:
|
||||
pass
|
||||
|
||||
def fts_search(self, query, kinds=(), profile=None, limit=10):
|
||||
profile = self._scope(profile)
|
||||
if not self.fts_ok:
|
||||
return []
|
||||
terms = Store.fts_escape(query)
|
||||
if terms is None:
|
||||
return []
|
||||
sql = "SELECT item, kind, sub, title, snippet(fts_docs, 4, '[', ']', '...', 12), bm25(fts_docs) FROM fts_docs WHERE fts_docs MATCH ? AND profile = ?"
|
||||
params = [Store.fts_match(terms), profile]
|
||||
if kinds:
|
||||
marks = ", ".join("?" * len(kinds))
|
||||
sql += " AND kind IN (%s)" % marks
|
||||
params += list(kinds)
|
||||
if self.seal.enabled:
|
||||
sql += " AND kind != 'event'"
|
||||
sql += " ORDER BY bm25(fts_docs) LIMIT ?"
|
||||
capped = max(1, min(limit, 100))
|
||||
params.append(capped)
|
||||
try:
|
||||
rows = self.db.execute(sql, tuple(params)).fetchall()
|
||||
except sqlite3.Error:
|
||||
return []
|
||||
hits = [{"item": row[0], "kind": row[1], "sub": row[2], "title": row[3], "snippet": row[4], "rank": row[5]} for row in rows]
|
||||
if self.seal.enabled and self.memdb is not None and (not kinds or "event" in kinds):
|
||||
self._sync_mem_events()
|
||||
try:
|
||||
mem_rows = self.memdb.execute("SELECT item, kind, sub, title, snippet(mem_events, 4, '[', ']', '...', 12), bm25(mem_events) FROM mem_events WHERE mem_events MATCH ? AND profile = ? ORDER BY bm25(mem_events) LIMIT ?", (Store.fts_match(terms), profile, capped)).fetchall()
|
||||
except sqlite3.Error:
|
||||
mem_rows = []
|
||||
hits += [{"item": row[0], "kind": row[1], "sub": row[2], "title": row[3], "snippet": row[4], "rank": row[5]} for row in mem_rows]
|
||||
hits.sort(key=lambda hit: hit["rank"])
|
||||
hits = hits[:capped]
|
||||
return hits
|
||||
|
||||
def audit_search(self, query, profile=None, limit=20):
|
||||
profile = self._scope(profile)
|
||||
hits = self.fts_search(query, ("audit",), profile, limit)
|
||||
found = []
|
||||
for hit in hits:
|
||||
try:
|
||||
row_id = int(str(hit["item"]).split(":", 1)[1])
|
||||
except (IndexError, ValueError):
|
||||
continue
|
||||
row = self.audit_get(row_id, profile)
|
||||
if row is not None:
|
||||
row["snippet"] = hit["snippet"]
|
||||
found.append(row)
|
||||
return found
|
||||
|
||||
def search_events(self, profile, query, limit=8, tags=()):
|
||||
clean = Text.normalize_tags(tags)
|
||||
if self.fts_ok:
|
||||
hits = self.fts_search(query, ("event",), profile, max(limit * 5, 20))
|
||||
if hits:
|
||||
ids = [int(hit["item"]) for hit in hits]
|
||||
marks = ", ".join("?" * len(ids))
|
||||
text_col = "tai_dec(text)" if self.seal.enabled else "text"
|
||||
rows = {row[0]: row[1:] for row in self.db.execute("SELECT id, ts, role, kind, %s, tags FROM events WHERE id IN (%s)" % (text_col, marks), tuple(ids)).fetchall()}
|
||||
found = []
|
||||
for hit in hits:
|
||||
row = rows.get(int(hit["item"]))
|
||||
if row is None:
|
||||
continue
|
||||
if clean and not all((" %s " % tag) in (" %s " % (row[4] or "")) for tag in clean):
|
||||
continue
|
||||
found.append(row[:4])
|
||||
if len(found) >= limit:
|
||||
break
|
||||
if found:
|
||||
return found
|
||||
return self.events_like(profile, query, limit, clean)
|
||||
|
||||
def events_like(self, profile, query, limit=8, tags=()):
|
||||
clean = Text.normalize_tags(tags)
|
||||
extra = "".join(" AND tags LIKE ?" for _tag in clean)
|
||||
wild = tuple("% %s %%" % tag for tag in clean)
|
||||
if self.seal.enabled:
|
||||
return self.db.execute("SELECT ts, role, kind, text FROM (SELECT id, ts, role, kind, tags, tai_dec(text) AS text FROM events WHERE profile = ? ORDER BY id DESC LIMIT 5000) WHERE text LIKE ?" + extra + " ORDER BY id DESC LIMIT ?", (profile, "%" + query + "%") + wild + (limit,)).fetchall()
|
||||
return self.db.execute("SELECT ts, role, kind, text FROM events WHERE profile = ? AND text LIKE ?" + extra + " ORDER BY id DESC LIMIT ?", (profile, "%" + query + "%") + wild + (limit,)).fetchall()
|
||||
|
||||
def _scope(self, profile):
|
||||
return profile or self.profile
|
||||
|
||||
def save_secret(self, name, value, meta=None, tags=None, profile=None):
|
||||
profile = self._scope(profile)
|
||||
self.secret_cache.pop(profile, None)
|
||||
stored = self.seal.lock(value) if self.seal.enabled else value
|
||||
blob = json.dumps(meta or {})
|
||||
locked_meta = self.seal.lock(blob) if self.seal.enabled else blob
|
||||
self.db.execute("INSERT OR REPLACE INTO secrets (profile, name, value, meta, updated) VALUES (?, ?, ?, ?, ?)", (profile, name, stored, locked_meta, Text.now_iso()))
|
||||
if tags is None:
|
||||
self.tag_item("secret:" + name, ["secret"], profile)
|
||||
else:
|
||||
self.set_item_tags("secret:" + name, ["secret"] + list(tags), profile)
|
||||
self.db.commit()
|
||||
|
||||
def secret_meta(self, name, profile=None):
|
||||
profile = self._scope(profile)
|
||||
row = self.db.execute("SELECT meta FROM secrets WHERE profile = ? AND name = ?", (profile, name)).fetchone()
|
||||
if row is None or not row[0]:
|
||||
return {}
|
||||
raw = self.seal.safe_unlock(row[0]) if row[0].startswith(Seal.PREFIX) else row[0]
|
||||
try:
|
||||
data = json.loads(raw)
|
||||
except ValueError:
|
||||
return {}
|
||||
return data if isinstance(data, dict) else {}
|
||||
|
||||
def list_secret_infos(self, profile=None):
|
||||
profile = self._scope(profile)
|
||||
rows = self.db.execute("SELECT name, meta, updated FROM secrets WHERE profile = ? ORDER BY name", (profile,)).fetchall()
|
||||
items = []
|
||||
for name, meta, updated in rows:
|
||||
raw = None
|
||||
if meta:
|
||||
raw = self.seal.safe_unlock(meta) if meta.startswith(Seal.PREFIX) else meta
|
||||
try:
|
||||
data = json.loads(raw or "{}")
|
||||
except ValueError:
|
||||
data = {}
|
||||
items.append({"name": name, "meta": data if isinstance(data, dict) else {}, "tags": self.item_tags("secret:" + name, profile), "updated": updated})
|
||||
return items
|
||||
|
||||
def known_tags(self, profile=None):
|
||||
profile = self._scope(profile)
|
||||
return {row[0] for row in self.db.execute("SELECT DISTINCT tag FROM tags WHERE profile = ?", (profile,)).fetchall()}
|
||||
|
||||
def canonical_tag(self, tag, vocab=None, profile=None):
|
||||
if vocab is None:
|
||||
vocab = self.known_tags(profile)
|
||||
if tag in vocab:
|
||||
return tag
|
||||
single = Text.singular_noun(tag)
|
||||
if single != tag and single in vocab:
|
||||
return single
|
||||
return tag
|
||||
|
||||
def canonical_tags(self, tags, vocab=None, profile=None):
|
||||
profile = self._scope(profile)
|
||||
if vocab is None:
|
||||
vocab = self.known_tags(profile)
|
||||
seen = []
|
||||
for tag in Text.normalize_tags(tags):
|
||||
clean = self.canonical_tag(tag, vocab)
|
||||
if clean not in seen:
|
||||
seen.append(clean)
|
||||
return seen
|
||||
|
||||
def tag_counts(self, prefix="", limit=50, profile=None):
|
||||
profile = self._scope(profile)
|
||||
clean = Text.normalize_tag(prefix)
|
||||
sql = "SELECT tag, COUNT(*) FROM tags WHERE profile = ?"
|
||||
params = [profile]
|
||||
if clean:
|
||||
sql += " AND tag LIKE ? ESCAPE '\\'"
|
||||
params.append(clean.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + "%")
|
||||
sql += " GROUP BY tag ORDER BY COUNT(*) DESC, tag LIMIT ?"
|
||||
params.append(max(1, min(limit, 200)))
|
||||
return [(row[0], row[1]) for row in self.db.execute(sql, tuple(params)).fetchall()]
|
||||
|
||||
def auto_tags_for(self, text, explicit=(), vocab=None):
|
||||
if vocab is None:
|
||||
vocab = self.known_tags()
|
||||
skip = set(explicit) | set(Consts.TAG_AUTO_SKIP)
|
||||
words = re.findall(r"[a-z0-9]+", (text or "").lower())
|
||||
found = []
|
||||
for pos, word in enumerate(words):
|
||||
if len(found) >= Consts.TAG_AUTO_MAX:
|
||||
break
|
||||
if pos + 1 < len(words):
|
||||
pair = word + "-" + words[pos + 1]
|
||||
if pair in vocab and pair not in skip and pair not in found:
|
||||
found.append(pair)
|
||||
continue
|
||||
if word in vocab and word not in skip and word not in found:
|
||||
found.append(word)
|
||||
continue
|
||||
single = Text.singular_noun(word)
|
||||
if single != word and single in vocab and single not in skip and single not in found:
|
||||
found.append(single)
|
||||
return found
|
||||
|
||||
def tag_neighbors(self, record_id, limit=None, profile=None):
|
||||
profile = self._scope(profile)
|
||||
rows = self.db.execute("SELECT r.id, MIN(t.tag) FROM records r JOIN tags t ON t.item = r.id WHERE t.tag IN (SELECT tag FROM tags WHERE item = ? AND profile = ? AND tag NOT IN ('record', 'file')) AND r.id != ? AND r.profile = ? AND t.profile = ? GROUP BY r.id ORDER BY r.updated DESC LIMIT ?", (record_id, profile, record_id, profile, profile, max(1, limit or Consts.TAG_LINK_MAX))).fetchall()
|
||||
return [(row[0], row[1]) for row in rows]
|
||||
|
||||
def tag_item(self, item, tags, profile=None):
|
||||
profile = self._scope(profile)
|
||||
for tag in self.canonical_tags(tags, None, profile):
|
||||
self.db.execute("INSERT OR IGNORE INTO tags (item, tag, profile) VALUES (?, ?, ?)", (item, tag, profile))
|
||||
self.db.commit()
|
||||
|
||||
def set_item_tags(self, item, tags, profile=None):
|
||||
profile = self._scope(profile)
|
||||
self.db.execute("DELETE FROM tags WHERE item = ? AND profile = ?", (item, profile))
|
||||
for tag in self.canonical_tags(tags, None, profile):
|
||||
self.db.execute("INSERT OR IGNORE INTO tags (item, tag, profile) VALUES (?, ?, ?)", (item, tag, profile))
|
||||
self.db.commit()
|
||||
|
||||
def item_tags(self, item, profile=None):
|
||||
profile = self._scope(profile)
|
||||
return [row[0] for row in self.db.execute("SELECT tag FROM tags WHERE item = ? AND profile = ? ORDER BY tag", (item, profile)).fetchall()]
|
||||
|
||||
def tagged_items(self, tags, profile=None):
|
||||
profile = self._scope(profile)
|
||||
clean = Text.normalize_tags(tags)
|
||||
if not clean:
|
||||
return None
|
||||
marks = ", ".join("?" * len(clean))
|
||||
rows = self.db.execute("SELECT item FROM tags WHERE profile = ? AND tag IN (%s) GROUP BY item HAVING COUNT(DISTINCT tag) = ?" % marks, (profile,) + tuple(clean) + (len(clean),)).fetchall()
|
||||
return [row[0] for row in rows]
|
||||
|
||||
def has_node(self, kind, key, profile=None):
|
||||
profile = self._scope(profile)
|
||||
if kind == "mem":
|
||||
row = self.db.execute("SELECT 1 FROM records WHERE profile = ? AND id = ?", (profile, "mem:" + key,)).fetchone()
|
||||
elif kind == "secret":
|
||||
row = self.db.execute("SELECT 1 FROM secrets WHERE profile = ? AND name = ?", (profile, key,)).fetchone()
|
||||
else:
|
||||
row = self.db.execute("SELECT 1 FROM schedules WHERE profile = ? AND id = ?", (profile, int(key),)).fetchone()
|
||||
return row is not None
|
||||
|
||||
def node_title(self, node, profile=None):
|
||||
profile = self._scope(profile)
|
||||
parsed = Text.parse_node_id(node)
|
||||
if parsed is None:
|
||||
return node
|
||||
kind, key = parsed
|
||||
if kind == "mem":
|
||||
row = self.db.execute("SELECT title FROM records WHERE profile = ? AND id = ?", (profile, node,)).fetchone()
|
||||
return row[0] or node if row else node
|
||||
if kind == "secret":
|
||||
return "secret:" + key
|
||||
row = self.db.execute("SELECT name, prompt FROM schedules WHERE profile = ? AND id = ?", (profile, int(key),)).fetchone()
|
||||
if not row:
|
||||
return node
|
||||
return row[0] or "schedule #%s" % key
|
||||
|
||||
def add_record(self, kind, title, content, tags=(), profile=None):
|
||||
profile = self._scope(profile)
|
||||
record_id = "mem:" + uuid.uuid4().hex[:16]
|
||||
while self.db.execute("SELECT 1 FROM records WHERE id = ?", (record_id,)).fetchone():
|
||||
record_id = "mem:" + uuid.uuid4().hex[:16]
|
||||
stamp = datetime.now(timezone.utc).isoformat()
|
||||
size = len(content.encode("utf-8"))
|
||||
self.db.execute("INSERT INTO records(id, kind, profile, title, content, size, reads, created, updated) VALUES(?, ?, ?, ?, ?, ?, 0, ?, ?)", (record_id, kind, profile, title, content, size, stamp, stamp))
|
||||
vocab = self.known_tags(profile)
|
||||
explicit = self.canonical_tags(tags, vocab, profile)
|
||||
scan = "%s\n%s" % (title, content.split("\n...[truncated", 1)[0])
|
||||
auto = self.auto_tags_for(scan, explicit, vocab)
|
||||
self.set_item_tags(record_id, ["record"] + explicit + auto, profile)
|
||||
for other, shared in self.tag_neighbors(record_id, None, profile):
|
||||
self.db.execute("INSERT OR IGNORE INTO edges(src, dst, relation, profile, created) VALUES(?, ?, ?, ?, ?)", (record_id, other, "shares-" + shared, profile, stamp))
|
||||
self.db.commit()
|
||||
return record_id
|
||||
|
||||
def add_chunks(self, parent_kind, parent_id, texts, profile=None):
|
||||
"""Replaces any prior chunks for this parent, so re-indexing is idempotent."""
|
||||
profile = self._scope(profile)
|
||||
self.db.execute("DELETE FROM chunks WHERE parent_kind = ? AND parent_id = ? AND profile = ?", (parent_kind, parent_id, profile))
|
||||
stamp = datetime.now(timezone.utc).isoformat()
|
||||
for idx, text in enumerate(texts):
|
||||
self.db.execute("INSERT INTO chunks(parent_kind, parent_id, idx, text, profile, created) VALUES (?, ?, ?, ?, ?, ?)", (parent_kind, parent_id, idx, text, profile, stamp))
|
||||
self.db.commit()
|
||||
|
||||
def get_chunks_by_ids(self, ids, profile=None):
|
||||
if not ids:
|
||||
return {}
|
||||
profile = self._scope(profile)
|
||||
marks = ", ".join("?" * len(ids))
|
||||
rows = self.db.execute("SELECT id, parent_kind, parent_id, idx, text FROM chunks WHERE profile = ? AND id IN (%s)" % marks, tuple([profile] + list(ids))).fetchall()
|
||||
return {row[0]: {"parent_kind": row[1], "parent_id": row[2], "idx": row[3], "text": row[4]} for row in rows}
|
||||
|
||||
def get_record(self, record_id, profile=None):
|
||||
profile = self._scope(profile)
|
||||
row = self.db.execute("SELECT id, kind, title, content, size, reads, created, updated FROM records WHERE profile = ? AND id = ?", (profile, record_id,)).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
self.db.execute("UPDATE records SET reads = reads + 1 WHERE id = ?", (record_id,))
|
||||
self.db.commit()
|
||||
return {"id": row[0], "kind": row[1], "title": row[2], "content": row[3], "size": Text.as_int(row[4]), "reads": Text.as_int(row[5]) + 1, "created": row[6], "updated": row[7], "tags": self.item_tags(row[0], profile)}
|
||||
|
||||
def read_record(self, record_id, offset=0, limit=4000, profile=None):
|
||||
profile = self._scope(profile)
|
||||
row = self.db.execute("SELECT content FROM records WHERE profile = ? AND id = ?", (profile, record_id,)).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
content = row[0]
|
||||
total = len(content)
|
||||
start = max(0, min(offset, total))
|
||||
end = min(total, start + max(1, limit))
|
||||
self.db.execute("UPDATE records SET reads = reads + 1 WHERE id = ?", (record_id,))
|
||||
self.db.commit()
|
||||
return {"id": record_id, "slice": content[start:end], "start": start, "end": end, "total": total, "tags": self.item_tags(record_id, profile)}
|
||||
|
||||
def search_records(self, query="", kind=None, tags=(), limit=8, profile=None):
|
||||
profile = self._scope(profile)
|
||||
clean = Text.normalize_tags(tags)
|
||||
terms = Store.fts_escape(query) if query else None
|
||||
if terms is not None and self.fts_ok:
|
||||
ranked = self.search_records_fts(terms, kind, clean, limit, profile)
|
||||
if ranked:
|
||||
return ranked
|
||||
sql = "SELECT id, kind, title, size, reads, updated FROM records"
|
||||
clauses = ["profile = ?"]
|
||||
params = [profile]
|
||||
if kind:
|
||||
clauses.append("kind = ?")
|
||||
params.append(kind)
|
||||
if query:
|
||||
clauses.append("(title LIKE ? OR content LIKE ?)")
|
||||
params += ["%" + query + "%", "%" + query + "%"]
|
||||
if clean:
|
||||
for tag in clean:
|
||||
options = Text.tag_variants(tag)
|
||||
marks = ", ".join("?" * len(options))
|
||||
clauses.append("id IN (SELECT item FROM tags WHERE profile = ? AND tag IN (%s))" % marks)
|
||||
params += [profile] + options
|
||||
sql += " WHERE " + " AND ".join(clauses)
|
||||
sql += " ORDER BY updated DESC LIMIT ?"
|
||||
params.append(max(1, min(limit, 50)))
|
||||
rows = self.db.execute(sql, tuple(params)).fetchall()
|
||||
return [{"id": row[0], "kind": row[1], "title": row[2], "size": row[3], "reads": row[4], "updated": row[5], "tags": self.item_tags(row[0], profile)} for row in rows]
|
||||
|
||||
def search_records_fts(self, terms, kind, clean, limit, profile):
|
||||
sql = "SELECT r.id, r.kind, r.title, r.size, r.reads, r.updated FROM records r JOIN fts_docs f ON f.item = r.id AND f.kind = 'record' WHERE r.profile = ? AND f MATCH ?"
|
||||
params = [profile, Store.fts_match(terms)]
|
||||
if kind:
|
||||
sql += " AND r.kind = ?"
|
||||
params.append(kind)
|
||||
for tag in clean:
|
||||
options = Text.tag_variants(tag)
|
||||
marks = ", ".join("?" * len(options))
|
||||
sql += " AND r.id IN (SELECT item FROM tags WHERE profile = ? AND tag IN (%s))" % marks
|
||||
params += [profile] + options
|
||||
sql += " ORDER BY bm25(f) LIMIT ?"
|
||||
params.append(max(1, min(limit, 50)))
|
||||
try:
|
||||
rows = self.db.execute(sql, tuple(params)).fetchall()
|
||||
except sqlite3.Error:
|
||||
return []
|
||||
return [{"id": row[0], "kind": row[1], "title": row[2], "size": row[3], "reads": row[4], "updated": row[5], "tags": self.item_tags(row[0], profile)} for row in rows]
|
||||
|
||||
def delete_record(self, record_id, profile=None):
|
||||
profile = self._scope(profile)
|
||||
done = self.db.execute("DELETE FROM records WHERE profile = ? AND id = ?", (profile, record_id,))
|
||||
if done.rowcount > 0:
|
||||
self.db.execute("DELETE FROM tags WHERE item = ? AND profile = ?", (record_id, profile))
|
||||
self.db.execute("DELETE FROM edges WHERE profile = ? AND (src = ? OR dst = ?)", (profile, record_id, record_id))
|
||||
self.db.execute("DELETE FROM chunks WHERE parent_kind = 'record' AND parent_id = ? AND profile = ?", (record_id, profile))
|
||||
self.db.commit()
|
||||
return True
|
||||
self.db.commit()
|
||||
return False
|
||||
|
||||
def add_edge(self, src, dst, relation="linked", profile=None):
|
||||
profile = self._scope(profile)
|
||||
for node in (src, dst):
|
||||
parsed = Text.parse_node_id(node)
|
||||
if parsed is None or not self.has_node(*parsed, profile):
|
||||
raise ValueError("unknown node " + node)
|
||||
clean = re.sub(r"[^a-z0-9]+", "-", str(relation or "").lower()).strip("-")[:32] or "linked"
|
||||
self.db.execute("INSERT OR IGNORE INTO edges(src, dst, relation, profile, created) VALUES(?, ?, ?, ?, ?)", (src, dst, clean, profile, datetime.now(timezone.utc).isoformat()))
|
||||
self.db.commit()
|
||||
return clean
|
||||
|
||||
def edges_for(self, node, profile=None):
|
||||
profile = self._scope(profile)
|
||||
rows = self.db.execute("SELECT src, dst, relation FROM edges WHERE profile = ? AND (src = ? OR dst = ?)", (profile, node, node)).fetchall()
|
||||
return [{"relation": row[2], "other": row[1] if row[0] == node else row[0], "direction": "out" if row[0] == node else "in"} for row in rows]
|
||||
|
||||
def remove_edges(self, node, profile=None):
|
||||
profile = self._scope(profile)
|
||||
done = self.db.execute("DELETE FROM edges WHERE profile = ? AND (src = ? OR dst = ?)", (profile, node, node))
|
||||
self.db.commit()
|
||||
return done.rowcount
|
||||
|
||||
def traverse(self, start, depth=2, limit=50, profile=None):
|
||||
profile = self._scope(profile)
|
||||
seen = {start}
|
||||
found = [{"node": start, "title": self.node_title(start, profile), "depth": 0, "via": ""}]
|
||||
frontier = [(start, 0, "")]
|
||||
capped = max(1, min(depth, 4))
|
||||
while frontier and len(found) < max(1, min(limit, 200)):
|
||||
node, level, _via = frontier.pop(0)
|
||||
if level >= capped:
|
||||
continue
|
||||
for edge in self.edges_for(node, profile):
|
||||
other = edge["other"]
|
||||
if other in seen:
|
||||
continue
|
||||
seen.add(other)
|
||||
hop = "%s -[%s]-> %s" % (node, edge["relation"], other) if edge["direction"] == "out" else "%s <-[%s]- %s" % (node, edge["relation"], other)
|
||||
found.append({"node": other, "title": self.node_title(other, profile), "depth": level + 1, "via": hop})
|
||||
frontier.append((other, level + 1, hop))
|
||||
if len(found) >= max(1, min(limit, 200)):
|
||||
break
|
||||
return found
|
||||
|
||||
def audit_event(self, actor, action, path, message="", old=None, new=None, tags=(), old_size=None, new_size=None, profile=None):
|
||||
profile = self._scope(profile)
|
||||
old_text, computed_old = Text.capped_text(old)
|
||||
new_text, computed_new = Text.capped_text(new)
|
||||
if old is not None and old_size is not None and old_size > len(old.encode("utf-8")):
|
||||
old_text = old[:Consts.AUDIT_MAX_CHARS] + Text.trunc_marker(old_size)
|
||||
computed_old = old_size
|
||||
if new is not None and new_size is not None and new_size > len(new.encode("utf-8")):
|
||||
new_text = new[:Consts.AUDIT_MAX_CHARS] + Text.trunc_marker(new_size)
|
||||
computed_new = new_size
|
||||
cursor = self.db.execute("INSERT INTO audit(profile, ts, actor, action, path, message, old_size, new_size, old, new, tags) VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", (profile, datetime.now(timezone.utc).isoformat(), actor, action, path, str(message or "")[:2000], computed_old, computed_new, old_text, new_text, " ".join(Text.normalize_tags(tags))))
|
||||
self.db.commit()
|
||||
return cursor.lastrowid
|
||||
|
||||
def audit_history(self, path=None, tag=None, limit=20, profile=None):
|
||||
profile = self._scope(profile)
|
||||
sql = "SELECT id, ts, actor, action, path, message, old_size, new_size, tags FROM audit"
|
||||
clauses = ["profile = ?"]
|
||||
params = [profile]
|
||||
if path:
|
||||
clauses.append("path = ?")
|
||||
params.append(path)
|
||||
if tag:
|
||||
clauses.append("tags LIKE ?")
|
||||
params.append("%" + tag + "%")
|
||||
sql += " WHERE " + " AND ".join(clauses)
|
||||
sql += " ORDER BY id DESC LIMIT ?"
|
||||
params.append(max(1, min(limit, 100)))
|
||||
rows = self.db.execute(sql, tuple(params)).fetchall()
|
||||
return [{"id": row[0], "ts": row[1], "actor": row[2], "action": row[3], "path": row[4], "message": row[5], "old_size": row[6], "new_size": row[7], "tags": row[8]} for row in rows]
|
||||
|
||||
def audit_get(self, row_id, profile=None):
|
||||
profile = self._scope(profile)
|
||||
row = self.db.execute("SELECT id, ts, actor, action, path, message, old_size, new_size, old, new, tags FROM audit WHERE profile = ? AND id = ?", (profile, row_id,)).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
return {"id": row[0], "ts": row[1], "actor": row[2], "action": row[3], "path": row[4], "message": row[5], "old_size": row[6], "new_size": row[7], "old": row[8], "new": row[9], "tags": row[10]}
|
||||
|
||||
def file_record_id(self, path, profile=None):
|
||||
profile = self._scope(profile)
|
||||
row = self.db.execute("SELECT id FROM records WHERE profile = ? AND kind = 'file' AND title = ?", (profile, path,)).fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
def upsert_file_record(self, path, content, true_size=None, profile=None):
|
||||
profile = self._scope(profile)
|
||||
full_bytes = len(content.encode("utf-8"))
|
||||
if true_size is not None and true_size > full_bytes:
|
||||
text = content + Text.trunc_marker(true_size)
|
||||
elif len(content) <= Consts.FILE_RECORD_MAX:
|
||||
text = content
|
||||
else:
|
||||
text = content[:Consts.FILE_RECORD_MAX] + Text.trunc_marker(full_bytes)
|
||||
ext = os.path.splitext(path)[1].lstrip(".").lower()
|
||||
explicit = ["file"] + ([ext] if ext and ext.isalnum() and len(ext) <= 5 else [])
|
||||
vocab = self.known_tags(profile)
|
||||
tags = explicit + self.auto_tags_for("%s\n%s" % (path, content), explicit, vocab)
|
||||
record_id = self.file_record_id(path, profile)
|
||||
stamp = datetime.now(timezone.utc).isoformat()
|
||||
if record_id is None:
|
||||
return self.add_record("file", path, text, tags, profile)
|
||||
self.db.execute("UPDATE records SET content = ?, size = ?, updated = ? WHERE id = ?", (text, len(text.encode("utf-8")), stamp, record_id))
|
||||
self.set_item_tags(record_id, tags, profile)
|
||||
self.db.commit()
|
||||
return record_id
|
||||
|
||||
def delete_file_record(self, path, profile=None):
|
||||
profile = self._scope(profile)
|
||||
record_id = self.file_record_id(path, profile)
|
||||
if record_id is None:
|
||||
return False
|
||||
return self.delete_record(record_id, profile)
|
||||
|
||||
def load_secret(self, name, profile=None):
|
||||
profile = self._scope(profile)
|
||||
row = self.db.execute("SELECT value FROM secrets WHERE profile = ? AND name = ?", (profile, name,)).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
if row[0].startswith(Seal.PREFIX):
|
||||
return self.seal.unlock(row[0]) if self.seal.enabled else None
|
||||
return row[0]
|
||||
|
||||
def delete_secret(self, name, profile=None):
|
||||
profile = self._scope(profile)
|
||||
self.secret_cache.pop(profile, None)
|
||||
done = self.db.execute("DELETE FROM secrets WHERE profile = ? AND name = ?", (profile, name,))
|
||||
if done.rowcount > 0:
|
||||
item = "secret:" + name
|
||||
self.db.execute("DELETE FROM tags WHERE item = ? AND profile = ?", (item, profile))
|
||||
self.db.execute("DELETE FROM edges WHERE profile = ? AND (src = ? OR dst = ?)", (profile, item, item))
|
||||
self.db.commit()
|
||||
return True
|
||||
self.db.commit()
|
||||
return False
|
||||
|
||||
def secret_values(self, profile=None):
|
||||
profile = self._scope(profile)
|
||||
if profile not in self.secret_cache:
|
||||
found = {}
|
||||
for name, value in self.db.execute("SELECT name, value FROM secrets WHERE profile = ?", (profile,)).fetchall():
|
||||
if value.startswith(Seal.PREFIX):
|
||||
if not self.seal.enabled:
|
||||
continue
|
||||
value = self.seal.unlock(value)
|
||||
if len(value) >= 4:
|
||||
found[name] = value
|
||||
self.secret_cache[profile] = found
|
||||
return self.secret_cache[profile]
|
||||
|
||||
def redact(self, text, profile=None):
|
||||
if not isinstance(text, str):
|
||||
return text
|
||||
for name, value in sorted(self.secret_values(profile).items(), key=lambda item: -len(item[1])):
|
||||
if value in text:
|
||||
text = text.replace(value, "[redacted:%s]" % name)
|
||||
return text
|
||||
|
||||
def add_schedule(self, name, prompt, profile, every_sec, next_run, timeout, tags=None):
|
||||
stored = self.seal.lock(prompt) if self.seal.enabled else prompt
|
||||
done = self.db.execute("INSERT INTO schedules (name, prompt, profile, every_sec, next_run, timeout, status, last_status, last_result, created, updated) VALUES (?, ?, ?, ?, ?, ?, 'pending', '', '', ?, ?)", (name, stored, profile, every_sec, next_run, timeout, Text.now_iso(), Text.now_iso()))
|
||||
item = "sched:%s" % done.lastrowid
|
||||
if tags is None:
|
||||
self.tag_item(item, ["schedule"], profile)
|
||||
else:
|
||||
self.set_item_tags(item, ["schedule"] + list(tags), profile)
|
||||
self.db.commit()
|
||||
return done.lastrowid
|
||||
|
||||
def remove_schedule(self, row_id, profile=None):
|
||||
profile = self._scope(profile)
|
||||
done = self.db.execute("DELETE FROM schedules WHERE id = ? AND profile = ?", (row_id, profile))
|
||||
if done.rowcount > 0:
|
||||
item = "sched:%s" % row_id
|
||||
self.db.execute("DELETE FROM tags WHERE item = ? AND profile = ?", (item, profile))
|
||||
self.db.execute("DELETE FROM edges WHERE profile = ? AND (src = ? OR dst = ?)", (profile, item, item))
|
||||
self.db.commit()
|
||||
return True
|
||||
self.db.commit()
|
||||
return False
|
||||
|
||||
def list_schedules(self, profile=None):
|
||||
profile = self._scope(profile)
|
||||
rows = self.db.execute("SELECT id, name, prompt, profile, every_sec, next_run, timeout, status, last_status, last_result, created FROM schedules WHERE profile = ? ORDER BY next_run", (profile,)).fetchall()
|
||||
items = []
|
||||
for row in rows:
|
||||
prompt = self.seal.safe_unlock(row[2]) if row[2].startswith(Seal.PREFIX) else row[2]
|
||||
items.append({"id": row[0], "name": row[1], "prompt": prompt, "profile": row[3], "every": row[4], "next_run": row[5], "timeout": row[6], "status": row[7], "last": row[8], "result": row[9], "created": row[10], "tags": self.item_tags("sched:%s" % row[0], row[3])})
|
||||
return items
|
||||
|
||||
def profile_path(self, name):
|
||||
return os.path.join(self.config.profiles_dir, name + ".sys.md")
|
||||
|
||||
def session_path(self, name):
|
||||
return os.path.join(self.config.profiles_dir, name + ".session.json")
|
||||
|
||||
def list_profiles(self):
|
||||
names = []
|
||||
for entry in sorted(os.listdir(self.config.profiles_dir)):
|
||||
if entry.endswith(".sys.md") and ".bot." not in entry:
|
||||
names.append(entry[:-7])
|
||||
return names
|
||||
|
||||
def unseal_file(self, raw):
|
||||
if raw.startswith(Seal.PREFIX):
|
||||
if not self.seal.enabled:
|
||||
raise SealError("profile is sealed, set TAI_PASSPHRASE")
|
||||
return self.seal.unlock(raw)
|
||||
return raw
|
||||
|
||||
def load_system(self, name):
|
||||
path = self.profile_path(name)
|
||||
if not os.path.exists(path):
|
||||
return None
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
return self.unseal_file(handle.read())
|
||||
except (OSError, ValueError):
|
||||
Release.quarantine(path)
|
||||
return None
|
||||
except SealError:
|
||||
if not self.seal.enabled:
|
||||
raise
|
||||
Release.quarantine(path)
|
||||
return None
|
||||
|
||||
def save_system(self, name, text):
|
||||
try:
|
||||
Release.sealed_write(self.profile_path(name), self.seal.lock(text))
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def load_session(self, name, bot="main"):
|
||||
path = self.session_path(name) if bot == "main" else self.bot_session_path(name, bot)
|
||||
if not os.path.exists(path):
|
||||
return []
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
raw = handle.read()
|
||||
data = json.loads(self.unseal_file(raw))
|
||||
return [item for item in data if isinstance(item, dict) and item.get("role") in ("user", "assistant", "tool")]
|
||||
except (OSError, ValueError):
|
||||
Release.quarantine(path)
|
||||
return []
|
||||
except SealError:
|
||||
if not self.seal.enabled:
|
||||
raise
|
||||
Release.quarantine(path)
|
||||
return []
|
||||
|
||||
def save_session(self, name, messages, bot="main"):
|
||||
path = self.session_path(name) if bot == "main" else self.bot_session_path(name, bot)
|
||||
try:
|
||||
Release.sealed_write(path, self.seal.lock(json.dumps([item for item in messages if item.get("role") != "system"])))
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def turn_state_path(self, profile, bot="main"):
|
||||
if bot == "main":
|
||||
return os.path.join(self.config.profiles_dir, "%s.turn.json" % profile)
|
||||
return os.path.join(self.config.profiles_dir, "%s.bot.%s.turn.json" % (profile, bot))
|
||||
|
||||
def load_turn_state(self, profile, bot="main"):
|
||||
path = self.turn_state_path(profile, bot)
|
||||
if not os.path.exists(path):
|
||||
return {}
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
data = json.loads(self.unseal_file(handle.read()))
|
||||
return data if isinstance(data, dict) else {}
|
||||
except (OSError, ValueError):
|
||||
Release.quarantine(path)
|
||||
return {}
|
||||
except SealError:
|
||||
if not self.seal.enabled:
|
||||
raise
|
||||
Release.quarantine(path)
|
||||
return {}
|
||||
|
||||
def save_turn_state(self, profile, bot, state):
|
||||
try:
|
||||
Release.sealed_write(self.turn_state_path(profile, bot), self.seal.lock(json.dumps(state)))
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def bot_system_path(self, profile, bot):
|
||||
return os.path.join(self.config.profiles_dir, "%s.bot.%s.sys.md" % (profile, bot))
|
||||
|
||||
def bot_session_path(self, profile, bot):
|
||||
return os.path.join(self.config.profiles_dir, "%s.bot.%s.session.json" % (profile, bot))
|
||||
|
||||
def bots_registry_path(self, profile):
|
||||
return os.path.join(self.config.profiles_dir, "%s.bots.json" % profile)
|
||||
|
||||
def load_bots(self, profile):
|
||||
path = self.bots_registry_path(profile)
|
||||
if not os.path.exists(path):
|
||||
return {}
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
data = json.loads(self.unseal_file(handle.read()))
|
||||
bots = data.get("bots") if isinstance(data, dict) else None
|
||||
return bots if isinstance(bots, dict) else {}
|
||||
except (OSError, ValueError):
|
||||
Release.quarantine(path)
|
||||
return {}
|
||||
except SealError:
|
||||
if not self.seal.enabled:
|
||||
raise
|
||||
Release.quarantine(path)
|
||||
return {}
|
||||
|
||||
def save_bots(self, profile, bots):
|
||||
try:
|
||||
Release.sealed_write(self.bots_registry_path(profile), self.seal.lock(json.dumps({"bots": bots})))
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def list_bots(self, profile):
|
||||
bots = self.load_bots(profile)
|
||||
found = [{"name": "main", "nicknames": []}]
|
||||
for name in sorted(bots):
|
||||
found.append({"name": name, "nicknames": (bots.get(name) or {}).get("nicknames", [])})
|
||||
return found
|
||||
|
||||
def resolve_bot(self, profile, mention):
|
||||
want = str(mention or "").strip().lower()
|
||||
if want in ("", "main"):
|
||||
return "main"
|
||||
bots = self.load_bots(profile)
|
||||
if want in bots:
|
||||
return want
|
||||
for name in sorted(bots):
|
||||
if want in (bots.get(name) or {}).get("nicknames", []):
|
||||
return name
|
||||
return None
|
||||
|
||||
def load_bot_system(self, profile, bot):
|
||||
if bot == "main":
|
||||
return self.load_system(profile)
|
||||
path = self.bot_system_path(profile, bot)
|
||||
if not os.path.exists(path):
|
||||
return None
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
return self.unseal_file(handle.read())
|
||||
except (OSError, ValueError):
|
||||
Release.quarantine(path)
|
||||
return None
|
||||
except SealError:
|
||||
if not self.seal.enabled:
|
||||
raise
|
||||
Release.quarantine(path)
|
||||
return None
|
||||
|
||||
def save_bot_system(self, profile, bot, text):
|
||||
try:
|
||||
Release.sealed_write(self.bot_system_path(profile, bot), self.seal.lock(text))
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def close(self):
|
||||
try:
|
||||
self.db.execute("PRAGMA wal_checkpoint(TRUNCATE)")
|
||||
except sqlite3.Error:
|
||||
pass
|
||||
try:
|
||||
self.db.close()
|
||||
except sqlite3.Error:
|
||||
pass
|
||||
if getattr(self, "memdb", None) is not None:
|
||||
try:
|
||||
self.memdb.close()
|
||||
except sqlite3.Error:
|
||||
pass
|
||||
self.memdb = None
|
||||
|
||||
@staticmethod
|
||||
def ensure_column(db, table, column, ddl):
|
||||
names = [row[1] for row in db.execute("PRAGMA table_info(%s)" % table).fetchall()]
|
||||
if column not in names:
|
||||
db.execute("ALTER TABLE %s ADD COLUMN %s %s" % (table, column, ddl))
|
||||
|
||||
@staticmethod
|
||||
def fts_escape(query):
|
||||
terms = re.findall(r"[a-z0-9]+", (query or "").lower())
|
||||
if not terms:
|
||||
return None
|
||||
return ['"%s"' % term for term in terms]
|
||||
|
||||
@staticmethod
|
||||
def fts_match(terms):
|
||||
return " AND ".join("{title body} : %s" % term for term in terms)
|
||||
|
||||
@staticmethod
|
||||
def open_mem_events():
|
||||
mem = sqlite3.connect(":memory:")
|
||||
try:
|
||||
mem.execute("CREATE VIRTUAL TABLE mem_events USING fts5(item, kind, sub, title, body, profile, tokenize='porter unicode61 remove_diacritics 2')")
|
||||
except sqlite3.Error:
|
||||
mem.execute("CREATE VIRTUAL TABLE mem_events USING fts5(item, kind, sub, title, body, profile)")
|
||||
return mem
|
||||
|
||||
@staticmethod
|
||||
def setup_fts(db):
|
||||
try:
|
||||
exists = db.execute("SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'fts_docs'").fetchone()
|
||||
if not exists:
|
||||
try:
|
||||
db.execute("CREATE VIRTUAL TABLE fts_docs USING fts5(item, kind, sub, title, body, profile, tokenize='porter unicode61 remove_diacritics 2')")
|
||||
except sqlite3.Error:
|
||||
db.execute("CREATE VIRTUAL TABLE fts_docs USING fts5(item, kind, sub, title, body, profile)")
|
||||
db.execute("""CREATE TRIGGER IF NOT EXISTS trg_records_ai AFTER INSERT ON records BEGIN
|
||||
INSERT INTO fts_docs(item, kind, sub, title, body, profile) VALUES(new.id, 'record', new.kind, new.title, new.content, new.profile);
|
||||
END""")
|
||||
db.execute("""CREATE TRIGGER IF NOT EXISTS trg_records_au AFTER UPDATE ON records WHEN old.title != new.title OR old.content != new.content OR old.profile != new.profile BEGIN
|
||||
DELETE FROM fts_docs WHERE item = old.id AND kind = 'record';
|
||||
INSERT INTO fts_docs(item, kind, sub, title, body, profile) VALUES(new.id, 'record', new.kind, new.title, new.content, new.profile);
|
||||
END""")
|
||||
db.execute("""CREATE TRIGGER IF NOT EXISTS trg_records_ad AFTER DELETE ON records BEGIN
|
||||
DELETE FROM fts_docs WHERE item = old.id AND kind = 'record';
|
||||
END""")
|
||||
db.execute("""CREATE TRIGGER IF NOT EXISTS trg_events_ai AFTER INSERT ON events BEGIN
|
||||
INSERT INTO fts_docs(item, kind, sub, title, body, profile) VALUES(new.id, 'event', new.role || '/' || new.kind, new.ts || ' ' || new.role || '/' || new.kind, new.text, new.profile);
|
||||
END""")
|
||||
db.execute("""CREATE TRIGGER IF NOT EXISTS trg_events_au AFTER UPDATE ON events WHEN old.text != new.text OR old.profile != new.profile BEGIN
|
||||
DELETE FROM fts_docs WHERE item = old.id AND kind = 'event';
|
||||
INSERT INTO fts_docs(item, kind, sub, title, body, profile) VALUES(new.id, 'event', new.role || '/' || new.kind, new.ts || ' ' || new.role || '/' || new.kind, new.text, new.profile);
|
||||
END""")
|
||||
db.execute("""CREATE TRIGGER IF NOT EXISTS trg_audit_ai AFTER INSERT ON audit BEGIN
|
||||
INSERT INTO fts_docs(item, kind, sub, title, body, profile) VALUES('audit:' || new.id, 'audit', new.action, new.path, new.message, new.profile);
|
||||
END""")
|
||||
db.execute("""CREATE TRIGGER IF NOT EXISTS trg_chunks_ai AFTER INSERT ON chunks BEGIN
|
||||
INSERT INTO fts_docs(item, kind, sub, title, body, profile) VALUES('chunk:' || new.id, 'chunk', new.parent_kind, new.parent_id, new.text, new.profile);
|
||||
END""")
|
||||
db.execute("""CREATE TRIGGER IF NOT EXISTS trg_chunks_ad AFTER DELETE ON chunks BEGIN
|
||||
DELETE FROM fts_docs WHERE item = 'chunk:' || old.id AND kind = 'chunk';
|
||||
END""")
|
||||
if db.execute("SELECT COUNT(*) FROM fts_docs").fetchone()[0] == 0:
|
||||
db.execute("INSERT INTO fts_docs(item, kind, sub, title, body, profile) SELECT id, 'record', kind, title, content, profile FROM records")
|
||||
db.execute("INSERT INTO fts_docs(item, kind, sub, title, body, profile) SELECT id, 'event', role || '/' || kind, ts || ' ' || role || '/' || kind, text, profile FROM events")
|
||||
db.execute("INSERT INTO fts_docs(item, kind, sub, title, body, profile) SELECT 'audit:' || id, 'audit', action, path, message, profile FROM audit")
|
||||
db.commit()
|
||||
return True
|
||||
except sqlite3.Error:
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def table_pk_columns(db, table):
|
||||
for row in db.execute("PRAGMA index_list(%s)" % table).fetchall():
|
||||
if row[3] == "pk":
|
||||
return [info[2] for info in db.execute("PRAGMA index_info(%s)" % row[1]).fetchall()]
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def rebuild_with_profile(db, table, columns, pk):
|
||||
if Store.table_pk_columns(db, table) == list(pk):
|
||||
return False
|
||||
names = [column.split()[0] for column in columns]
|
||||
have = {row[1] for row in db.execute("PRAGMA table_info(%s)" % table).fetchall()}
|
||||
keep = [name for name in names if name in have]
|
||||
db.execute("CREATE TABLE %s_new (%s, PRIMARY KEY (%s))" % (table, ", ".join(columns), ", ".join(pk)))
|
||||
if keep:
|
||||
db.execute("INSERT OR IGNORE INTO %s_new (%s) SELECT %s FROM %s" % (table, ", ".join(keep), ", ".join(keep), table))
|
||||
db.execute("DROP TABLE %s" % table)
|
||||
db.execute("ALTER TABLE %s_new RENAME TO %s" % (table, table))
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def migrate_profiles(db):
|
||||
for table in ("secrets", "records", "tags", "edges", "audit", "events", "schedules"):
|
||||
Store.ensure_column(db, table, "profile", "TEXT")
|
||||
db.execute("UPDATE %s SET profile = 'default' WHERE profile IS NULL" % table)
|
||||
Store.rebuild_with_profile(db, "secrets", ("profile TEXT", "name TEXT", "value TEXT", "meta TEXT", "updated TEXT"), ("profile", "name"))
|
||||
Store.rebuild_with_profile(db, "tags", ("item TEXT", "tag TEXT", "profile TEXT"), ("item", "tag", "profile"))
|
||||
Store.rebuild_with_profile(db, "edges", ("src TEXT", "dst TEXT", "relation TEXT", "profile TEXT", "created TEXT"), ("src", "dst", "relation", "profile"))
|
||||
|
||||
@staticmethod
|
||||
def ensure_schema(db):
|
||||
notes = []
|
||||
for table, columns, pk in Consts.SCHEMA:
|
||||
exists = db.execute("SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?", (table,)).fetchone()
|
||||
if not exists:
|
||||
extra = ", PRIMARY KEY (%s)" % ", ".join(pk) if pk else ""
|
||||
db.execute("CREATE TABLE %s (%s%s)" % (table, ", ".join(columns), extra))
|
||||
continue
|
||||
have = {row[1] for row in db.execute("PRAGMA table_info(%s)" % table).fetchall()}
|
||||
for column in columns:
|
||||
parts = column.split()
|
||||
if parts[0] not in have:
|
||||
db.execute("ALTER TABLE %s ADD COLUMN %s" % (table, " ".join(parts[:2])))
|
||||
notes.append("added %s.%s" % (table, parts[0]))
|
||||
for index, target in Consts.SCHEMA_INDEXES:
|
||||
db.execute("CREATE INDEX IF NOT EXISTS %s ON %s" % (index, target))
|
||||
db.commit()
|
||||
return notes
|
||||
@@ -0,0 +1,135 @@
|
||||
"""Host inspection probes and report assembly."""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.sandbox import Sandbox
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["SysInfo", "SYSINFO_CHECKS"]
|
||||
|
||||
|
||||
class SysInfo:
|
||||
"""SysInfo."""
|
||||
|
||||
@staticmethod
|
||||
def tmux_pane_header():
|
||||
try:
|
||||
info = subprocess.run(["tmux", "display-message", "-p", "#{session_name}:#{window_index}.#{pane_index} #{pane_current_command}"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=10)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
return ""
|
||||
if info.returncode == 0 and info.stdout.strip():
|
||||
return info.stdout.strip()
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def sysinfo_check_os():
|
||||
try:
|
||||
uname = os.uname()
|
||||
base = "%s %s %s" % (uname.sysname, uname.release, uname.machine)
|
||||
except (AttributeError, OSError):
|
||||
base = sys.platform
|
||||
try:
|
||||
with open("/etc/os-release", "r", encoding="utf-8") as handle:
|
||||
fields = dict(line.strip().split("=", 1) for line in handle if "=" in line)
|
||||
pretty = fields.get("PRETTY_NAME", "").strip('"')
|
||||
if pretty:
|
||||
return "%s (%s)" % (pretty, base)
|
||||
except OSError:
|
||||
pass
|
||||
return base
|
||||
|
||||
@staticmethod
|
||||
def sysinfo_check_python():
|
||||
return "%s (%s)" % (sys.version.split()[0], sys.executable)
|
||||
|
||||
@staticmethod
|
||||
def sysinfo_check_venv():
|
||||
env_path = os.environ.get("VIRTUAL_ENV") or ""
|
||||
if env_path:
|
||||
return "yes (%s)" % env_path
|
||||
if sys.prefix != sys.base_prefix:
|
||||
return "yes (%s)" % sys.prefix
|
||||
return "no (system python at %s)" % sys.prefix
|
||||
|
||||
@staticmethod
|
||||
def sysinfo_check_root():
|
||||
geteuid = getattr(os, "geteuid", None)
|
||||
if geteuid is None:
|
||||
return "unknown (no geteuid on %s)" % sys.platform
|
||||
uid = geteuid()
|
||||
return "yes (uid 0)" if uid == 0 else "no (uid %d)" % uid
|
||||
|
||||
@staticmethod
|
||||
def sysinfo_check_container():
|
||||
engine = Sandbox.container_engine()
|
||||
if not engine:
|
||||
return "no engine (install podman or docker)"
|
||||
return "%s, box %s" % (engine, Sandbox.box_state(engine))
|
||||
|
||||
@staticmethod
|
||||
def sysinfo_check_binaries():
|
||||
found = []
|
||||
lost = []
|
||||
for name in Consts.SYSINFO_BINARIES:
|
||||
(found if shutil.which(name) else lost).append(name)
|
||||
return "present: %s; missing: %s" % (", ".join(found) or "none", ", ".join(lost) or "none")
|
||||
|
||||
@staticmethod
|
||||
def sysinfo_check_cpu():
|
||||
count = os.cpu_count() or 0
|
||||
try:
|
||||
one, five, fifteen = os.getloadavg()
|
||||
return "%d cores, load %.2f/%.2f/%.2f" % (count, one, five, fifteen)
|
||||
except OSError:
|
||||
return "%d cores" % count
|
||||
|
||||
@staticmethod
|
||||
def sysinfo_check_disk():
|
||||
usage = shutil.disk_usage(os.getcwd())
|
||||
return "%.1f GiB free of %.1f GiB at %s" % (usage.free / 2 ** 30, usage.total / 2 ** 30, os.getcwd())
|
||||
|
||||
@staticmethod
|
||||
def collect_sysinfo(names=None):
|
||||
wanted = [pair for pair in SYSINFO_CHECKS if names is None or pair[0] in names]
|
||||
results = {}
|
||||
|
||||
def target(name, func):
|
||||
started = time.time()
|
||||
try:
|
||||
value = func()
|
||||
except Exception as exc:
|
||||
value = "error: " + Text.short_error(exc)
|
||||
results[name] = (value, int((time.time() - started) * 1000))
|
||||
|
||||
workers = [threading.Thread(target=target, args=(name, func), daemon=True) for name, func in wanted]
|
||||
for worker in workers:
|
||||
worker.start()
|
||||
deadline = time.time() + Consts.SYSINFO_TIMEOUT
|
||||
for worker in workers:
|
||||
worker.join(timeout=max(0, deadline - time.time()))
|
||||
lines = []
|
||||
for name, _func in wanted:
|
||||
if name in results:
|
||||
value, elapsed = results[name]
|
||||
lines.append("%s: %s (%dms)" % (name, value, elapsed))
|
||||
else:
|
||||
lines.append("%s: timed out after %ds" % (name, Consts.SYSINFO_TIMEOUT))
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
SYSINFO_CHECKS = (
|
||||
("os", SysInfo.sysinfo_check_os),
|
||||
("python", SysInfo.sysinfo_check_python),
|
||||
("venv", SysInfo.sysinfo_check_venv),
|
||||
("root", SysInfo.sysinfo_check_root),
|
||||
("container", SysInfo.sysinfo_check_container),
|
||||
("binaries", SysInfo.sysinfo_check_binaries),
|
||||
("cpu", SysInfo.sysinfo_check_cpu),
|
||||
("disk", SysInfo.sysinfo_check_disk),
|
||||
)
|
||||
@@ -0,0 +1,123 @@
|
||||
"""Telegram bot transport and update handling."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import urllib.request
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.errors import BackendError
|
||||
from tai_pkg.sandbox import Sandbox
|
||||
from tai_pkg.scheduler import Scheduler
|
||||
from tai_pkg.text import Text
|
||||
|
||||
__all__ = ["Telegram", "TelegramBot"]
|
||||
|
||||
|
||||
class Telegram:
|
||||
def __init__(self, token):
|
||||
self.token = token
|
||||
|
||||
def call(self, method, params, timeout=70):
|
||||
body = json.dumps(params).encode("utf-8")
|
||||
request = urllib.request.Request(Consts.TELEGRAM_API + self.token + "/" + method, data=body, headers={"Content-Type": "application/json"}, method="POST")
|
||||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||||
return json.loads(response.read().decode("utf-8", "replace"))
|
||||
|
||||
|
||||
class TelegramBot:
|
||||
"""TelegramBot."""
|
||||
|
||||
@staticmethod
|
||||
def telegram_env_path(home):
|
||||
return os.path.join(home, "telegram.env")
|
||||
|
||||
@staticmethod
|
||||
def load_telegram_token(home):
|
||||
token = os.environ.get("TELEGRAM_BOT_TOKEN") or ""
|
||||
if token:
|
||||
return token
|
||||
path = TelegramBot.telegram_env_path(home)
|
||||
if not os.path.exists(path):
|
||||
return ""
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
for line in handle.read().splitlines():
|
||||
if line.startswith("TELEGRAM_BOT_TOKEN="):
|
||||
return line.split("=", 1)[1].strip().strip("'\"")
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def telegram_send(bot, chat_id, text):
|
||||
for pos in range(0, len(text), 4000):
|
||||
bot.call("sendMessage", {"chat_id": chat_id, "text": text[pos:pos + 4000] or "(empty reply)"})
|
||||
|
||||
@staticmethod
|
||||
def telegram_transcribe(bot, file_id):
|
||||
try:
|
||||
info = bot.call("getFile", {"file_id": file_id})
|
||||
remote = (info.get("result") or {}).get("file_path", "")
|
||||
if not remote:
|
||||
return "error: telegram returned no file"
|
||||
with urllib.request.urlopen(Consts.TELEGRAM_FILE_API + bot.token + "/" + remote, timeout=60) as response:
|
||||
audio = response.read(20000000)
|
||||
engine = Sandbox.ensure_box()
|
||||
return Sandbox.box_transcribe(engine, audio)
|
||||
except (OSError, ValueError, BackendError) as exc:
|
||||
return "error: transcription failed: " + Text.short_error(exc)
|
||||
|
||||
@staticmethod
|
||||
def handle_telegram_update(agent, bot, update):
|
||||
message = update.get("message") or {}
|
||||
chat_id = (message.get("chat") or {}).get("id")
|
||||
if not chat_id:
|
||||
return
|
||||
text = message.get("text") or ""
|
||||
if text == "/start":
|
||||
TelegramBot.telegram_send(bot, chat_id, "tai online. Send any message, /new for a fresh start.")
|
||||
return
|
||||
if text == "/new":
|
||||
agent.reset_history()
|
||||
TelegramBot.telegram_send(bot, chat_id, "fresh start.")
|
||||
return
|
||||
voice = message.get("voice") or message.get("audio")
|
||||
if voice:
|
||||
bot.call("sendChatAction", {"chat_id": chat_id, "action": "typing"})
|
||||
text = TelegramBot.telegram_transcribe(bot, voice.get("file_id", ""))
|
||||
if text.startswith("error"):
|
||||
TelegramBot.telegram_send(bot, chat_id, text)
|
||||
return
|
||||
if not text.strip():
|
||||
return
|
||||
bot.call("sendChatAction", {"chat_id": chat_id, "action": "typing"})
|
||||
reply = agent.run_turn(text, capture=True)
|
||||
TelegramBot.telegram_send(bot, chat_id, reply)
|
||||
|
||||
@staticmethod
|
||||
def run_telegram_bot(agent):
|
||||
token = TelegramBot.load_telegram_token(agent.config.home)
|
||||
if not token:
|
||||
print("missing TELEGRAM_BOT_TOKEN", file=sys.stderr)
|
||||
return 2
|
||||
bot = Telegram(token)
|
||||
try:
|
||||
me = bot.call("getMe", {})
|
||||
except (OSError, ValueError) as exc:
|
||||
print("telegram auth failed: %s" % Text.short_error(exc), file=sys.stderr)
|
||||
return 2
|
||||
print("telegram bot online as @%s" % (me.get("result") or {}).get("username", "?"))
|
||||
Scheduler.start_scheduler(agent.config, agent.store.seal)
|
||||
offset = 0
|
||||
while True:
|
||||
try:
|
||||
data = bot.call("getUpdates", {"offset": offset, "timeout": 50})
|
||||
except (OSError, ValueError) as exc:
|
||||
print("poll error: %s, retrying" % Text.short_error(exc))
|
||||
time.sleep(5)
|
||||
continue
|
||||
for update in data.get("result") or []:
|
||||
offset = update.get("update_id", offset) + 1
|
||||
try:
|
||||
TelegramBot.handle_telegram_update(agent, bot, update)
|
||||
except Exception as exc:
|
||||
print("update error: %s" % Text.short_error(exc))
|
||||
+252
@@ -0,0 +1,252 @@
|
||||
"""Pure string, time, token, record, and tag helpers."""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
|
||||
__all__ = ["Text"]
|
||||
|
||||
|
||||
class Text:
|
||||
"""Text."""
|
||||
|
||||
@staticmethod
|
||||
# Python 3.6 compatibility notes:
|
||||
# - subprocess 'text' and 'capture_output' kwargs were added in 3.7, so this
|
||||
# module only uses 'universal_newlines' plus explicit PIPEs (works on 3.6+).
|
||||
# - datetime.fromisoformat was added in 3.7, so all parsing goes through
|
||||
# _fromisoformat below (native when available, strptime fallback on 3.6).
|
||||
|
||||
|
||||
def _fromisoformat(text):
|
||||
"""datetime.fromisoformat backport for Python 3.6 (handles Z and +HH:MM)."""
|
||||
native = getattr(datetime, "fromisoformat", None)
|
||||
if native is not None:
|
||||
raw = str(text or "").strip()
|
||||
if raw[-1:] in ("Z", "z"):
|
||||
raw = raw[:-1] + "+00:00"
|
||||
return native(raw)
|
||||
raw = str(text or "").strip()
|
||||
if not raw:
|
||||
raise ValueError("empty datetime")
|
||||
if raw[-1:] in ("Z", "z"):
|
||||
raw = raw[:-1] + "+00:00"
|
||||
tzinfo = None
|
||||
core = raw
|
||||
match = re.match(r"^(\d{4}-\d{2}-\d{2}(?:[T ]\d{2}:\d{2}(?::\d{2}(?:\.\d{1,6})?)?)?)([+-]\d{2}:?\d{2}(?::\d{2})?|[+-]\d{4}|[+-]\d{2})$", raw)
|
||||
if match is not None and match.group(2):
|
||||
core, tzpart = match.group(1), match.group(2)
|
||||
sign = 1 if tzpart[0] == "+" else -1
|
||||
digits = tzpart[1:].replace(":", "")
|
||||
hours = int(digits[:2])
|
||||
minutes = int(digits[2:4]) if len(digits) >= 4 else 0
|
||||
tzinfo = timezone(sign * timedelta(hours=hours, minutes=minutes))
|
||||
for fmt in ("%Y-%m-%dT%H:%M:%S.%f", "%Y-%m-%dT%H:%M:%S", "%Y-%m-%dT%H:%M",
|
||||
"%Y-%m-%d %H:%M:%S.%f", "%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M",
|
||||
"%Y-%m-%d"):
|
||||
try:
|
||||
parsed = datetime.strptime(core, fmt)
|
||||
except ValueError:
|
||||
continue
|
||||
return parsed.replace(tzinfo=tzinfo)
|
||||
raise ValueError("invalid datetime, use ISO like 2026-10-08T09:00")
|
||||
|
||||
@staticmethod
|
||||
def curl_hint(command):
|
||||
if re.search(r"\bcurl\b|\bwget\b", command or "") and "http" in (command or "").lower():
|
||||
return "\n" + Consts.CURL_HINT
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def now_iso():
|
||||
return datetime.now(timezone.utc).isoformat(timespec="seconds")
|
||||
|
||||
@staticmethod
|
||||
def short_error(exc):
|
||||
text = str(exc) or type(exc).__name__
|
||||
return text[:300]
|
||||
|
||||
@staticmethod
|
||||
def truncate(text, head=1500, tail=500):
|
||||
if len(text) <= head + tail + 50:
|
||||
return text
|
||||
skipped = len(text) - head - tail
|
||||
return text[:head] + "\n...[%d chars truncated]...\n" % skipped + text[-tail:]
|
||||
|
||||
@staticmethod
|
||||
def capped_text(text):
|
||||
if text is None:
|
||||
return None, None
|
||||
if len(text) <= Consts.AUDIT_MAX_CHARS:
|
||||
return text, len(text.encode("utf-8"))
|
||||
return text[:Consts.AUDIT_MAX_CHARS] + "\n...[truncated, full size %d bytes]..." % len(text.encode("utf-8")), len(text.encode("utf-8"))
|
||||
|
||||
@staticmethod
|
||||
def trunc_marker(true_bytes):
|
||||
return "\n...[truncated, full size %d bytes]..." % true_bytes
|
||||
|
||||
@staticmethod
|
||||
def read_capped(path, cap=Consts.AUDIT_MAX_CHARS):
|
||||
try:
|
||||
size = os.path.getsize(path)
|
||||
except OSError:
|
||||
return None, None
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8", errors="replace") as handle:
|
||||
text = handle.read(cap + 1)
|
||||
except OSError:
|
||||
return None, None
|
||||
if len(text) > cap:
|
||||
return text[:cap], size
|
||||
return text, None
|
||||
|
||||
@staticmethod
|
||||
def chunk_lines(text, max_lines=Consts.CHUNK_MAX_LINES, overlap_lines=Consts.CHUNK_OVERLAP_LINES, max_chars=Consts.CHUNK_MAX_CHARS):
|
||||
"""Generic RAG chunker: bounded windows of whole lines with small overlap, never splitting mid-line.
|
||||
|
||||
Used for any content that gets indexed for retrieval (terminal captures, long records, ...),
|
||||
so there is one chunking rule in the codebase instead of one per content type.
|
||||
"""
|
||||
lines = text.splitlines()
|
||||
n = len(lines)
|
||||
if n == 0:
|
||||
return []
|
||||
chunks = []
|
||||
start = 0
|
||||
while start < n:
|
||||
end = start
|
||||
size = 0
|
||||
count = 0
|
||||
while end < n and count < max_lines and size < max_chars:
|
||||
size += len(lines[end]) + 1
|
||||
count += 1
|
||||
end += 1
|
||||
piece = "\n".join(lines[start:end]).strip()
|
||||
if piece:
|
||||
chunks.append(piece)
|
||||
if end >= n:
|
||||
break
|
||||
start = max(start + 1, end - overlap_lines)
|
||||
return chunks
|
||||
|
||||
@staticmethod
|
||||
def collapse_repeated_lines(text, threshold=3):
|
||||
"""Collapses runs of 3+ identical lines (progress bars, retry spam) to one line + a count note."""
|
||||
lines = text.splitlines()
|
||||
out = []
|
||||
i = 0
|
||||
n = len(lines)
|
||||
while i < n:
|
||||
j = i
|
||||
while j < n and lines[j] == lines[i]:
|
||||
j += 1
|
||||
run = j - i
|
||||
if run >= threshold:
|
||||
out.append(lines[i])
|
||||
out.append("... (repeated %d more times) ..." % (run - 1))
|
||||
else:
|
||||
out.extend(lines[i:j])
|
||||
i = j
|
||||
return "\n".join(out)
|
||||
|
||||
@staticmethod
|
||||
def estimate_tokens(text):
|
||||
return max(1, len(text or "") // 4)
|
||||
|
||||
@staticmethod
|
||||
def as_int(value, default=0):
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
@staticmethod
|
||||
def tool_args_text(raw_arguments):
|
||||
if not raw_arguments:
|
||||
return "{}"
|
||||
if isinstance(raw_arguments, str):
|
||||
return raw_arguments
|
||||
try:
|
||||
return json.dumps(raw_arguments, default=str)
|
||||
except (TypeError, ValueError):
|
||||
return str(raw_arguments)
|
||||
|
||||
@staticmethod
|
||||
def action_signature(name, raw_arguments):
|
||||
try:
|
||||
flat = json.dumps(json.loads(Text.tool_args_text(raw_arguments)), sort_keys=True, default=str)
|
||||
except ValueError:
|
||||
flat = str(raw_arguments or "")
|
||||
return "%s %s" % (name, flat[:500])
|
||||
|
||||
@staticmethod
|
||||
def result_digest(result):
|
||||
return hashlib.sha1((result or "").encode("utf-8", "replace")).hexdigest()
|
||||
|
||||
@staticmethod
|
||||
def human_size(num):
|
||||
num = max(0, int(num))
|
||||
if num < 1024:
|
||||
return "%d B" % num
|
||||
if num < 1048576:
|
||||
return "%.1f KB" % (num / 1024)
|
||||
return "%.1f MB" % (num / 1048576)
|
||||
|
||||
@staticmethod
|
||||
def is_free_model(model):
|
||||
model = (model or "").lower()
|
||||
return model.endswith(":free") or model.endswith("-free")
|
||||
|
||||
@staticmethod
|
||||
def secret_env_name(name):
|
||||
return "TAI_SECRET_" + re.sub(r"[^A-Z0-9_]", "_", name.upper())
|
||||
|
||||
@staticmethod
|
||||
def normalize_tag(raw):
|
||||
return re.sub(r"[^a-z0-9]+", "-", str(raw or "").strip().lower()).strip("-")[:32]
|
||||
|
||||
@staticmethod
|
||||
def normalize_tags(items):
|
||||
seen = []
|
||||
for item in items:
|
||||
cleaned = Text.normalize_tag(item)
|
||||
if cleaned and cleaned not in seen:
|
||||
seen.append(cleaned)
|
||||
return seen[:20]
|
||||
|
||||
@staticmethod
|
||||
def singular_noun(word):
|
||||
if word in Consts.SINGULAR_KEEP or len(word) <= 3:
|
||||
return word
|
||||
if word.endswith("ies"):
|
||||
return word[:-3] + "y"
|
||||
if word.endswith(("ses", "xes", "zes", "ches", "shes")):
|
||||
return word[:-2]
|
||||
if word.endswith("s") and not word.endswith(("ss", "us", "is", "os")):
|
||||
return word[:-1]
|
||||
return word
|
||||
|
||||
@staticmethod
|
||||
def tag_variants(tag):
|
||||
found = {tag, Text.singular_noun(tag)}
|
||||
if not tag.endswith("s"):
|
||||
found.add(tag + "s")
|
||||
return sorted(found)
|
||||
|
||||
@staticmethod
|
||||
def parse_node_id(raw):
|
||||
match = re.fullmatch(r"(mem|secret|sched):(.+)", str(raw or "").strip())
|
||||
if not match:
|
||||
return None
|
||||
kind, key = match.group(1), match.group(2)
|
||||
if kind == "mem" and re.fullmatch(r"[0-9a-f]{16}", key):
|
||||
return kind, key
|
||||
if kind == "secret" and Consts.SECRET_RE.match(key):
|
||||
return kind, key
|
||||
if kind == "sched" and key.isdigit():
|
||||
return kind, key
|
||||
return None
|
||||
+1432
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,116 @@
|
||||
"""Agent tool registry, schemas, and lazy selection."""
|
||||
|
||||
import os
|
||||
import re
|
||||
|
||||
from tai_pkg.consts import Consts
|
||||
from tai_pkg.plexer import PlexerControllerFactory
|
||||
|
||||
__all__ = ["ToolRegistry", "TOOL_SCHEMAS", "TOOL_ENV_GATES"]
|
||||
|
||||
|
||||
class ToolRegistry:
|
||||
"""ToolRegistry."""
|
||||
|
||||
@staticmethod
|
||||
def tool_schema(name, description, properties, required):
|
||||
return {"type": "function", "function": {"name": name, "description": description, "parameters": {"type": "object", "properties": properties, "required": required}}}
|
||||
|
||||
@staticmethod
|
||||
def tool_env_ok(name):
|
||||
gate = TOOL_ENV_GATES.get(name)
|
||||
return gate is None or gate()
|
||||
|
||||
@staticmethod
|
||||
def tool_catalog():
|
||||
lines = ["", "", "## Tool catalog (core loads always; name a lazy tool or topic to load it)"]
|
||||
for schema in TOOL_SCHEMAS:
|
||||
name = schema["function"]["name"]
|
||||
if not ToolRegistry.tool_env_ok(name):
|
||||
continue
|
||||
if name in Consts.CORE_TOOLS:
|
||||
continue
|
||||
first = schema["function"]["description"].split(".")[0][:80]
|
||||
lines.append("- %s: %s" % (name, first))
|
||||
return "\n".join(lines)
|
||||
|
||||
@staticmethod
|
||||
def select_tools(text):
|
||||
words = set(re.findall(r"[a-z0-9]+", (text or "").lower()))
|
||||
picked = []
|
||||
for schema in TOOL_SCHEMAS:
|
||||
name = schema["function"]["name"]
|
||||
if not ToolRegistry.tool_env_ok(name):
|
||||
continue
|
||||
if name in Consts.CORE_TOOLS:
|
||||
picked.append(schema)
|
||||
continue
|
||||
tokens = [item for item in name.split("_") if len(item) >= 4 and item not in Consts.LAZY_NAME_STOP]
|
||||
if tokens and set(tokens) & words:
|
||||
picked.append(schema)
|
||||
continue
|
||||
if set(Consts.TOOL_TAGS.get(name, ())) & words:
|
||||
picked.append(schema)
|
||||
continue
|
||||
desc = set(word for word in re.findall(r"[a-z0-9]+", schema["function"]["description"].lower()) if len(word) >= 6)
|
||||
if desc & words:
|
||||
picked.append(schema)
|
||||
return picked
|
||||
|
||||
|
||||
TOOL_SCHEMAS = [
|
||||
ToolRegistry.tool_schema("shell", "Run a shell command. Output is truncated. On home, destructive commands ask the user first. On sandbox, commands run isolated in the container and need no approval. Pass secrets to expose named vault secrets as TAI_SECRET_<NAME> variables for that command only; using secrets asks the user first. For local commands only: use web_fetch for HTTP(S), never curl or wget.", {"command": {"type": "string"}, "workdir": {"type": "string"}, "secrets": {"type": "array"}}, ["command"]),
|
||||
ToolRegistry.tool_schema("read_file", "Read a text file. Large files are truncated.", {"path": {"type": "string"}}, ["path"]),
|
||||
ToolRegistry.tool_schema("write_file", "Write content to a file, creating parent directories. Overwrites existing files.", {"path": {"type": "string"}, "content": {"type": "string"}}, ["path", "content"]),
|
||||
ToolRegistry.tool_schema("edit_file", "Replace one exact text match in a file. Fails unless the match is unique.", {"path": {"type": "string"}, "find": {"type": "string"}, "replace": {"type": "string"}}, ["path", "find", "replace"]),
|
||||
ToolRegistry.tool_schema("web_search", "Search the web. Set images for image search, content to include fetched page text.", {"query": {"type": "string"}, "images": {"type": "boolean"}, "content": {"type": "boolean"}}, ["query"]),
|
||||
ToolRegistry.tool_schema("web_fetch", "HTTP client: fetch a URL with any method and return status, content type, size, and content. GET by default; pass method (POST, PUT, PATCH, DELETE, HEAD, OPTIONS), headers, and a string body for APIs. HTML becomes readable text, JSON stays raw, raw forces the raw body. Use this for all HTTP(S), never curl or wget via shell. Pass auth_secret to send a vault secret as an Authorization header; the value stays sealed and secret use asks the user first.", {"url": {"type": "string"}, "method": {"type": "string"}, "headers": {"type": "object"}, "body": {"type": "string"}, "timeout": {"type": "integer"}, "raw": {"type": "boolean"}, "auth_secret": {"type": "string"}, "auth_header": {"type": "string"}, "auth_prefix": {"type": "string"}}, ["url"]),
|
||||
ToolRegistry.tool_schema("speak", "Synthesize text to speech with a free voice, save the MP3, and play it when a player exists.", {"text": {"type": "string"}, "voice": {"type": "string"}}, ["text"]),
|
||||
ToolRegistry.tool_schema("listen", "Record from the microphone for some seconds and transcribe it. Needs OS audio tools.", {"seconds": {"type": "integer"}}, []),
|
||||
ToolRegistry.tool_schema("remember", "Update your own system message with new knowledge or behavior. The instruction is merged into the existing system message, unrelated parts stay intact. Use it to forget as well by instructing what to remove. Execute by default when behavior, preferences, or durable facts need updating. Never pass passwords, tokens, or secret values here; store those with store_secret instead.", {"instruction": {"type": "string"}}, ["instruction"]),
|
||||
ToolRegistry.tool_schema("recall", "Search past session memory of the current profile by keyword, optionally filtered by tags.", {"query": {"type": "string"}, "tags": {"type": "array"}}, ["query"]),
|
||||
ToolRegistry.tool_schema("load_skill", "Load a skill by name. Returns full instructions plus bundled file paths. Missing skills with a blueprint are researched and built on demand.", {"name": {"type": "string"}}, ["name"]),
|
||||
ToolRegistry.tool_schema("get_current_terminal_content", "Capture visible text of the current tmux pane including scrollback. Works inside tmux or against a running tmux server.", {"lines": {"type": "integer"}}, []),
|
||||
ToolRegistry.tool_schema("export_terminal_log", "Only available inside a live tmux session. Exports the full pane scrollback (tmux capture-pane -pS -, the entire history, not just what's visible) to a log file under tai's home dir, then chunks it and indexes the chunks for retrieval with SQLite FTS5 + bm25 (lexical only, no embeddings). From then on, use the search tool (kind chunk) to answer questions about anything that happened in this terminal; it returns the actual chunk text, not just a snippet.", {"tags": {"type": "array"}}, []),
|
||||
ToolRegistry.tool_schema("plexer_send", "Only available when a tai-managed terminal session is active (either tai's own native wrapper, or an active tmux pane). Sends literal keystrokes into that session's shell, as if typed. Set enter to also press Enter and actually run it. No output is returned by this call; follow up with plexer_read.", {"text": {"type": "string"}, "enter": {"type": "boolean"}}, ["text"]),
|
||||
ToolRegistry.tool_schema("plexer_read", "Only available when a tai-managed terminal session is active. Reads back the last N lines the session's shell has produced (its actual rendered output, not a tool-call summary).", {"lines": {"type": "integer"}}, []),
|
||||
ToolRegistry.tool_schema("plexer_type", "Only available when a tai-managed terminal session is active. Types text into the session one character at a time with a delay between keystrokes, for a demo/animated-typing effect, then presses Enter unless enter is set to false.", {"text": {"type": "string"}, "delay": {"type": "number"}, "enter": {"type": "boolean"}}, ["text"]),
|
||||
ToolRegistry.tool_schema("plexer_status", "Only available when a tai-managed terminal session is active. Reports which backend is managing the current session (tai's own native pty wrapper, or tmux) and backend-specific details (pid/socket/log, or pane id).", {}, []),
|
||||
ToolRegistry.tool_schema("fork", "Spawn a background subagent with its own context that works while you continue. Returns an agent id immediately. Collect its summarized result with poll. Subagents get a smaller step budget and a time limit.", {"task": {"type": "string"}, "timeout": {"type": "integer"}, "profile": {"type": "string"}}, ["task"]),
|
||||
ToolRegistry.tool_schema("poll", "Collect a background subagent result by id. Waits up to wait seconds, then reports running or the result.", {"id": {"type": "integer"}, "wait": {"type": "integer"}}, ["id"]),
|
||||
ToolRegistry.tool_schema("sysinfo", "Inspect the host machine where tai runs. Runs os, python, venv, root, container, binaries, cpu, and disk checks in parallel and reports each with timing. Pass checks to run a named subset.", {"checks": {"type": "array"}}, []),
|
||||
ToolRegistry.tool_schema("create_skill", "Create a new agent skill by name from a brief. Runs a dedicated deep-research worker on this machine (it calls sysinfo first, then researches with web search and shell until it has verified information) and writes SKILL.md plus supporting files into the skill directory. Runs synchronously and can take many minutes. Scope project writes under ./.tai/skills (the default), scope home under ~/.tai/skills.", {"name": {"type": "string"}, "brief": {"type": "string"}, "scope": {"type": "string"}, "timeout": {"type": "integer"}}, ["name", "brief"]),
|
||||
ToolRegistry.tool_schema("install_skill", "Install a skill that already exists somewhere else, instead of building one from scratch. source can be a local directory containing SKILL.md (or one subdirectory that does), a bare SKILL.md file, a local .zip, or a git URL (needs git installed; most sources need nothing beyond the standard library). Validates the SKILL.md frontmatter, then lands it under tai's own skill directories: copy vendors a private snapshot, symlink points at the source in place (the default when the source already sits in a recognized skills folder like .claude/skills or .agents/skills, so there is one canonical copy instead of a drifting duplicate). Scope defaults like create_skill: project when the cwd looks like a project, else home. Overwriting an existing skill of the same name asks the user first unless force is set.", {"source": {"type": "string"}, "name": {"type": "string"}, "scope": {"type": "string"}, "mode": {"type": "string"}, "force": {"type": "boolean"}}, ["source"]),
|
||||
ToolRegistry.tool_schema("store_secret", "Store a password, token, or secret in the sealed vault under a name. Values are encrypted at rest, never shown back, and only usable by reference; prefer /secret set in the REPL so the value never enters the conversation. Optional username, host, port, notes, expires (ISO datetime), and tags describe it; name and value suffice, never pester the user for more.", {"name": {"type": "string"}, "value": {"type": "string"}, "username": {"type": "string"}, "host": {"type": "string"}, "port": {"type": "integer"}, "notes": {"type": "string"}, "expires": {"type": "string"}, "tags": {"type": "array"}}, ["name", "value"]),
|
||||
ToolRegistry.tool_schema("list_secrets", "List vault secret names with metadata and tags. Values are never revealed.", {}, []),
|
||||
ToolRegistry.tool_schema("delete_secret", "Delete a vault secret by name. Asks the user first.", {"name": {"type": "string"}}, ["name"]),
|
||||
ToolRegistry.tool_schema("schedule", "Schedule a prompt to run at a future time or on a repeating interval. Exactly one of at (ISO datetime, naive means local time) or every (interval seconds, minimum 60) is required. When due, a background subagent runs the prompt without blocking anyone; inspect outcomes with schedules. Asks the user first. Optional tags label it.", {"at": {"type": "string"}, "every": {"type": "integer"}, "prompt": {"type": "string"}, "name": {"type": "string"}, "timeout": {"type": "integer"}, "profile": {"type": "string"}, "tags": {"type": "array"}}, ["prompt"]),
|
||||
ToolRegistry.tool_schema("unschedule", "Delete a scheduled prompt by id. Asks the user first.", {"id": {"type": "integer"}}, ["id"]),
|
||||
ToolRegistry.tool_schema("schedules", "List scheduled prompts with due times and last outcomes.", {}, []),
|
||||
ToolRegistry.tool_schema("record_save", "Save text as a tagged record in the vault and get back a mem:<hash> id. Use for anything too big or too durable for chat: research notes, command output, transcripts. Kind is one of note, output, file, research, transcript. Tag rules: lowercase singular, reuse tags from the tags tool, at most 5; known words in the text attach automatically and the record links to recent same-tag records.", {"title": {"type": "string"}, "content": {"type": "string"}, "kind": {"type": "string"}, "tags": {"type": "array"}}, ["content"]),
|
||||
ToolRegistry.tool_schema("record_read", "Read a slice of a record by mem:<hash> id. Offset and limit page through huge records without loading them whole.", {"id": {"type": "string"}, "offset": {"type": "integer"}, "limit": {"type": "integer"}}, ["id"]),
|
||||
ToolRegistry.tool_schema("record_search", "Search records by text, kind, and tags. Returns mem:<hash> ids with sizes so you can page in only what you need.", {"query": {"type": "string"}, "kind": {"type": "string"}, "tags": {"type": "array"}, "limit": {"type": "integer"}}, []),
|
||||
ToolRegistry.tool_schema("record_delete", "Delete a record by mem:<hash> id. Asks the user first.", {"id": {"type": "string"}}, ["id"]),
|
||||
ToolRegistry.tool_schema("graph_link", "Link two vault nodes with a relation: mem:<hash> records, secret:<name> secrets, sched:<id> schedules. Both endpoints must exist.", {"src": {"type": "string"}, "dst": {"type": "string"}, "relation": {"type": "string"}}, ["src", "dst"]),
|
||||
ToolRegistry.tool_schema("graph_query", "Show the neighborhood of a vault node (mem:<hash>, secret:<name>, sched:<id>) by breadth-first traversal. Depth caps at 4.", {"node": {"type": "string"}, "depth": {"type": "integer"}, "limit": {"type": "integer"}}, ["node"]),
|
||||
ToolRegistry.tool_schema("delete_file", "Delete a file. Requires reading it first; the full pre-image stays in the audit trail so restore can undelete. Asks the user first.", {"path": {"type": "string"}}, ["path"]),
|
||||
ToolRegistry.tool_schema("audit", "Show the file audit trail for time travel: every write, edit, delete, shell snapshot, release, and restore with sizes and messages. Filter by exact path or tag, or rank by a full-text query over paths and messages.", {"path": {"type": "string"}, "tag": {"type": "string"}, "query": {"type": "string"}, "limit": {"type": "integer"}}, []),
|
||||
ToolRegistry.tool_schema("restore", "Restore a file to an audit row: the post-image for writes and edits, the pre-image for deletes and shell snapshots. Refuses truncated images. Asks the user first.", {"id": {"type": "integer"}}, ["id"]),
|
||||
ToolRegistry.tool_schema("release", "Cut a release of the running script: bump VERSION by part (patch for fixes, minor for features, major for breaking changes), keep a versioned backup, log the message in the audit trail. Asks the user first.", {"part": {"type": "string"}, "message": {"type": "string"}}, ["part", "message"]),
|
||||
ToolRegistry.tool_schema("tags", "List vault tags with usage counts, most used first. Consult before tagging so new items reuse established tags instead of inventing synonyms.", {"prefix": {"type": "string"}, "limit": {"type": "integer"}}, []),
|
||||
ToolRegistry.tool_schema("install", "Manage tai installations: binary, command-not-found hook and plexer auto-wrap hook for bash/zsh/fish, venv, scheduler service, telegram service, container. Status reports what exists. Install adds missing pieces (skipping a shell hook loudly if that shell isn't on the machine), upgrade refreshes in place and restarts services, reinstall rebuilds artifacts, uninstall removes them. Service changes take effect immediately. The vault is never touched by any action. Non-status actions ask the user first.", {"action": {"type": "string"}, "targets": {"type": "array"}}, ["action"]),
|
||||
ToolRegistry.tool_schema("create_bot", "Create a bot in the current profile: a name plus its own system message and history. Give at least one of description, rules, or behavior; optional nicknames register as @mention aliases (a short form of the name registers automatically). Names are lowercase.", {"name": {"type": "string"}, "description": {"type": "string"}, "rules": {"type": "string"}, "behavior": {"type": "string"}, "nicknames": {"type": "array"}}, ["name"]),
|
||||
ToolRegistry.tool_schema("search", "Search everything at once with ranked full-text matching (SQLite FTS5 + bm25, no vectors): records, episodic events, the audit trail, and RAG chunks. Chunks are the retrieval-grade granular hits (e.g. one slice of an exported terminal log) and come back with their actual text, not just a snippet; use the chunk's parent id with record_read to see more surrounding context. One call replaces paging through each store separately. expand adds one graph hop per record hit.", {"query": {"type": "string"}, "kinds": {"type": "array"}, "limit": {"type": "integer"}, "expand": {"type": "boolean"}}, ["query"]),
|
||||
]
|
||||
|
||||
|
||||
# Tools gated on top of the lazy keyword match: a gate must pass before a tool can even be
|
||||
# offered to the model, regardless of relevance. Used for tools that only make sense in a
|
||||
# specific environment (e.g. a tmux pane to capture).
|
||||
TOOL_ENV_GATES = {
|
||||
"export_terminal_log": lambda: bool(os.environ.get("TMUX")),
|
||||
"plexer_send": lambda: PlexerControllerFactory.detect() is not None,
|
||||
"plexer_read": lambda: PlexerControllerFactory.detect() is not None,
|
||||
"plexer_type": lambda: PlexerControllerFactory.detect() is not None,
|
||||
"plexer_status": lambda: PlexerControllerFactory.detect() is not None,
|
||||
}
|
||||
+196
-23
@@ -122,7 +122,7 @@ class SysinfoTests(unittest.TestCase):
|
||||
return "ok"
|
||||
|
||||
stubs = tuple(("slow%d" % pos, slow) for pos in range(4))
|
||||
with mock.patch.object(tai, "SYSINFO_CHECKS", stubs):
|
||||
with mock.patch("tai_pkg.sysinfo.SYSINFO_CHECKS", stubs):
|
||||
started = time.time()
|
||||
report = tai.collect_sysinfo()
|
||||
self.assertLess(time.time() - started, 1.0)
|
||||
@@ -133,8 +133,8 @@ class SysinfoTests(unittest.TestCase):
|
||||
time.sleep(5)
|
||||
return "never"
|
||||
|
||||
with mock.patch.object(tai, "SYSINFO_CHECKS", (("stuck", stuck),)):
|
||||
with mock.patch.object(tai, "SYSINFO_TIMEOUT", 1):
|
||||
with mock.patch("tai_pkg.sysinfo.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)
|
||||
|
||||
@@ -142,7 +142,7 @@ class SysinfoTests(unittest.TestCase):
|
||||
def broken():
|
||||
raise RuntimeError("boom")
|
||||
|
||||
with mock.patch.object(tai, "SYSINFO_CHECKS", (("broken", broken),)):
|
||||
with mock.patch("tai_pkg.sysinfo.SYSINFO_CHECKS", (("broken", broken),)):
|
||||
report = tai.collect_sysinfo()
|
||||
self.assertIn("broken: error: boom", report)
|
||||
|
||||
@@ -165,7 +165,7 @@ class SysinfoTests(unittest.TestCase):
|
||||
def test_sysinfo_command(self):
|
||||
self.assertIn("sysinfo", tai.COMMANDS)
|
||||
self.assertIn("/sysinfo", tai.HELP_TEXT)
|
||||
with mock.patch.object(tai, "collect_sysinfo", return_value="fake-report") as collector:
|
||||
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()
|
||||
@@ -1616,7 +1616,7 @@ class InstallOpsTests(unittest.TestCase):
|
||||
tai._BOX_PYTHON_CACHE.clear()
|
||||
|
||||
def test_status_reports_all_targets(self):
|
||||
with mock.patch.object(tai, "container_engine", return_value=None):
|
||||
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:",
|
||||
@@ -1652,7 +1652,7 @@ class InstallOpsTests(unittest.TestCase):
|
||||
|
||||
def test_zsh_hook_install_remove_idempotent(self):
|
||||
path = os.path.join(self.home, ".zshrc")
|
||||
with mock.patch.object(tai, "shell_present", return_value=True):
|
||||
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()
|
||||
@@ -1668,7 +1668,7 @@ class InstallOpsTests(unittest.TestCase):
|
||||
|
||||
def test_fish_hook_install_remove_idempotent(self):
|
||||
path = os.path.join(self.home, ".config", "fish", "config.fish")
|
||||
with mock.patch.object(tai, "shell_present", return_value=True):
|
||||
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:
|
||||
@@ -1682,7 +1682,7 @@ class InstallOpsTests(unittest.TestCase):
|
||||
self.assertTrue(os.path.isfile(path + ".bak-tai"))
|
||||
|
||||
def test_shell_hook_skips_loudly_when_shell_missing(self):
|
||||
with mock.patch.object(tai, "shell_present", return_value=False):
|
||||
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(
|
||||
@@ -1695,15 +1695,15 @@ class InstallOpsTests(unittest.TestCase):
|
||||
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, "shell_present", return_value=True):
|
||||
with mock.patch.object(tai.Install, "shell_present", return_value=True):
|
||||
tai.op_zsh_hook("install", self.config.home)
|
||||
with mock.patch.object(tai, "shell_present", return_value=False):
|
||||
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, "shell_present", return_value=True):
|
||||
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:
|
||||
@@ -1715,7 +1715,7 @@ class InstallOpsTests(unittest.TestCase):
|
||||
|
||||
def test_fish_plexer_hook_uses_fish_syntax(self):
|
||||
path = os.path.join(self.home, ".config", "fish", "config.fish")
|
||||
with mock.patch.object(tai, "shell_present", return_value=True):
|
||||
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()
|
||||
@@ -1727,7 +1727,7 @@ class InstallOpsTests(unittest.TestCase):
|
||||
self.assertNotIn("[ -z \"$TAI_PLEXER_ACTIVE\" ]", content)
|
||||
|
||||
def test_install_self_detects_current_shell(self):
|
||||
with mock.patch.object(tai, "shell_present", return_value=True):
|
||||
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()
|
||||
@@ -1749,7 +1749,7 @@ class InstallOpsTests(unittest.TestCase):
|
||||
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, "venv_available", return_value=False):
|
||||
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):
|
||||
@@ -1764,7 +1764,7 @@ class InstallOpsTests(unittest.TestCase):
|
||||
outcome.stderr = ""
|
||||
return outcome
|
||||
|
||||
with mock.patch.object(tai, "venv_available", return_value=True):
|
||||
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))
|
||||
|
||||
@@ -1789,7 +1789,7 @@ class InstallOpsTests(unittest.TestCase):
|
||||
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, "ensure_venv", return_value=(True, "created fake")):
|
||||
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())
|
||||
@@ -1798,7 +1798,7 @@ class InstallOpsTests(unittest.TestCase):
|
||||
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, "container_engine", return_value=None):
|
||||
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
|
||||
@@ -1819,13 +1819,13 @@ class InstallOpsTests(unittest.TestCase):
|
||||
self.assertIn("python3-venv", tai.BOX_CONTAINERFILE)
|
||||
probe = mock.Mock()
|
||||
probe.returncode = 0
|
||||
with mock.patch.object(tai, "box_exec", return_value=probe) as runner:
|
||||
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, "box_exec", return_value=probe):
|
||||
with mock.patch.object(tai.Sandbox, "box_exec", return_value=probe):
|
||||
self.assertEqual(tai.box_python("podman"), "python3")
|
||||
|
||||
def test_unit_templates_format(self):
|
||||
@@ -1837,7 +1837,7 @@ class InstallOpsTests(unittest.TestCase):
|
||||
|
||||
def test_install_repl(self):
|
||||
self.app.tools = self.tools
|
||||
with mock.patch.object(tai, "container_engine", return_value=None):
|
||||
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")
|
||||
@@ -2400,7 +2400,7 @@ class ProgressTests(unittest.TestCase):
|
||||
|
||||
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, "TOTAL_STEP_CAP", 6):
|
||||
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, "")
|
||||
@@ -3275,7 +3275,7 @@ class MarkdownTests(unittest.TestCase):
|
||||
with mock.patch("sys.stdout") as fake_out:
|
||||
fake_out.isatty.return_value = True
|
||||
with mock.patch("builtins.print") as printer:
|
||||
with mock.patch.object(tai, "Spinner"):
|
||||
with mock.patch("tai_pkg.render.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)
|
||||
@@ -3290,5 +3290,178 @@ class MarkdownTests(unittest.TestCase):
|
||||
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_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 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 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)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user