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:
+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()
|
||||
Reference in New Issue
Block a user