343 lines
12 KiB
Python
343 lines
12 KiB
Python
|
"""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()
|