"""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)