tai/tai_pkg/agent.py

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