From b37281390fb2944a76bb4c4124dbc452e1e749e3 Mon Sep 17 00:00:00 2001 From: retoor Date: Mon, 12 Oct 2026 19:10:06 +0200 Subject: [PATCH] Per-shell conversation contexts: ShellSession id per shell instance, TAI_SESSION var, vault session tree, /ctx Each bash/zsh/fish instance resolves a stable dated id (no install hook); transcripts and turn state scope per session so every new shell starts fresh, while vault data stays shared. Parent links accrete in the shell_sessions table with stub rows for outer shells, and /ctx renders the ancestor path plus children. TAI_SESSION pins or overrides the scope. --- README.md | 1 + tai.py | 409 +++++++++++++++++++++++++++++++++++++++++++++++----- test_tai.py | 237 ++++++++++++++++++++++++++++++ 3 files changed, 611 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index 0f3f65f..7bebd54 100644 --- a/README.md +++ b/README.md @@ -351,6 +351,7 @@ The interactive UI toolkit lives in `tai.py` itself (stdlib-only, no dependencie | `DEVPLACE_API_KEY` | Fallback backend credential | Empty, fallback disabled | | `TAI_VOICE` | Edge voice name | `en-US-EmmaMultilingualNeural` | | `TAI_PASSPHRASE` | Personal seal key | Unset: built-in key; empty: plaintext | +| `TAI_SESSION` | Shell session id (context scope) | Auto-derived per shell; fixed word shares | | `TELEGRAM_BOT_TOKEN` | Bot token, or set during install | Empty | ## Telegram diff --git a/tai.py b/tai.py index 63dcc2c..e5864de 100755 --- a/tai.py +++ b/tai.py @@ -99,7 +99,7 @@ def _fromisoformat(text): return parsed.replace(tzinfo=tzinfo) raise ValueError("invalid datetime, use ISO like 2026-10-08T09:00") -VERSION = "1.34.0" +VERSION = "1.35.0" WORKER_STEPS = 12 CREATE_SKILL_STEPS = 40 FORK_MAX_DEPTH = 2 @@ -480,6 +480,7 @@ HELP_TEXT = ( " /bots list bots, current marked with *\n" " /bot switch bot, history resumes\n" " /env [home|sandbox] show or switch execution environment\n" + " /ctx show shell session tree and recent sessions\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" @@ -2660,6 +2661,295 @@ def parse_backends(spec): return found +class ShellSession: + """One interactive shell is one conversation context. + + Spec: every bash/zsh/fish instance gets a stable logical id. All tai runs + inside that shell share one conversation history; a new shell starts fresh + (system message only), so the opening messages are identical on every start. + Vault data (records, secrets, memory, audit) stays shared across shells; + only the chat transcript and turn state are scoped per session. + + Id format (derived): YYYYMMDD-HHMMSS-p-b: local date and time + of the shell's start, the shell pid, and 8 hex chars of the kernel boot id. + Deterministic per shell (no state file, no install hook), unique across + reboots. Fallback with no shell ancestor or no /proc is ephemeral: + YYYYMMDD-HHMMSS-e-, fresh per process. + + Resolution (ShellSession.current): an explicit TAI_SESSION wins when it is + a plain word (a pinned, deliberately shared context) or when its embedded + -p- matches the nearest shell ancestor (the same shell coming back). + An id that embeds a different pid means a nested shell that inherited the + variable, so a fresh child id is derived and the inherited value becomes + its parent link. boot() exports the resolved id back into os.environ, which + makes TAI_SESSION the global var every child process (shell tool, nested + tai) inherits. Set TAI_SESSION to one fixed word to share a single context + across all shells; unset it for strict per-shell isolation. + + Parentage (for fun, queryable): register() upserts one shell_sessions row + per boot plus stub rows for outer shell ancestors, whose ids derive from + the same formula, so the whole nesting tree accretes over time. Dangling + parents after pruning are fine; this table is a log, not a constraint. + /ctx (ShellSession.report) renders the ancestor path and children of the + current session from these links. + + Caps: sid charset [A-Za-z0-9_.-], SID_MAX chars; CHAIN_CAP process levels; + FILES_CAP scoped transcript files per profile and bot (oldest pruned on + save); ROWS_CAP shell_sessions rows (oldest pruned on register). + """ + + ENV = "TAI_SESSION" + SHELLS = frozenset(("bash", "zsh", "fish", "sh", "dash", "ksh", "tcsh", "nu")) + SID_MAX = 64 + PID_RE = re.compile(r"-p(\d+)-") + CHAIN_CAP = 64 + FILES_CAP = 20 + ROWS_CAP = 1000 + TREE_DEPTH = 8 + TREE_KIDS = 10 + TREE_LINES = 40 + + def __init__(self, sid, shell="", pid=0, started="", parent=None, chain=()): + self.sid = sid + self.shell = shell + self.pid = pid + self.started = started + self.parent = parent + self.chain = list(chain) + + @classmethod + def sanitize(cls, value): + return re.sub(r"[^A-Za-z0-9_.-]", "", str(value or ""))[:cls.SID_MAX] + + @classmethod + def coerce(cls, value): + if isinstance(value, ShellSession): + return value.sid + return cls.sanitize(value) + + @staticmethod + def read_text(path): + try: + with open(path, "r", encoding="utf-8", errors="replace") as handle: + return handle.read(4096) + except OSError: + return "" + + @classmethod + def parse_stat(cls, text): + head, _, tail = (text or "").rpartition(")") + if not head or "(" not in head: + return None + comm = head.split("(", 1)[1] + fields = tail.split() + try: + return comm, int(fields[1]), int(fields[19]) + except (IndexError, ValueError): + return None + + @classmethod + def boot_time(cls, proc_root): + try: + with open(os.path.join(proc_root, "stat"), "r", encoding="utf-8") as handle: + for line in handle: + if line.startswith("btime "): + return float(line.split()[1]) + except (OSError, ValueError): + return None + return None + + @classmethod + def clk_tick(cls): + try: + return int(os.sysconf("SC_CLK_TCK")) + except (AttributeError, OSError, ValueError): + return 100 + + @classmethod + def boot_tag(cls, proc_root): + digest = re.sub(r"[^0-9a-f]", "", cls.read_text(os.path.join(proc_root, "sys/kernel/random/boot_id")).lower())[:8] + if digest: + return digest + btime = cls.boot_time(proc_root) + if btime: + return hashlib.md5(str(btime).encode("utf-8")).hexdigest()[:8] + return "deadbeef" + + @classmethod + def shell_name(cls, argv0, comm): + for cand in (os.path.basename((argv0 or "").split("\0")[0]).lstrip("-").lower(), (comm or "").lstrip("-").lower()): + if cand in cls.SHELLS: + return cand + return "" + + @classmethod + def read_chain(cls, pid, proc_root, btime, tick, cap): + chain = [] + seen = set() + current = pid + while current and current not in seen and len(chain) < cap: + seen.add(current) + parsed = cls.parse_stat(cls.read_text(os.path.join(proc_root, str(current), "stat"))) + if parsed is None: + break + comm, ppid, ticks = parsed + argv0 = os.path.basename(cls.read_text(os.path.join(proc_root, str(current), "cmdline")).split("\0")[0])[:64] + start = (btime + ticks / tick) if btime else 0.0 + chain.append({"pid": current, "ppid": ppid, "comm": comm[:64], "argv0": argv0, + "name": cls.shell_name(argv0, comm), "start": start, + "start_iso": datetime.fromtimestamp(start, tz=timezone.utc).isoformat(timespec="seconds") if start else "", + "sid": ""}) + if ppid <= 0 or ppid == current: + break + current = ppid + return chain + + @classmethod + def shell_id(cls, start_ts, pid, boot): + return "%s-p%d-%s" % (time.strftime("%Y%m%d-%H%M%S", time.localtime(start_ts)), pid, boot) + + @classmethod + def ephemeral_id(cls, now_ts, pid): + return "%s-e%d-%s" % (time.strftime("%Y%m%d-%H%M%S", time.localtime(now_ts)), pid, uuid.uuid4().hex[:4]) + + @classmethod + def current(cls, env=None, pid=None, proc_root="/proc", now=None, clk_tck=None): + source = os.environ if env is None else env + raw = cls.sanitize(source.get(cls.ENV) or "") + pid = os.getpid() if pid is None else pid + now = time.time() if now is None else now + btime = cls.boot_time(proc_root) + tick = clk_tck or cls.clk_tick() + boot = cls.boot_tag(proc_root) + chain = cls.read_chain(pid, proc_root, btime, tick, cls.CHAIN_CAP) + for entry in chain: + if entry["name"] and entry["start"]: + entry["sid"] = cls.shell_id(entry["start"], entry["pid"], boot) + shells = [entry for entry in chain if entry["sid"]] + inner = shells[0] if shells else None + outer = shells[1] if len(shells) > 1 else None + outer_sid = outer["sid"] if outer else None + if raw: + match = cls.PID_RE.search(raw) + if match is None: + return cls(raw, shell=inner["name"] if inner else "", pid=inner["pid"] if inner else pid, + started=inner["start_iso"] if inner else "", chain=chain) + if inner is not None and int(match.group(1)) == inner["pid"]: + return cls(raw, shell=inner["name"], pid=inner["pid"], started=inner["start_iso"], + parent=outer_sid, chain=chain) + fresh = inner["sid"] if inner else cls.ephemeral_id(now, pid) + started = inner["start_iso"] if inner else datetime.fromtimestamp(now, tz=timezone.utc).isoformat(timespec="seconds") + return cls(fresh, shell=inner["name"] if inner else "", pid=inner["pid"] if inner else pid, + started=started, parent=raw, chain=chain) + if inner is not None: + return cls(inner["sid"], shell=inner["name"], pid=inner["pid"], started=inner["start_iso"], + parent=outer_sid, chain=chain) + return cls(cls.ephemeral_id(now, pid), pid=pid, + started=datetime.fromtimestamp(now, tz=timezone.utc).isoformat(timespec="seconds"), chain=chain) + + def export(self): + os.environ[self.ENV] = self.sid + + def register(self, store): + try: + stamp = now_iso() + shells = [entry for entry in self.chain if entry.get("sid")] + for pos, entry in enumerate(shells): + next_sid = shells[pos + 1]["sid"] if pos + 1 < len(shells) else None + store.db.execute("INSERT OR IGNORE INTO shell_sessions (sid, parent, shell, pid, started) VALUES (?, ?, ?, ?, ?)", + (entry["sid"], next_sid, entry["name"], entry["pid"], entry["start_iso"])) + store.db.execute("INSERT OR IGNORE INTO shell_sessions (sid, parent, shell, pid, started, last_seen, cwd, chain) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", + (self.sid, self.parent, self.shell, self.pid, self.started, stamp, dir_key(), json.dumps(self.chain))) + store.db.execute("UPDATE shell_sessions SET last_seen = ?, cwd = ?, parent = COALESCE(parent, ?), chain = ? WHERE sid = ?", + (stamp, dir_key(), self.parent, json.dumps(self.chain), self.sid)) + store.db.execute("DELETE FROM shell_sessions WHERE sid NOT IN (SELECT sid FROM shell_sessions ORDER BY COALESCE(last_seen, started, '') DESC LIMIT ?)", + (self.ROWS_CAP,)) + store.db.commit() + except (sqlite3.Error, OSError, ValueError): + pass + + @classmethod + def tree_lines(cls, store, sid): + """Hierarchy around sid: ancestor path root-first, then the children + below. Cycle-safe, capped by TREE_DEPTH/TREE_KIDS/TREE_LINES.""" + try: + rows = store.db.execute("SELECT sid, parent, shell, pid FROM shell_sessions").fetchall() + except sqlite3.Error: + return [] + by_sid = {} + kids = {} + for item in rows: + by_sid[item[0]] = item + if item[1]: + kids.setdefault(item[1], []).append(item[0]) + for key in kids: + kids[key].sort() + if sid not in by_sid: + return [] + + def label(key): + item = by_sid.get(key) + if item is None: + return "%s (pruned)" % key + return "%s %s%s" % (key, item[2] or "?", (" pid %s" % item[3]) if item[3] else "") + + upwards = [] + seen = {sid} + node = sid + while len(upwards) < cls.TREE_DEPTH: + parent = by_sid[node][1] if node in by_sid else None + if not parent or parent in seen: + break + seen.add(parent) + upwards.append(parent) + node = parent + lines = [" " * depth + label(key) for depth, key in enumerate(reversed(upwards))] + lines.append(" " * len(upwards) + "* " + label(sid)) + + def below(key, depth, seen_here): + out = [] + if depth > cls.TREE_DEPTH or len(lines) + len(out) >= cls.TREE_LINES: + return out + for child in kids.get(key, [])[:cls.TREE_KIDS]: + if child in seen_here or len(lines) + len(out) >= cls.TREE_LINES: + continue + out.append(" " * depth + label(child)) + out.extend(below(child, depth + 1, seen_here | {child})) + return out + + lines.extend(below(sid, len(upwards) + 1, seen)) + return lines[:cls.TREE_LINES] + + @classmethod + def report(cls, store, sid): + try: + row = store.db.execute("SELECT sid, parent, shell, pid, started, last_seen, cwd FROM shell_sessions WHERE sid = ?", (sid,)).fetchone() if sid else None + except sqlite3.Error: + return "ctx unavailable (vault error)" + if row: + lines = ["ctx: %s" % row[0], + "shell: %s (pid %s, started %s)" % (row[2] or "?", row[3] or "?", row[4] or "?"), + "parent: %s" % (row[1] or "(none)"), + "cwd: %s" % (row[6] or "?")] + elif sid: + lines = ["ctx: %s (unregistered)" % sid] + else: + lines = ["ctx: (legacy, no session)"] + tree = cls.tree_lines(store, sid) + if tree: + lines.append("tree:") + lines.extend(tree) + try: + rows = store.db.execute("SELECT sid, shell, last_seen FROM shell_sessions ORDER BY COALESCE(last_seen, started, '') DESC LIMIT 8").fetchall() + except sqlite3.Error: + rows = [] + if rows: + lines.append("recent:") + for item in rows: + lines.append("%s %s %s %s" % ("*" if item[0] == sid else " ", item[0], item[1] or "?", item[2] or "")) + return "\n".join(lines) + + class Config: def __init__(self, args): self.home = os.environ.get("TAI_HOME") or os.path.join(os.path.expanduser("~"), ".tai") @@ -2806,6 +3096,7 @@ SCHEMA = ( ("chunks", ("id INTEGER PRIMARY KEY", "parent_kind TEXT", "parent_id TEXT", "idx INTEGER", "text TEXT", "profile TEXT", "created TEXT"), ()), ("dir_defaults", ("dir TEXT PRIMARY KEY", "provider TEXT", "model TEXT", "bot TEXT", "updated TEXT"), ()), ("input_history", ("id INTEGER PRIMARY KEY", "profile TEXT", "bot TEXT", "ts TEXT", "text TEXT"), ()), + ("shell_sessions", ("sid TEXT PRIMARY KEY", "parent TEXT", "shell TEXT", "pid INTEGER", "started TEXT", "last_seen TEXT", "cwd TEXT", "chain TEXT"), ()), ) SCHEMA_INDEXES = ( @@ -2821,6 +3112,7 @@ SCHEMA_INDEXES = ( ("idx_chunks_parent", "chunks(parent_kind, parent_id)"), ("idx_chunks_profile", "chunks(profile)"), ("idx_input_history_profile", "input_history(profile, id)"), + ("idx_shell_sessions_parent", "shell_sessions(parent)"), ) @@ -3549,7 +3841,10 @@ class Store: def profile_path(self, name): return os.path.join(self.config.profiles_dir, name + ".sys.md") - def session_path(self, name): + def session_path(self, name, ctx=None): + sid = ShellSession.sanitize(ctx) if ctx else "" + if sid: + return os.path.join(self.config.profiles_dir, "%s.session.%s.json" % (name, sid)) return os.path.join(self.config.profiles_dir, name + ".session.json") def list_profiles(self): @@ -3588,8 +3883,8 @@ class Store: 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) + def load_session(self, name, bot="main", ctx=None): + path = self.session_path(name, ctx=ctx) if bot == "main" else self.bot_session_path(name, bot, ctx=ctx) if not os.path.exists(path): return [] try: @@ -3606,12 +3901,35 @@ class Store: 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) + def save_session(self, name, messages, bot="main", ctx=None): + path = self.session_path(name, ctx=ctx) if bot == "main" else self.bot_session_path(name, bot, ctx=ctx) try: sealed_write(path, self.seal.lock(json.dumps([item for item in messages if item.get("role") != "system"]))) except OSError: pass + if ShellSession.sanitize(ctx) if ctx else "": + stem = name if bot == "main" else "%s.bot.%s" % (name, bot) + self.prune_ctx_files("%s.session." % stem, "%s.session.json" % stem) + + def prune_ctx_files(self, prefix, legacy, cap=0): + try: + entries = [entry for entry in os.listdir(self.config.profiles_dir) + if entry.startswith(prefix) and entry.endswith(".json") and entry != legacy] + except OSError: + return + keep = cap or ShellSession.FILES_CAP + if len(entries) <= keep: + return + paths = [os.path.join(self.config.profiles_dir, entry) for entry in entries] + try: + paths.sort(key=lambda item: os.path.getmtime(item)) + except OSError: + return + for stale in paths[:len(paths) - keep]: + try: + os.remove(stale) + except OSError: + pass def load_dir_defaults(self, directory): try: @@ -3645,13 +3963,15 @@ class Store: except sqlite3.Error: 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 turn_state_path(self, profile, bot="main", ctx=None): + sid = ShellSession.sanitize(ctx) if ctx else "" + stem = profile if bot == "main" else "%s.bot.%s" % (profile, bot) + if sid: + return os.path.join(self.config.profiles_dir, "%s.turn.%s.json" % (stem, sid)) + return os.path.join(self.config.profiles_dir, "%s.turn.json" % stem) - def load_turn_state(self, profile, bot="main"): - path = self.turn_state_path(profile, bot) + def load_turn_state(self, profile, bot="main", ctx=None): + path = self.turn_state_path(profile, bot, ctx=ctx) if not os.path.exists(path): return {} try: @@ -3667,16 +3987,22 @@ class Store: quarantine(path) return {} - def save_turn_state(self, profile, bot, state): + def save_turn_state(self, profile, bot, state, ctx=None): try: - sealed_write(self.turn_state_path(profile, bot), self.seal.lock(json.dumps(state))) + sealed_write(self.turn_state_path(profile, bot, ctx=ctx), self.seal.lock(json.dumps(state))) except OSError: pass + if ShellSession.sanitize(ctx) if ctx else "": + stem = profile if bot == "main" else "%s.bot.%s" % (profile, bot) + self.prune_ctx_files("%s.turn." % stem, "%s.turn.json" % stem) 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): + def bot_session_path(self, profile, bot, ctx=None): + sid = ShellSession.sanitize(ctx) if ctx else "" + if sid: + return os.path.join(self.config.profiles_dir, "%s.bot.%s.session.%s.json" % (profile, bot, sid)) return os.path.join(self.config.profiles_dir, "%s.bot.%s.session.json" % (profile, bot)) def bots_registry_path(self, profile): @@ -14530,10 +14856,11 @@ class Denied(Exception): class Agent: - def __init__(self, config, store, persist=True, quiet=False, depth=0, auto=False, yolo=False): + def __init__(self, config, store, persist=True, quiet=False, depth=0, auto=False, yolo=False, ctx=None): self.config = config self.auto = auto self.yolo = yolo + self.ctx_id = ShellSession.coerce(ctx) self.store = store self.chat = ChatClient(config) if not persist or depth > 0: @@ -14577,7 +14904,7 @@ class Agent: def switch_profile(self, name, silent=False): if self.messages and self.persist: - self.store.save_session(self.profile, self.messages, self.bot) + self.store.save_session(self.profile, self.messages, self.bot, ctx=self.ctx_id) self.profile = name self.bot = "main" self.store.profile = name @@ -14591,7 +14918,7 @@ class Agent: loaded = loaded.replace(LEGACY_PASSWORD_NOTE, VAULT_MEMORY_NOTE) self.store.save_system(name, loaded) self.system_message = loaded - restored = self.store.load_session(name) if self.persist else [] + restored = self.store.load_session(name, ctx=self.ctx_id) if self.persist else [] self.messages = [{"role": "system", "content": self.system_message}] + restored self.apply_system() self.sync_chat_scope() @@ -14612,7 +14939,7 @@ class Agent: 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.store.save_session(self.profile, self.messages, self.bot, ctx=self.ctx_id) self.bot = target loaded = self.store.load_bot_system(self.profile, target) if loaded is None: @@ -14622,7 +14949,7 @@ class Agent: 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 [] + restored = self.store.load_session(self.profile, target, ctx=self.ctx_id) if self.persist else [] self.messages = [{"role": "system", "content": self.system_message}] + restored self.apply_system() self.sync_chat_scope() @@ -14762,18 +15089,18 @@ class Agent: def checkpoint(self): if not self.persist: return - self.store.save_session(self.profile, self.messages, self.bot) + self.store.save_session(self.profile, self.messages, self.bot, ctx=self.ctx_id) 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": now_iso()}) + self.store.save_turn_state(self.profile, self.bot, {"open": True, "goal": self._turn_goal, "steps": self._turn_steps, "tools": self._turn_tools, "updated": now_iso()}, ctx=self.ctx_id) def close_turn(self): if not self.persist: self._turn_open = False return - self.store.save_session(self.profile, self.messages, self.bot) + self.store.save_session(self.profile, self.messages, self.bot, ctx=self.ctx_id) 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": now_iso()}) + self.store.save_turn_state(self.profile, self.bot, {"open": False, "goal": self._turn_goal, "steps": self._turn_steps, "tools": self._turn_tools, "updated": now_iso()}, ctx=self.ctx_id) def mention_error(self, tag): bots = ", ".join(item["name"] for item in self.store.list_bots(self.profile)) @@ -15013,8 +15340,8 @@ class Agent: def run_mention(self, target, rest, original, capture=False, max_steps=MAX_STEPS): if self.persist: - self.store.save_session(self.profile, self.messages, self.bot) - other = self.store.load_session(self.profile, target) + self.store.save_session(self.profile, self.messages, self.bot, ctx=self.ctx_id) + other = self.store.load_session(self.profile, target, ctx=self.ctx_id) else: other = [] system = self.store.load_bot_system(self.profile, target) or DEFAULT_SYSTEM @@ -15028,7 +15355,7 @@ class Agent: 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.store.save_session(self.profile, self.messages, target, ctx=self.ctx_id) self.messages, self.system_message, self.bot = saved_messages, saved_system, saved_bot self.apply_system() self.sync_chat_scope() @@ -15036,7 +15363,7 @@ class Agent: 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) + self.store.save_session(self.profile, self.messages, self.bot, ctx=self.ctx_id) return answer @@ -15259,7 +15586,7 @@ def start_scheduler(config, seal): return stop_flag -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", "demo", "quit", "exit") +COMMANDS = ("help", "profile", "profiles", "bots", "bot", "env", "ctx", "skills", "sysinfo", "models", "secret", "schedule", "schedules", "unschedule", "records", "record", "graph", "tags", "tools", "install", "search", "audit", "restore", "release", "fork", "agents", "agent", "compact", "clear", "demo", "quit", "exit") def complete_command(text, state): @@ -15361,6 +15688,8 @@ def handle_command(agent, text): else: agent.env = "home" print("environment: home") + elif name == "ctx": + print(ShellSession.report(agent.store, getattr(agent, "ctx_id", ""))) elif name == "skills": bits = arg.split() if bits and bits[0] == "install": @@ -16369,6 +16698,9 @@ def boot(args): store = Store(config, seal) if store.schema_notes: print("vault schema upgraded: %s" % ", ".join(store.schema_notes)) + ctx = ShellSession.current() + ctx.register(store) + ctx.export() try: backed_up = ensure_self_backup(config, store) except (OSError, sqlite3.Error) as exc: @@ -16376,7 +16708,7 @@ def boot(args): backed_up = None if backed_up is not None: print("self-backup: %s" % os.path.basename(backed_up)) - agent = Agent(config, store, yolo=args.yolo or args.auto, auto=args.auto) + agent = Agent(config, store, yolo=args.yolo or args.auto, auto=args.auto, ctx=ctx) try: loaded = apply_dir_defaults(agent) except OSError: @@ -22077,7 +22409,7 @@ def tui_spawn_agent(base, name): """A fresh Agent for another bot, sharing the vault with base. Starts with empty messages so switch_bot never saves a stale copy of base's conversation over the real one.""" - child = Agent(base.config, base.store, yolo=base.yolo, auto=base.auto) + child = Agent(base.config, base.store, yolo=base.yolo, auto=base.auto, ctx=getattr(base, "ctx_id", "")) child.messages = [] child.switch_profile(base.profile, silent=True) child.messages = [] @@ -22371,7 +22703,7 @@ def tui_build_chat(app, agent, spawn=None): def bot_snippet(name): try: - msgs = agent.messages if name == agent.bot else agent.store.load_session(agent.profile, name) + msgs = agent.messages if name == agent.bot else agent.store.load_session(agent.profile, name, ctx=getattr(agent, "ctx_id", "")) except Exception: return "" for msg in reversed(msgs or []): @@ -22803,7 +23135,7 @@ def tui_chat(agent, theme_name): scheduler_stop.set() owners = [item.agent for item in refs["sessions"].values()] if refs is not None else [agent] for owner in owners: - owner.store.save_session(owner.profile, owner.messages, owner.bot) + owner.store.save_session(owner.profile, owner.messages, owner.bot, ctx=getattr(owner, "ctx_id", "")) def build_parser(): @@ -22848,18 +23180,23 @@ def banner(agent): print(" %sseal:%s %s" % (Ansi.DIM, Ansi.RESET, seal_state)) else: print(" seal: %s" % seal_state) + if getattr(agent, "ctx_id", ""): + if Glow.enabled(): + print(" %sctx:%s %s" % (Ansi.DIM, Ansi.RESET, agent.ctx_id)) + else: + print(" ctx: %s" % agent.ctx_id) for label, _chat_url, models_url, key, _model in agent.chat.backends(): if Glow.enabled(): print(" %s%s:%s %s" % (Ansi.DIM, label, Ansi.RESET, probe_backend(models_url, key))) else: print(" %s: %s" % (label, probe_backend(models_url, key))) - if agent.store.load_turn_state(agent.profile, agent.bot).get("open"): + if agent.store.load_turn_state(agent.profile, agent.bot, ctx=getattr(agent, "ctx_id", "")).get("open"): print(paint("interrupted turn available, restart with --continue to pick it up", Ansi.YELLOW)) print(paint("type /help for commands", Ansi.DIM)) def resume_turn(agent): - state = agent.store.load_turn_state(agent.profile, agent.bot) + state = agent.store.load_turn_state(agent.profile, agent.bot, ctx=getattr(agent, "ctx_id", "")) if not state.get("open"): print("nothing to resume: last turn for %s/%s completed normally" % (agent.profile, agent.bot)) return "" @@ -22908,7 +23245,7 @@ def repl(agent): agent.checkpoint() finally: scheduler_stop.set() - agent.store.save_session(agent.profile, agent.messages, agent.bot) + agent.store.save_session(agent.profile, agent.messages, agent.bot, ctx=getattr(agent, "ctx_id", "")) def update_models_command(args): diff --git a/test_tai.py b/test_tai.py index e9a0948..e8aa5bd 100644 --- a/test_tai.py +++ b/test_tai.py @@ -9864,5 +9864,242 @@ class MplexPaneEnvTests(unittest.TestCase): self.assertEqual(controller.pane, "7") +class ShellSessionTests(unittest.TestCase): + BTIME = 1700000000.0 + BOOT_ID = "12345678-9abc-def0-1234-56789abcdef0\n" + + def write_proc(self, root, pid, ppid, comm, argv0, ticks): + folder = os.path.join(root, str(pid)) + os.makedirs(folder, exist_ok=True) + stat = "%d (%s) R %d %s %d 0 0" % (pid, comm, ppid, " ".join(["0"] * 17), ticks) + with open(os.path.join(folder, "stat"), "w", encoding="utf-8") as handle: + handle.write(stat) + with open(os.path.join(folder, "cmdline"), "w", encoding="utf-8") as handle: + handle.write(argv0 + "\0") + + def proc_root(self, tmp, members=((50, 20, "python3", "python3", 100000), + (20, 10, "bash", "/bin/bash", 50000), + (10, 1, "bash", "/bin/bash", 10000), + (1, 0, "systemd", "/sbin/init", 100))): + root = os.path.join(tmp, "proc") + for pid, ppid, comm, argv0, ticks in members: + self.write_proc(root, pid, ppid, comm, argv0, ticks) + with open(os.path.join(root, "stat"), "w", encoding="utf-8") as handle: + handle.write("cpu 0\nbtime %d\n" % int(self.BTIME)) + boot_dir = os.path.join(root, "sys/kernel/random") + os.makedirs(boot_dir, exist_ok=True) + with open(os.path.join(boot_dir, "boot_id"), "w", encoding="utf-8") as handle: + handle.write(self.BOOT_ID) + return root + + def test_derived_id_is_stable_dated_and_pid_scoped(self): + with tempfile.TemporaryDirectory() as tmp: + root = self.proc_root(tmp) + first = tai.ShellSession.current(env={}, pid=50, proc_root=root, clk_tck=100) + second = tai.ShellSession.current(env={}, pid=50, proc_root=root, clk_tck=100) + self.assertEqual(first.sid, second.sid) + self.assertEqual(first.sid, tai.ShellSession.shell_id(self.BTIME + 500, 20, "12345678")) + self.assertIn("-p20-", first.sid) + self.assertEqual((first.shell, first.pid), ("bash", 20)) + self.assertEqual(first.parent, tai.ShellSession.shell_id(self.BTIME + 100, 10, "12345678")) + + def test_env_reused_when_pid_matches_shell(self): + with tempfile.TemporaryDirectory() as tmp: + root = self.proc_root(tmp) + derived = tai.ShellSession.current(env={}, pid=50, proc_root=root, clk_tck=100).sid + again = tai.ShellSession.current(env={tai.ShellSession.ENV: derived}, pid=50, proc_root=root, clk_tck=100) + self.assertEqual(again.sid, derived) + self.assertEqual(again.shell, "bash") + + def test_env_superseded_by_nested_shell_links_parent(self): + with tempfile.TemporaryDirectory() as tmp: + root = self.proc_root(tmp, members=((50, 30, "python3", "python3", 100000), + (30, 20, "bash", "/bin/bash", 80000), + (20, 1, "bash", "/bin/bash", 50000), + (1, 0, "systemd", "/sbin/init", 100))) + outer = tai.ShellSession.shell_id(self.BTIME + 500, 20, "12345678") + ctx = tai.ShellSession.current(env={tai.ShellSession.ENV: outer}, pid=50, proc_root=root, clk_tck=100) + self.assertIn("-p30-", ctx.sid) + self.assertEqual(ctx.parent, outer) + + def test_pinned_word_is_used_verbatim(self): + with tempfile.TemporaryDirectory() as tmp: + root = self.proc_root(tmp) + ctx = tai.ShellSession.current(env={tai.ShellSession.ENV: "main"}, pid=50, proc_root=root, clk_tck=100) + self.assertEqual(ctx.sid, "main") + + def test_ephemeral_without_proc(self): + first = tai.ShellSession.current(env={}, pid=4242, proc_root="/nonexistent-proc", now=1700000000.0) + second = tai.ShellSession.current(env={}, pid=4242, proc_root="/nonexistent-proc", now=1700000000.0) + self.assertIn("-e4242-", first.sid) + self.assertNotEqual(first.sid, second.sid) + + def test_login_shell_argv0_detected(self): + with tempfile.TemporaryDirectory() as tmp: + root = self.proc_root(tmp, members=((50, 20, "python3", "python3", 100000), + (20, 1, "bash", "-bash", 50000), + (1, 0, "systemd", "/sbin/init", 100))) + ctx = tai.ShellSession.current(env={}, pid=50, proc_root=root, clk_tck=100) + self.assertEqual(ctx.shell, "bash") + self.assertIn("-p20-", ctx.sid) + + def test_sanitize_coerce_and_parse_stat(self): + self.assertEqual(tai.ShellSession.sanitize("../../x;rm A"), "....xrmA") + self.assertEqual(tai.ShellSession.sanitize("x" * 100), "x" * tai.ShellSession.SID_MAX) + self.assertEqual(tai.ShellSession.coerce(None), "") + ctx = tai.ShellSession("abc") + self.assertEqual(tai.ShellSession.coerce(ctx), "abc") + self.assertEqual(tai.ShellSession.coerce("a/b"), "ab") + self.assertIsNone(tai.ShellSession.parse_stat("garbage")) + self.assertIsNone(tai.ShellSession.parse_stat("")) + self.assertEqual(tai.ShellSession.parse_stat("9 (a(b) R 3 " + " ".join(["0"] * 17) + " 42 0")[1:], (3, 42)) + + def test_export_sets_global_var(self): + with mock.patch.dict(os.environ, {}, clear=False): + os.environ.pop(tai.ShellSession.ENV, None) + tai.ShellSession("s1").export() + self.assertEqual(os.environ[tai.ShellSession.ENV], "s1") + + +class CtxStoreTests(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + os.environ["TAI_HOME"] = self.tmp.name + + class FakeArgs: + profile = "t" + yes = False + + self.config = tai.Config(FakeArgs()) + self.store = tai.Store(self.config, tai.Seal(self.config.home, "")) + + def tearDown(self): + self.store.close() + self.tmp.cleanup() + os.environ.pop("TAI_HOME", None) + + def test_scoped_paths_leave_legacy_untouched(self): + self.assertTrue(self.store.session_path("t", ctx="abc").endswith("t.session.abc.json")) + self.assertTrue(self.store.session_path("t").endswith("t.session.json")) + self.assertTrue(self.store.bot_session_path("t", "coder", ctx="abc").endswith("t.bot.coder.session.abc.json")) + self.assertTrue(self.store.bot_session_path("t", "coder").endswith("t.bot.coder.session.json")) + self.assertTrue(self.store.turn_state_path("t", "main", ctx="abc").endswith("t.turn.abc.json")) + self.assertTrue(self.store.turn_state_path("t").endswith("t.turn.json")) + + def test_scoped_sessions_are_isolated_from_each_other_and_legacy(self): + self.store.save_session("t", [{"role": "user", "content": "legacy"}]) + self.store.save_session("t", [{"role": "user", "content": "one"}], ctx="one") + self.store.save_session("t", [{"role": "user", "content": "two"}], ctx="two") + self.assertEqual([item["content"] for item in self.store.load_session("t")], ["legacy"]) + self.assertEqual([item["content"] for item in self.store.load_session("t", ctx="one")], ["one"]) + self.assertEqual([item["content"] for item in self.store.load_session("t", ctx="two")], ["two"]) + + def test_ctx_in_path_is_filename_safe(self): + path = self.store.session_path("t", ctx="../../x;rm") + self.assertEqual(os.path.dirname(path), self.config.profiles_dir) + self.assertNotIn("/", os.path.basename(path)) + + def test_prune_ctx_files_keeps_newest_and_legacy(self): + self.store.save_session("t", [{"role": "user", "content": "legacy"}]) + for num in range(5): + self.store.save_session("t", [{"role": "user", "content": "m%d" % num}], ctx="c%d" % num) + self.store.prune_ctx_files("t.session.", "t.session.json", cap=3) + left = sorted(os.listdir(self.config.profiles_dir)) + self.assertIn("t.session.json", left) + self.assertEqual(len([entry for entry in left if entry.startswith("t.session.") and entry != "t.session.json"]), 3) + + def test_corrupt_scoped_session_quarantined(self): + path = self.store.session_path("t", ctx="one") + self.store.save_session("t", [{"role": "user", "content": "good"}], ctx="one") + with open(path, "w", encoding="utf-8") as handle: + handle.write("{not json") + self.assertEqual(self.store.load_session("t", ctx="one"), []) + leftovers = [entry for entry in os.listdir(self.config.profiles_dir) if ".corrupt-" in entry] + self.assertEqual(len(leftovers), 1) + + def test_agent_ctx_isolation(self): + first = tai.Agent(self.config, self.store, persist=True, quiet=True, ctx="one") + first.messages.append({"role": "user", "content": "hello one"}) + first.checkpoint() + same = tai.Agent(self.config, self.store, persist=True, quiet=True, ctx="one") + self.assertTrue(any("hello one" in (item.get("content") or "") for item in same.messages)) + other = tai.Agent(self.config, self.store, persist=True, quiet=True, ctx="two") + self.assertEqual([item["role"] for item in other.messages], ["system"]) + legacy = tai.Agent(self.config, self.store, persist=True, quiet=True) + self.assertEqual([item["role"] for item in legacy.messages], ["system"]) + self.assertEqual(legacy.ctx_id, "") + + def test_agent_ctx_turn_state_scoped(self): + agent = tai.Agent(self.config, self.store, persist=True, quiet=True, ctx="one") + agent._turn_open = True + agent._turn_goal = "goal one" + agent.checkpoint() + self.assertTrue(self.store.load_turn_state("t", "main", ctx="one")["open"]) + self.assertEqual(self.store.load_turn_state("t", "main", ctx="two"), {}) + self.assertEqual(self.store.load_turn_state("t", "main"), {}) + + def test_register_builds_parent_tree_and_report(self): + ctx = tai.ShellSession("inner", shell="bash", pid=20, started="s", parent="outer", + chain=[{"pid": 20, "ppid": 10, "comm": "bash", "argv0": "bash", "name": "bash", + "start": 1.0, "start_iso": "s", "sid": "inner"}, + {"pid": 10, "ppid": 1, "comm": "bash", "argv0": "bash", "name": "bash", + "start": 0.5, "start_iso": "s0", "sid": "outer"}]) + ctx.register(self.store) + rows = dict(self.store.db.execute("SELECT sid, parent FROM shell_sessions").fetchall()) + self.assertEqual(rows["inner"], "outer") + self.assertIn("outer", rows) + ctx.register(self.store) + self.assertEqual(self.store.db.execute("SELECT COUNT(*) FROM shell_sessions").fetchone()[0], 2) + text = tai.ShellSession.report(self.store, "inner") + self.assertIn("ctx: inner", text) + self.assertIn("parent: outer", text) + self.assertIn("recent:", text) + + def test_register_prunes_rows_and_report_handles_legacy(self): + with mock.patch.object(tai.ShellSession, "ROWS_CAP", 3): + for num in range(5): + self.store.db.execute("INSERT INTO shell_sessions (sid, started) VALUES (?, ?)", ("old%d" % num, "2020-01-0%d" % (num + 1))) + self.store.db.commit() + tai.ShellSession("fresh", chain=[]).register(self.store) + self.assertLessEqual(self.store.db.execute("SELECT COUNT(*) FROM shell_sessions").fetchone()[0], 3) + self.assertIn("(legacy", tai.ShellSession.report(self.store, "")) + + def seed_tree(self, rows): + for sid, parent in rows: + self.store.db.execute("INSERT INTO shell_sessions (sid, parent, shell, pid) VALUES (?, ?, ?, ?)", + (sid, parent, "bash", 1)) + self.store.db.commit() + + def test_tree_shows_lineage_and_children(self): + self.seed_tree([("root", None), ("mid", "root"), ("leaf", "mid"), ("sib", "root"), ("kid", "leaf")]) + lines = tai.ShellSession.tree_lines(self.store, "mid") + self.assertEqual([line.strip()[:4] for line in lines], ["root", "* mi", "leaf", "kid "]) + self.assertTrue(lines[0].startswith("root")) + self.assertTrue(lines[1].startswith(" * mid")) + self.assertTrue(lines[2].startswith(" leaf")) + text = tai.ShellSession.report(self.store, "mid") + self.assertIn("tree:", text) + self.assertIn("* mid", text) + + def test_tree_cycle_and_orphan_safe(self): + self.seed_tree([("a", "b"), ("b", "a")]) + lines = tai.ShellSession.tree_lines(self.store, "a") + self.assertTrue(any(line.strip().startswith("* a") for line in lines)) + self.store.db.execute("INSERT INTO shell_sessions (sid, parent, shell) VALUES (?, ?, ?)", ("orph", "gone", "sh")) + self.store.db.commit() + lines = tai.ShellSession.tree_lines(self.store, "orph") + self.assertTrue(any("(pruned)" in line for line in lines)) + self.assertEqual(tai.ShellSession.tree_lines(self.store, "missing"), []) + + def test_tree_caps_fanout_and_depth(self): + with mock.patch.object(tai.ShellSession, "TREE_KIDS", 2), mock.patch.object(tai.ShellSession, "TREE_LINES", 5): + rows = [("top", None)] + [("c%02d" % num, "top") for num in range(5)] + self.seed_tree(rows) + lines = tai.ShellSession.tree_lines(self.store, "top") + self.assertLessEqual(len(lines), 5) + self.assertTrue(any("c00" in line for line in lines)) + self.assertFalse(any("c04" in line for line in lines)) + + if __name__ == "__main__": unittest.main()