diff --git a/tai.py b/tai.py index 5cfc288..4b63341 100755 --- a/tai.py +++ b/tai.py @@ -96,7 +96,7 @@ def _fromisoformat(text): return parsed.replace(tzinfo=tzinfo) raise ValueError("invalid datetime, use ISO like 2026-10-08T09:00") -VERSION = "1.22.1" +VERSION = "1.26.0" WORKER_STEPS = 12 CREATE_SKILL_STEPS = 40 FORK_MAX_DEPTH = 2 @@ -321,6 +321,18 @@ def orientation_block(directory=None, sandbox=False): % dir_key(directory)) +def capability_block(yolo=False, auto=False): + """Live approval state for the system prompt. Grounds the model in what + this session can actually do, so it never invents a read-only jail.""" + if yolo or auto: + return ("\n\n## Session capabilities\nApprovals are off and every tool is callable, " + "including shell and file writes. This session is not read-only: only a failed " + "tool call proves a limit, so act first and adjust on evidence.") + return ("\n\n## Session capabilities\nEvery tool is callable, including shell and file writes; " + "some commands ask for approval first. This session is not read-only: only a failed " + "tool call proves a limit, so act first and adjust on evidence.") + + def split_leading_tags(text): tags = [] rest = text or "" @@ -1032,13 +1044,21 @@ class Spinner: def spin(self): pos = 0 + started = time.time() + glow = Glow.stderr_enabled() + widest = len(self.label) + 2 while not self.stop_flag.is_set(): frame = self.FRAMES[pos % len(self.FRAMES)] - sys.stderr.write("\r%s %s" % (frame, self.label)) + if glow: + clock = "· %.1fs" % (time.time() - started) + sys.stderr.write("\r%s%s%s %s %s%s%s" % (Ansi.BOLD, Ansi.CYAN, frame, Glow.shimmer(self.label, pos), Ansi.DIM, clock, Ansi.RESET)) + widest = max(widest, len(self.label) + len(clock) + 3) + else: + sys.stderr.write("\r%s %s" % (frame, self.label)) sys.stderr.flush() pos += 1 time.sleep(0.08) - sys.stderr.write("\r%s\r" % (" " * (len(self.label) + 2))) + sys.stderr.write("\r%s\r" % (" " * widest)) sys.stderr.flush() def stop(self): @@ -1175,11 +1195,91 @@ def run_live(argv, timeout=SHELL_TIMEOUT, shell=False, cwd=None, env=None, strea return proc.returncode, "".join(chunks), time.time() - started, timed_out +class LaneBoard: + """Live lane block for poll waits: ticking elapsed, pulsing dots. + + Prints K lines once, then rewrites them in place at FPS until stop() + leaves a steady final frame. Rows only flip running to finished + mid-wait; if the row count ever changes the board degrades to fresh + static lines instead of mis-drawing over scrollback. + """ + + FPS = 3 + PULSE = (Ansi.DIM, "", Ansi.BOLD, "") + + def __init__(self, profile, width=0, cap=6, file=None): + self.profile = profile + self.width = width or Glow.width() + self.cap = cap + self.file = file + self.stop_flag = threading.Event() + self.worker = None + self.painted = 0 + + @staticmethod + def available(app): + return not getattr(app, "quiet", False) and Glow.enabled() + + def rows(self, pulse=""): + return Glow.lanes(list_agents(self.profile), self.cap, self.width, pulse) + + def _paint(self, pulse): + out = self.file if self.file is not None else sys.stdout + rows = self.rows(pulse) + if not rows: + return + if self.painted and self.painted != len(rows): + for row in rows: + out.write("%s\n" % row) + out.flush() + self.painted = len(rows) + return + if self.painted: + out.write("\033[%dA" % self.painted) + for row in rows: + out.write("\r\033[K%s\n" % row) + out.flush() + self.painted = len(rows) + + def _tick(self): + frame = 1 + while not self.stop_flag.is_set(): + time.sleep(1.0 / self.FPS) + if self.stop_flag.is_set(): + break + try: + self._paint(self.PULSE[frame % len(self.PULSE)]) + except Exception: + break + frame += 1 + + def start(self): + if self.worker is not None: + return + self._paint("") + self.worker = threading.Thread(target=self._tick, daemon=True) + self.worker.start() + + def stop(self): + self.stop_flag.set() + if self.worker is not None: + self.worker.join(timeout=2) + self.worker = None + self._paint("") + + DIFF_CONTEXT = 3 DIFF_MAX_LINES = 120 -BG_ADD = "\033[48;5;28m" -BG_DEL = "\033[48;5;88m" +BG_ADD = "\033[48;5;236m" +BG_DEL = "\033[48;5;236m" +FG_ADD = "\033[38;5;222m" +FG_DEL = "\033[38;5;203m" +FG_FILE = "\033[38;5;111m" +FG_BRIGHT = "\033[97m" +FG_YOLO = "\033[38;5;211m" FG_ORANGE = "\033[38;5;214m" +FG_RUN = "\033[38;5;80m" +FG_FORK = "\033[38;5;135m" PY_KEYWORDS = ("False", "None", "True", "and", "as", "assert", "async", "await", "break", "class", "continue", "def", "del", "elif", "else", "except", "finally", "for", "from", "global", "if", "import", "in", "is", "lambda", "nonlocal", "not", "or", "pass", "raise", "return", "try", "while", "with", "yield", "match", "case") @@ -1254,17 +1354,20 @@ def render_diff(path, old, new, width=0, colors=True): added = sum(1 for kind, _a, _b, _t in rows if kind == "add") removed = sum(1 for kind, _a, _b, _t in rows if kind == "del") if new is None: - head = "── %s (deleted, %d lines)" % (path, len(old_lines)) + verb, meta = "Deleted", "(deleted, %d lines)" % len(old_lines) elif not old_lines: - head = "── %s (new file, %d lines)" % (path, len(new_lines)) + verb, meta = "Created", "(new file, %d lines)" % len(new_lines) else: - head = "── %s (+%d -%d)" % (path, added, removed) + verb, meta = "Edited", "(+%d -%d)" % (added, removed) python = str(path).lower().endswith(".py") num_width = max(4, len(str(max(len(old_lines), len(new_lines), 1)))) code_width = max(20, width - num_width - 4) fence = None out = [] - out.append(Ansi.DIM + "│ " + head + Ansi.RESET if colors else "│ " + head) + if colors: + out.append(Glow.event_row("◆", verb, str(path), meta, width)) + else: + out.append("│ ── %s %s" % (path, meta)) for kind, old_no, new_no, text in rows[:DIFF_MAX_LINES]: if kind == "gap": out.append(Ansi.DIM + "···" + Ansi.RESET if colors else "···") @@ -1294,10 +1397,10 @@ def render_diff(path, old, new, width=0, colors=True): out.append("%s %s" % (dimmed, truncate_ansi(code, code_width))) elif kind == "del": row = "%s %s %s" % (dimmed, Ansi.RED + Ansi.BOLD + "-" + Ansi.RESET, truncate_ansi(code, code_width)) - out.append(pad_tinted(style_span(row, BG_DEL), width)) + out.append(pad_tinted(style_span(row, BG_DEL, FG_DEL), width)) else: - row = "%s %s %s" % (dimmed, Ansi.GREEN + Ansi.BOLD + "+" + Ansi.RESET, truncate_ansi(code, code_width)) - out.append(pad_tinted(style_span(row, BG_ADD), width)) + row = "%s %s %s" % (dimmed, FG_ADD + Ansi.BOLD + "+" + Ansi.RESET, truncate_ansi(code, code_width)) + out.append(pad_tinted(style_span(row, BG_ADD, FG_ADD), width)) if len(rows) > DIFF_MAX_LINES: note = "··· %d more lines hidden" % (len(rows) - DIFF_MAX_LINES) out.append(Ansi.DIM + note + Ansi.RESET if colors else note) @@ -1579,6 +1682,293 @@ def render_markdown(text, width=None, color=None): return "\n".join(out) +class Glow: + """Screenshot-style CLI event renderer: bold verbs, blue targets, dim meta. + + Every styling method always emits codes; call sites gate on enabled() + so pipes, NO_COLOR, and tests keep the plain format byte for byte. + """ + + @staticmethod + def enabled(): + return color_enabled() + + @staticmethod + def stderr_enabled(): + return not os.environ.get("NO_COLOR") and sys.stderr.isatty() + + @staticmethod + def width(fallback=80): + try: + return shutil.get_terminal_size().columns + except OSError: + return fallback + + @staticmethod + def accent_for(name): + """Kind color for the card bar and open icon. Small surface on purpose.""" + if name == "shell": + return FG_RUN + if name == "fork": + return FG_FORK + if name == "poll": + return FG_FILE + if name in ("write_file", "edit_file", "delete_file"): + return Ansi.GREEN + if name == "ask_user": + return FG_ORANGE + if name in ("store_secret", "list_secrets", "delete_secret"): + return FG_YOLO + return Ansi.DIM + + @staticmethod + def bar(name): + return "%s▌%s " % (Glow.accent_for(name), Ansi.RESET) + + @staticmethod + def event_row(icon, verb, target, meta="", width=0, icon_color=""): + row = "%s%s%s%s %s%s%s" % (Ansi.BOLD, icon_color, icon, Ansi.RESET, Ansi.BOLD, verb, Ansi.RESET) + if target: + row += " %s%s%s" % (FG_FILE, target, Ansi.RESET) + if meta: + row += " %s%s%s" % (Ansi.DIM, meta, Ansi.RESET) + return truncate_ansi(row, width) if width else row + + @staticmethod + def preview_row(text, width=0): + row = "%s│ %s%s" % (Ansi.DIM, text, Ansi.RESET) + return truncate_ansi(row, width) if width else row + + @staticmethod + def icon_for(name): + """Small open-row icon set: execute, spawn, poll, everything else.""" + if name == "shell": + return "▶" + if name == "fork": + return "✦" + if name == "poll": + return "⟳" + return "◇" + + @staticmethod + def ellipsis(text, cap=60, head=False): + """Truncate long targets: middle cut by default, head cut for sentences + whose discriminating words sit at the front (agent tasks).""" + if len(text or "") <= cap or cap < 8: + return text + if head: + return text[:cap - 1] + "…" + keep = cap - 1 + front = keep // 2 + return text[:front] + "…" + text[len(text) - (keep - front):] + + @staticmethod + def done_row(verb, target, chars, elapsed_ms, failed=False, width=0): + elapsed = "%.1fs" % (elapsed_ms / 1000.0) if elapsed_ms >= 1000 else "%dms" % elapsed_ms + mark = Ansi.RED + "✗" + Ansi.RESET if failed else Ansi.GREEN + "✓" + Ansi.RESET + tail = "%s· %s · %d chars%s" % (Ansi.DIM, elapsed, chars, Ansi.RESET) + left = "%s◆%s %s%s%s" % (Ansi.BOLD, Ansi.RESET, Ansi.BOLD, verb, Ansi.RESET) + if target: + left += " %s%s%s" % (FG_FILE, target, Ansi.RESET) + if width: + plain_left, plain_right = len(strip_ansi(left)), len(strip_ansi(mark + " " + tail)) + if plain_left + 2 + plain_right <= width: + return left + " " * (width - plain_left - plain_right) + mark + " " + tail + return truncate_ansi(left + " " + Ansi.DIM + "·" + Ansi.RESET + " " + mark + " " + tail, width) + return left + " " + Ansi.DIM + "·" + Ansi.RESET + " " + mark + " " + tail + + @staticmethod + def hairline(width=0): + width = width or Glow.width() + return "%s%s%s" % (Ansi.DIM, "─" * max(1, width), Ansi.RESET) + + @staticmethod + def shimmer(text, pos, span=3): + """Traveling bright window over dim text: the live-line light effect.""" + if not text: + return "" + cycle = len(text) + span * 2 + center = (pos % cycle) - span + out = [] + for pos_, char in enumerate(text): + dist = abs(pos_ - center) + if dist == 0: + out.append(Ansi.BOLD + FG_BRIGHT + char) + elif dist <= span: + out.append(Ansi.BOLD + char) + else: + out.append(Ansi.DIM + char) + return "".join(out) + Ansi.RESET + + @staticmethod + def describe_call(name, raw_arguments): + """Map a tool call to (verb, target, extra, preview) for event rows.""" + try: + args = raw_arguments if isinstance(raw_arguments, dict) else json.loads(raw_arguments or "{}") + parsed = True + except (TypeError, ValueError): + args = {} + parsed = False + if not isinstance(args, dict): + args = {} + parsed = False + if name == "store_secret" and "value" in args: + args = dict(args) + args["value"] = "[hidden]" + + def head(key, cap=80): + value = args.get(key) + if value is None: + return "" + return " ".join(str(value).split())[:cap] + + def base(key): + value = head(key) + return os.path.basename(value.rstrip("/")) or value + + path = base("path") or base("file") + command = head("command") + query = head("query") + url = head("url") + prompt = head("prompt", 60) + task = head("task", 60) + who = head("name") + ident = head("id") or head("target") or head("agent_id") + when = head("at") or ("every %s" % head("every") if head("every") else "") + table = { + "shell": ("Run", command, ""), + "read_file": ("Read", path, ""), + "write_file": ("Write", path, ""), + "edit_file": ("Edit", path, ""), + "delete_file": ("Delete", path, ""), + "web_search": ("Search", query, ""), + "web_fetch": ("Fetch", url, ""), + "describe_image": ("See", path or head("image"), ""), + "speak": ("Speak", head("text", 60), ""), + "listen": ("Listen", head("seconds"), ""), + "remember": ("Remember", head("instruction", 60), ""), + "recall": ("Recall", query, ""), + "load_skill": ("Skill", who, ""), + "create_skill": ("Skill", who, ""), + "install_skill": ("Install", who or head("source"), ""), + "get_current_terminal_content": ("Capture", "terminal", ""), + "list_terminals": ("Terminals", "", ""), + "tail_terminals": ("Tail", "terminals", ""), + "export_terminal_log": ("Export", "terminal log", ""), + "plexer_send": ("Send", ident, ""), + "plexer_read": ("Read", ident, ""), + "plexer_type": ("Type", ident, ""), + "plexer_status": ("Plexer", "", ""), + "fork": ("Fork", task, ""), + "poll": ("Poll", ident, ""), + "ask_user": ("Ask", head("question", 60), ""), + "sysinfo": ("Probe", "", ""), + "store_secret": ("Seal", who, ""), + "list_secrets": ("Seal", "list", ""), + "delete_secret": ("Delete", who, ""), + "schedule": ("Schedule", prompt, when), + "unschedule": ("Unschedule", ident, ""), + "schedules": ("Schedules", "", ""), + "record_save": ("Save", head("title"), ""), + "record_read": ("Read", ident, ""), + "record_search": ("Search", query, ""), + "record_delete": ("Delete", ident, ""), + "graph_link": ("Link", "", ""), + "graph_query": ("Graph", head("node"), ""), + "audit": ("Audit", path, ""), + "restore": ("Restore", ident, ""), + "release": ("Release", head("part"), head("message", 60)), + "tags": ("Tags", head("prefix"), ""), + "install": ("Install", head("component"), ""), + "create_bot": ("Bot", who, ""), + "search": ("Search", query, ""), + } + primary = {"shell": "command", "read_file": "path", "write_file": "path", "edit_file": "path", + "delete_file": "path", "web_search": "query", "web_fetch": "url", "describe_image": "path", + "speak": "text", "remember": "instruction", "recall": "query", "load_skill": "name", + "create_skill": "name", "install_skill": "source", "fork": "task", "ask_user": "question", + "delete_secret": "name", "schedule": "prompt", "unschedule": "id", "record_save": "title", + "record_read": "id", "record_search": "query", "record_delete": "id", "graph_query": "node", + "audit": "path", "restore": "id", "release": "part", "tags": "prefix", + "install": "component", "create_bot": "name", "search": "query"} + if name not in table: + return name, "", "", " ".join(tool_args_text(raw_arguments).split()) + verb, target, extra = table[name] + bits = ([extra] if extra else []) + [name] + if name == "poll" and ident: + preview = "agent %s" % ident + elif name in primary: + preview = head(primary[name], 200) + else: + preview = "" + if not preview: + preview = " ".join(tool_args_text(args if parsed else raw_arguments).split()) + return verb, Glow.ellipsis(target, 60), " · ".join(bits), preview + + @staticmethod + def statusbar(version, profile, directory, yolo=False, agents=0, width=0): + home = os.path.expanduser("~") + short = directory + if directory == home: + short = "~" + elif home and directory.startswith(home + os.sep): + short = "~" + directory[len(home):] + segs = [ + "%s%stai %s%s" % (Ansi.BOLD, FG_FILE, version, Ansi.RESET), + "%s%s%s" % (FG_FILE, profile, Ansi.RESET), + "%s%s%s" % (Ansi.DIM, short, Ansi.RESET), + ] + if yolo: + segs.append("%s%sYOLO%s" % (Ansi.BOLD, FG_YOLO, Ansi.RESET)) + if agents: + segs.append("%s+%d%s" % (Ansi.YELLOW, agents, Ansi.RESET)) + dot = "%s·%s" % (Ansi.DIM, Ansi.RESET) + width = width or Glow.width() + flow = dot.join(" %s " % seg for seg in segs).strip() + if (yolo or agents) and len(segs) > 3: + left = dot.join(" %s " % seg for seg in segs[:3]).strip() + right = dot.join(" %s " % seg for seg in segs[3:]).strip() + gap = width - len(strip_ansi(left)) - len(strip_ansi(right)) + if gap >= 4: + return left + " " * gap + right + flow = dot.join(" %s " % seg for seg in [segs[0], segs[1]] + segs[3:]).strip() + if len(strip_ansi(flow)) <= width: + return flow + core = truncate_ansi(flow, width) + return re.sub(r"(?:\033\[[0-9;]*m|[ ·])+$", "", core) + Ansi.RESET + + @staticmethod + def prompt_text(): + return "%s%s› %s" % (Ansi.BOLD, FG_ORANGE, Ansi.RESET) + + @staticmethod + def agent_row(record, width=0, pulse="", suffix=""): + dots = {"running": Ansi.YELLOW, "done": Ansi.GREEN, "error": Ansi.RED, "timeout": FG_ORANGE} + dot = dots.get(record.get("status"), Ansi.DIM) + task = Glow.ellipsis(" ".join(str(record.get("task") or "").split()), 50, head=True) + row = " #%d %s%s●%s %s %s· %s · %s%s" % (record.get("id"), pulse, dot, Ansi.RESET, task, Ansi.DIM, + record.get("status"), format_delay(record.get("elapsed") or 0), Ansi.RESET) + if suffix: + row += suffix + return truncate_ansi(row, width) if width else row + + @staticmethod + def lanes(records, cap=6, width=0, pulse=""): + running = [record for record in records if record["status"] == "running"] + pending = [record for record in records if record["status"] in ("done", "timeout", "error") and not record.get("collected")] + view = running + pending + head = "◇ lanes · %d running" % len(running) + if pending: + head += " · %d to collect" % len(pending) + rows = [truncate_ansi("%s%s%s" % (Ansi.DIM, head, Ansi.RESET), width) if width else "%s%s%s" % (Ansi.DIM, head, Ansi.RESET)] + for record in view[:cap]: + suffix = " %s· to collect%s" % (Ansi.DIM, Ansi.RESET) if record in pending else "" + rows.append(Glow.agent_row(record, width, pulse if record["status"] == "running" else "", suffix)) + if len(view) > cap: + rows.append("%s · +%d more%s" % (Ansi.DIM, len(view) - cap, Ansi.RESET)) + return rows + + class SealError(Exception): pass @@ -5930,6 +6320,7 @@ class Tools: self.read_files = set() self.live = False self.diff_preview = None + self.board_shown = False def note_diff(self, path, old, new): self.diff_preview = {"path": path, "old": old, "new": new} @@ -6629,7 +7020,21 @@ class Tools: wait = max(0, min(120, int(args.get("wait") or 0))) except (TypeError, ValueError): wait = 0 - status, text = poll_agent(agent_id, wait, self.active_profile) + status, text = poll_agent(agent_id, 0, self.active_profile) + if status == "running" and wait > 0: + board = LaneBoard(self.active_profile) if LaneBoard.available(self.app) else None + if board is not None: + board.start() + try: + status, text = poll_agent(agent_id, wait, self.active_profile) + finally: + if board is not None: + board.stop() + self.board_shown = True + if status in ("done", "timeout", "error"): + with AGENTS_LOCK: + if agent_id in AGENTS: + AGENTS[agent_id]["collected"] = True if status in ("done", "timeout") and self.app.store is not None: with AGENTS_LOCK: record = dict(AGENTS.get(agent_id) or {}) @@ -7474,11 +7879,13 @@ class Agent: self._turn_goal = "" self._turn_steps = 0 self._turn_tools = 0 + self._last_event = None self.switch_profile(config.profile, silent=True) def apply_system(self): content = self.system_message + skill_catalog(self.skills) + tool_catalog() content += orientation_block(sandbox=self.env == "sandbox") + content += capability_block(self.yolo, self.auto) if self.auto: content += "\n\n" + AUTONOMOUS_NOTE self.messages[0] = {"role": "system", "content": content} @@ -7588,7 +7995,6 @@ class Agent: def show_call(self, call): if self.quiet: return - print(paint("┌─ %s" % call["name"], Ansi.MAGENTA)) preview_args = tool_args_text(call["arguments"]) if call["name"] == "store_secret": try: @@ -7597,7 +8003,19 @@ class Agent: preview_args = json.dumps(masked) except ValueError: preview_args = "[hidden]" + if Glow.enabled(): + verb, target, extra, pretty = Glow.describe_call(call["name"], preview_args) + target = self.store.redact(target, self.profile) + extra = self.store.redact(extra, self.profile) + preview = self.store.redact(pretty, self.profile)[:400].replace("\n", " ") + self._last_event = (call["name"], verb, target) + print() + inner = max(20, Glow.width() - 2) + print(Glow.bar(call["name"]) + Glow.event_row(Glow.icon_for(call["name"]), verb, target, extra, inner, Glow.accent_for(call["name"]))) + print(Glow.bar(call["name"]) + Glow.preview_row(preview, inner)) + return preview = self.store.redact(preview_args, self.profile)[:400].replace("\n", " ") + print(paint("┌─ %s" % call["name"], Ansi.MAGENTA)) print(paint("│ %s" % preview, Ansi.DIM)) def show_result(self, result, elapsed_ms): @@ -7608,8 +8026,41 @@ class Agent: if note: self.tools.result_note = None print(paint(note, Ansi.DIM)) + if Glow.enabled(): + name, verb, target = self._last_event or ("", "Done", "") + failed = result == "denied by user" or (result or "").startswith(("error:", "[error]", "[missing]")) + if name == "fork": + match = re.match(r"agent (\d+) started", result or "") + if match: + target = "%s · #%s" % (target, match.group(1)) if target else "#%s" % match.group(1) + inner = max(20, Glow.width() - 2) + print(Glow.bar(name) + Glow.done_row(verb, target, len(result), elapsed_ms, failed, inner)) + self._last_event = None + self.show_lanes(name, result or "") + return print(paint("└─ %d chars · %dms" % (len(result), elapsed_ms), Ansi.DIM)) + def show_lanes(self, name, result): + boarded = getattr(self.tools, "board_shown", False) + self.tools.board_shown = False + started = name == "fork" and result.startswith("agent ") + collected = name == "poll" and result.startswith(("[done]", "[timeout]")) + watching = name == "poll" and result.startswith("[running]") + if not (started or collected or watching): + return + records = list_agents(self.profile) + running = [record for record in records if record["status"] == "running"] + pending = [record for record in records if record["status"] in ("done", "timeout", "error") and not record.get("collected")] + if boarded: + if collected and not running and not pending: + print(Glow.event_row("◆", "lanes", "", "all collected")) + return + if len(running) >= 2 or pending: + for row in Glow.lanes(records): + print(row) + elif collected and not running and not pending: + print(Glow.event_row("◆", "lanes", "", "all collected")) + def show_diff(self): preview = getattr(self.tools, "diff_preview", None) if preview is None: @@ -7706,6 +8157,7 @@ class Agent: return prompt, None def run_turn(self, text, capture=False, max_steps=MAX_STEPS, _routed=False): + self.apply_system() if not _routed: text, result = self.apply_mentions(text, capture, max_steps) if result is not None: @@ -7851,6 +8303,8 @@ class Agent: self.close_turn() total = time.time() - started if loud: + if Glow.enabled(): + print(Glow.hairline()) print(paint("tokens≈%d · %s · %.1fs" % (conversation_tokens(self.messages), "+".join(backends), total), Ansi.DIM)) return last_text @@ -7899,7 +8353,7 @@ def spawn_agent(task, profile, timeout, config, seal, depth=0, runner=None): with AGENTS_LOCK: agent_id = AGENTS_NEXT[0] AGENTS_NEXT[0] += 1 - record = {"id": agent_id, "task": task, "profile": profile, "status": "running", "result": "", "started": time.time(), "ended": None, "thread": None, "spilled": None} + record = {"id": agent_id, "task": task, "profile": profile, "status": "running", "result": "", "started": time.time(), "ended": None, "thread": None, "spilled": None, "collected": False} AGENTS[agent_id] = record def target(): @@ -7947,7 +8401,7 @@ def poll_agent(agent_id, wait=0, profile=None): def list_agents(profile=None): with AGENTS_LOCK: - return [{"id": record["id"], "task": record["task"], "status": record["status"], "elapsed": int((record["ended"] or time.time()) - record["started"])} for record in AGENTS.values() if profile is None or record.get("profile") == profile] + return [{"id": record["id"], "task": record["task"], "status": record["status"], "elapsed": int((record["ended"] or time.time()) - record["started"]), "collected": record.get("collected", False)} for record in AGENTS.values() if profile is None or record.get("profile") == profile] def clear_agents(profile=None): @@ -8364,8 +8818,16 @@ def handle_command(agent, text): records = list_agents(agent.profile) if not records: print("(no agents)") - for record in records: - print("#%d [%s] %ds %s" % (record["id"], record["status"], record["elapsed"], record["task"][:60])) + elif Glow.enabled(): + running = sum(1 for record in records if record["status"] == "running") + print(Glow.event_row("◆", "agents", "%d total" % len(records), "%d running" % running)) + for record in records[:50]: + print(Glow.agent_row(record)) + if len(records) > 50: + print(paint(" · +%d more" % (len(records) - 50), Ansi.DIM)) + else: + for record in records: + print("#%d [%s] %ds %s" % (record["id"], record["status"], record["elapsed"], record["task"][:60])) elif name == "agent": if arg == "clear": print("purged %d finished agents" % clear_agents(agent.profile)) @@ -13308,7 +13770,10 @@ def build_parser(): def banner(agent): - print(paint("tai %s · profile %s" % (VERSION, agent.profile), Ansi.BOLD, Ansi.CYAN)) + if Glow.enabled(): + print("%s◆%s %stai%s %s%s%s %s· profile %s%s" % (Ansi.BOLD, Ansi.RESET, Ansi.BOLD, Ansi.RESET, FG_FILE, VERSION, Ansi.RESET, Ansi.DIM, agent.profile, Ansi.RESET)) + else: + print(paint("tai %s · profile %s" % (VERSION, agent.profile), Ansi.BOLD, Ansi.CYAN)) if agent.yolo: print(paint("YOLO MODE: all approvals off", Ansi.BOLD, Ansi.RED)) if agent.auto: @@ -13316,9 +13781,15 @@ def banner(agent): seal_state = "off" if agent.store.seal.enabled: seal_state = "on (default key)" if agent.store.seal.default_key else "on (TAI_PASSPHRASE)" - print(" seal: %s" % seal_state) + if Glow.enabled(): + print(" %sseal:%s %s" % (Ansi.DIM, Ansi.RESET, seal_state)) + else: + print(" seal: %s" % seal_state) for label, _chat_url, models_url, key, _model in agent.chat.backends(): - print(" %s: %s" % (label, probe_backend(models_url, key))) + 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"): print(paint("interrupted turn available, restart with --continue to pick it up", Ansi.YELLOW)) print(paint("type /help for commands", Ansi.DIM)) @@ -13350,7 +13821,11 @@ def repl(agent): busy = running_agents() base = agent.profile + ("!" if agent.yolo else "") tag = "%s+%d" % (base, busy) if busy else base - line = input(paint("tai[%s]› " % tag, Ansi.BOLD, Ansi.GREEN)) + if Glow.enabled(): + print(Glow.statusbar(VERSION, agent.profile, dir_key(), agent.yolo, busy)) + line = input(Glow.prompt_text()) + else: + line = input(paint("tai[%s]› " % tag, Ansi.BOLD, Ansi.GREEN)) except EOFError: print() break diff --git a/test_tai.py b/test_tai.py index 94da4af..c44a980 100644 --- a/test_tai.py +++ b/test_tai.py @@ -2536,6 +2536,53 @@ class OrientationTests(unittest.TestCase): self.assertIn("never invoke assistant or harness tools yourself", prompt) +class CapabilityTests(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + os.environ["TAI_HOME"] = self.tmp.name + + class FakeArgs: + profile = "t" + yes = True + + self.config = tai.Config(FakeArgs()) + self.store = tai.Store(self.config, tai.Seal(self.config.home, "cap-test-1")) + self.agent = tai.Agent(self.config, self.store, persist=True, quiet=True) + + def tearDown(self): + self.store.close() + self.tmp.cleanup() + os.environ.pop("TAI_HOME", None) + + def test_yolo_variant(self): + block = tai.capability_block(yolo=True) + self.assertIn("Approvals are off", block) + self.assertIn("not read-only", block) + self.assertIn("only a failed tool call proves a limit", block) + + def test_default_variant(self): + block = tai.capability_block() + self.assertIn("ask for approval", block) + self.assertIn("not read-only", block) + + def test_apply_system_injects_capabilities(self): + system = self.agent.messages[0]["content"] + self.assertIn("## Session capabilities", system) + self.assertIn("ask for approval", system) + self.agent.yolo = True + self.agent.apply_system() + self.assertIn("Approvals are off", self.agent.messages[0]["content"]) + + def test_run_turn_refreshes_capabilities(self): + self.agent.messages[0] = {"role": "system", "content": "stale"} + reply = {"role": "assistant", "content": "done", "reasoning": "", "tool_calls": [], "backend": "x"} + with mock.patch.object(self.agent.chat, "complete", return_value=dict(reply)): + self.agent.run_turn("hi", capture=True) + system = self.agent.messages[0]["content"] + self.assertIn("## Session capabilities", system) + self.assertNotIn("stale", system) + + class FTSSearchTests(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() @@ -3197,6 +3244,342 @@ class DiffTests(unittest.TestCase): shown.assert_not_called() +class GlowRenderingTests(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + os.environ["TAI_HOME"] = self.tmp.name + + class FakeArgs: + profile = "t" + yes = True + + self.config = tai.Config(FakeArgs()) + self.store = tai.Store(self.config, tai.Seal(self.config.home, "glow-test-1")) + + def tearDown(self): + self.store.close() + self.tmp.cleanup() + os.environ.pop("TAI_HOME", None) + + def test_describe_call_verbs(self): + self.assertEqual(tai.Glow.describe_call("shell", json.dumps({"command": "make test\nsecond"})), + ("Run", "make test second", "shell", "make test second")) + self.assertEqual(tai.Glow.describe_call("read_file", json.dumps({"path": "/tmp/x/main.py"})), + ("Read", "main.py", "read_file", "/tmp/x/main.py")) + self.assertEqual(tai.Glow.describe_call("shell", {"command": "ls"}), ("Run", "ls", "shell", "ls")) + self.assertEqual(tai.Glow.describe_call("shell", "nope{"), ("Run", "", "shell", "nope{")) + self.assertEqual(tai.Glow.describe_call("mystery", "{}"), ("mystery", "", "", "{}")) + self.assertEqual(tai.Glow.describe_call("poll", {"id": 3}), ("Poll", "3", "poll", "agent 3")) + self.assertEqual(tai.Glow.describe_call("write_file", json.dumps({"path": "n.txt", "content": "x" * 500}))[3], "n.txt") + verb, target, _extra, preview = tai.Glow.describe_call("store_secret", json.dumps({"name": "api", "value": "s3cr3t"})) + self.assertEqual(verb, "Seal") + self.assertNotIn("s3cr3t", target) + self.assertNotIn("s3cr3t", preview) + self.assertIn("[hidden]", preview) + + def test_shimmer_sweeps(self): + out = tai.Glow.shimmer("thinking", 2) + self.assertIn(tai.Ansi.BOLD, out) + self.assertIn(tai.Ansi.DIM, out) + self.assertEqual(tai.strip_ansi(out), "thinking") + self.assertEqual(tai.strip_ansi(tai.Glow.shimmer("thinking", 99)), "thinking") + self.assertEqual(tai.Glow.shimmer("", 5), "") + + def test_event_and_done_rows(self): + row = tai.Glow.event_row("◆", "Edited", "a.py", "(+1 -1)") + self.assertIn("◆", row) + self.assertIn(tai.FG_FILE, row) + self.assertEqual(tai.strip_ansi(row), "◆ Edited a.py (+1 -1)") + done = tai.Glow.done_row("Run", "make test", 12, 234) + self.assertIn("✓", done) + self.assertIn("234ms", done) + self.assertIn("12 chars", done) + failed = tai.Glow.done_row("Run", "", 5, 1500, failed=True) + self.assertIn("✗", failed) + self.assertIn("1.5s", failed) + short = tai.Glow.event_row("◆", "Edited", "a.py", "(+1 -1)", width=10) + self.assertEqual(len(tai.strip_ansi(short)), 10) + + def test_kind_icons_and_ellipsis(self): + self.assertEqual(tai.Glow.icon_for("shell"), "▶") + self.assertEqual(tai.Glow.icon_for("fork"), "✦") + self.assertEqual(tai.Glow.icon_for("poll"), "⟳") + self.assertEqual(tai.Glow.icon_for("read_file"), "◇") + self.assertEqual(tai.Glow.icon_for("mystery"), "◇") + self.assertEqual(tai.Glow.ellipsis("short", 60), "short") + long = "a" * 30 + "/" + "b" * 30 + cut = tai.Glow.ellipsis(long, 20) + self.assertEqual(len(cut), 20) + self.assertIn("…", cut) + self.assertTrue(cut.startswith("a" * 9)) + self.assertTrue(cut.endswith("b" * 10)) + self.assertEqual(tai.Glow.ellipsis(long, 4), long) + self.assertEqual(tai.Glow.ellipsis("a" * 60, 20, head=True), "a" * 19 + "…") + + def test_accents_and_bars(self): + self.assertEqual(tai.Glow.accent_for("shell"), tai.FG_RUN) + self.assertEqual(tai.Glow.accent_for("fork"), tai.FG_FORK) + self.assertEqual(tai.Glow.accent_for("poll"), tai.FG_FILE) + self.assertEqual(tai.Glow.accent_for("edit_file"), tai.Ansi.GREEN) + self.assertEqual(tai.Glow.accent_for("ask_user"), tai.FG_ORANGE) + self.assertEqual(tai.Glow.accent_for("store_secret"), tai.FG_YOLO) + self.assertEqual(tai.Glow.accent_for("read_file"), tai.Ansi.DIM) + bar = tai.Glow.bar("shell") + self.assertIn("▌", bar) + self.assertIn(tai.FG_RUN, bar) + tinted = tai.Glow.event_row("▶", "Run", "x", "", 0, tai.FG_RUN) + self.assertIn(tai.Ansi.BOLD + tai.FG_RUN + "▶", tinted) + self.assertEqual(tai.strip_ansi(tinted), "▶ Run x") + + def test_done_row_right_aligns_meta(self): + row = tai.Glow.done_row("Run", "x", 12, 234, width=60) + plain = tai.strip_ansi(row) + self.assertEqual(len(plain), 60) + self.assertTrue(plain.startswith("◆ Run x")) + self.assertTrue(plain.endswith("✓ · 234ms · 12 chars")) + narrow = tai.Glow.done_row("Run", "averylongtargetname", 12, 234, width=10) + self.assertEqual(len(tai.strip_ansi(narrow)), 10) + + def test_agent_row_and_lanes(self): + row = tai.Glow.agent_row({"id": 3, "task": "translate p1 to german", "status": "running", "elapsed": 75}) + self.assertIn("#3", row) + self.assertIn(tai.Ansi.YELLOW + "●" + tai.Ansi.RESET, row) + self.assertIn("translate p1 to german", tai.strip_ansi(row)) + self.assertIn("1m", tai.strip_ansi(row)) + done = tai.Glow.agent_row({"id": 4, "task": "t", "status": "done", "elapsed": 3}) + self.assertIn(tai.Ansi.GREEN + "●" + tai.Ansi.RESET, done) + many = [{"id": num, "task": "t%d" % num, "status": "running", "elapsed": 1} for num in range(8)] + rows = tai.Glow.lanes(many) + self.assertEqual(len(rows), 1 + 6 + 1) + self.assertIn("8 running", tai.strip_ansi(rows[0])) + self.assertIn("+2 more", tai.strip_ansi(rows[-1])) + + def test_lane_tasks_keep_distinguishing_head(self): + task = "Translate README.md lines 94-186 to /tmp/xxx_p2.md. Rules: " + "x" * 80 + row = tai.strip_ansi(tai.Glow.agent_row({"id": 1, "task": task, "status": "running", "elapsed": 0})) + self.assertIn("lines 94-186", row) + self.assertIn("…", row) + wide = tai.Glow.agent_row({"id": 1, "task": task, "status": "running", "elapsed": 0}, width=40) + self.assertLessEqual(len(tai.strip_ansi(wide)), 40) + pulsed = tai.Glow.agent_row({"id": 1, "task": "t", "status": "running", "elapsed": 0}, 0, tai.Ansi.BOLD) + self.assertIn(tai.Ansi.BOLD + tai.Ansi.YELLOW + "●", pulsed) + + def test_lanes_show_pending_pickups(self): + records = [{"id": 1, "task": "a", "status": "running", "elapsed": 5, "collected": False}, + {"id": 2, "task": "b", "status": "done", "elapsed": 9, "collected": False}, + {"id": 3, "task": "c", "status": "done", "elapsed": 9, "collected": True}] + rows = tai.Glow.lanes(records) + self.assertIn("1 running · 1 to collect", tai.strip_ansi(rows[0])) + body = "\n".join(tai.strip_ansi(row) for row in rows[1:]) + self.assertIn("#1", body) + self.assertIn("#2", body) + self.assertNotIn("#3", body) + self.assertIn("to collect", body) + + def test_show_call_truncates_to_width(self): + agent = tai.Agent(self.config, self.store, persist=False, quiet=False) + with mock.patch.object(sys.stdout, "isatty", return_value=True): + with mock.patch.dict(os.environ, {"NO_COLOR": ""}): + with mock.patch("builtins.print") as shown: + agent.show_call({"name": "shell", "arguments": json.dumps({"command": "echo " + "x" * 300})}) + width = tai.Glow.width() + for call in shown.call_args_list: + if call.args: + self.assertLessEqual(len(tai.strip_ansi(str(call.args[0]))), width) + + def test_fork_done_row_carries_agent_id(self): + agent = tai.Agent(self.config, self.store, persist=False, quiet=False) + with mock.patch.object(sys.stdout, "isatty", return_value=True): + with mock.patch.dict(os.environ, {"NO_COLOR": ""}): + with mock.patch("builtins.print") as shown: + agent.show_call({"name": "fork", "arguments": json.dumps({"task": "translate p2"})}) + with mock.patch("tai.list_agents", return_value=[]): + agent.show_result("agent 7 started: translate p2", 5) + printed = "\n".join(str(call.args[0]) for call in shown.call_args_list if call.args) + self.assertIn("#7", printed) + + def test_show_lanes_snapshot_and_collect(self): + agent = tai.Agent(self.config, self.store, persist=False, quiet=False) + running = [{"id": 3, "task": "translate p1 to german", "status": "running", "elapsed": 5}, + {"id": 4, "task": "translate p2 to german", "status": "running", "elapsed": 5}] + with mock.patch.object(sys.stdout, "isatty", return_value=True): + with mock.patch.dict(os.environ, {"NO_COLOR": ""}): + with mock.patch("builtins.print") as shown: + agent.show_call({"name": "fork", "arguments": json.dumps({"task": "translate p2"})}) + with mock.patch("tai.list_agents", return_value=running): + agent.show_result("agent 4 started: translate p2", 5) + with mock.patch("builtins.print") as collected: + agent.show_call({"name": "poll", "arguments": json.dumps({"id": 3})}) + with mock.patch("tai.list_agents", return_value=[]): + agent.show_result("[done] translated text", 5) + with mock.patch("builtins.print") as quiet: + agent.show_call({"name": "poll", "arguments": json.dumps({"id": 3})}) + with mock.patch("tai.list_agents", return_value=running[:1]): + agent.show_result("[done] translated text", 5) + with mock.patch("builtins.print") as missing: + agent.show_call({"name": "poll", "arguments": json.dumps({"id": 9})}) + agent.show_result("[missing] no agent 9", 5) + snap = "\n".join(str(call.args[0]) for call in shown.call_args_list if call.args) + self.assertIn("✦", snap) + self.assertIn("lanes", snap) + self.assertIn("#3", snap) + self.assertIn("#4", snap) + done = "\n".join(str(call.args[0]) for call in collected.call_args_list if call.args) + self.assertIn("⟳", done) + self.assertIn("all collected", done) + lone = "\n".join(str(call.args[0]) for call in quiet.call_args_list if call.args) + self.assertNotIn("lanes", lone) + gone = "\n".join(str(call.args[0]) for call in missing.call_args_list if call.args) + self.assertIn("✗", gone) + + def test_show_lanes_pending_and_boarded(self): + agent = tai.Agent(self.config, self.store, persist=False, quiet=False) + leftover = [{"id": 2, "task": "translate p2", "status": "done", "elapsed": 9, "collected": False}] + picked = [{"id": 2, "task": "translate p2", "status": "done", "elapsed": 9, "collected": True}] + pair = [{"id": 3, "task": "t3", "status": "running", "elapsed": 5, "collected": False}, + {"id": 4, "task": "t4", "status": "running", "elapsed": 5, "collected": False}] + with mock.patch.object(sys.stdout, "isatty", return_value=True): + with mock.patch.dict(os.environ, {"NO_COLOR": ""}): + with mock.patch("builtins.print") as pending: + agent.show_call({"name": "poll", "arguments": json.dumps({"id": 1})}) + with mock.patch("tai.list_agents", return_value=leftover): + agent.show_result("[done] translated text", 5) + with mock.patch("builtins.print") as clean: + agent.show_call({"name": "poll", "arguments": json.dumps({"id": 2})}) + with mock.patch("tai.list_agents", return_value=picked): + agent.show_result("[done] translated text", 5) + agent.tools.board_shown = True + with mock.patch("builtins.print") as boarded: + agent.show_call({"name": "poll", "arguments": json.dumps({"id": 3})}) + with mock.patch("tai.list_agents", return_value=pair): + agent.show_result("[running] agent 3 still running (5s elapsed)", 5000) + held = "\n".join(str(call.args[0]) for call in pending.call_args_list if call.args) + self.assertIn("to collect", held) + self.assertNotIn("all collected", held) + swept = "\n".join(str(call.args[0]) for call in clean.call_args_list if call.args) + self.assertIn("all collected", swept) + live = "\n".join(str(call.args[0]) for call in boarded.call_args_list if call.args) + self.assertNotIn("lanes", live) + self.assertFalse(agent.tools.board_shown) + + def test_agents_command_lanes(self): + agent = tai.Agent(self.config, self.store, persist=False, quiet=False) + records = [{"id": 7, "task": "translate p3", "status": "running", "elapsed": 9}] + with mock.patch.object(sys.stdout, "isatty", return_value=True): + with mock.patch.dict(os.environ, {"NO_COLOR": ""}): + with mock.patch("tai.list_agents", return_value=records): + with mock.patch("builtins.print") as shown: + tai.handle_command(agent, "/agents") + printed = "\n".join(str(call.args[0]) for call in shown.call_args_list) + self.assertIn("agents", printed) + self.assertIn("#7", printed) + self.assertIn("translate p3", printed) + + def test_statusbar_segments(self): + bar = tai.Glow.statusbar("9.9.9", "work", os.path.expanduser("~/proj"), yolo=True, agents=2, width=200) + plain = tai.strip_ansi(bar) + for bit in ("tai 9.9.9", "work", "~/proj", "YOLO", "+2"): + self.assertIn(bit, plain) + self.assertIn(tai.FG_YOLO, bar) + slim = tai.strip_ansi(tai.Glow.statusbar("9.9.9", "work", "/tmp", width=200)) + self.assertNotIn("YOLO", slim) + self.assertIn("/tmp", slim) + compact = tai.strip_ansi(tai.Glow.statusbar("1.24.0", "t", "/tmp", yolo=True, agents=2, width=34)) + self.assertEqual(compact, "tai 1.24.0 · t · YOLO · +2") + tiny = tai.strip_ansi(tai.Glow.statusbar("1.24.0", "t", "/tmp", yolo=True, agents=2, width=20)) + self.assertFalse(tiny.endswith("·") or tiny.endswith(" ")) + + def test_prompt_and_hairline(self): + self.assertEqual(tai.strip_ansi(tai.Glow.prompt_text()), "› ") + self.assertIn(tai.FG_ORANGE, tai.Glow.prompt_text()) + self.assertEqual(tai.strip_ansi(tai.Glow.hairline(10)), "─" * 10) + + def test_render_diff_glow_header_and_rows(self): + out = tai.render_diff("a.py", "x = 1\n", "x = 2\n", width=60, colors=True) + self.assertTrue(tai.strip_ansi(out.splitlines()[0]).startswith("◆ Edited a.py (+1 -1)")) + self.assertIn(tai.FG_ADD, out) + self.assertIn(tai.FG_DEL, out) + + def test_show_call_result_glow_branch(self): + agent = tai.Agent(self.config, self.store, persist=False, quiet=False) + self.store.save_secret("api", "known-token-5") + with mock.patch.object(sys.stdout, "isatty", return_value=True): + with mock.patch.dict(os.environ, {"NO_COLOR": ""}): + with mock.patch("builtins.print") as shown: + agent.show_call({"name": "shell", "arguments": json.dumps({"command": "curl known-token-5"})}) + agent.show_call({"name": "read_file", "arguments": json.dumps({"path": "/tmp/x"})}) + agent.show_result("ok", 250) + with mock.patch("builtins.print") as failed: + agent.show_call({"name": "shell", "arguments": "{}"}) + agent.show_result("error: boom", 5) + printed = "\n".join(str(call.args[0]) for call in shown.call_args_list if call.args) + self.assertIn("▶", printed) + self.assertIn("◇", printed) + self.assertIn("◆", printed) + self.assertIn("▌", printed) + self.assertIn("✓", printed) + self.assertNotIn('{"command"', printed) + self.assertNotIn("known-token-5", printed) + self.assertIn("[redacted:api]", printed) + busted = "\n".join(str(call.args[0]) for call in failed.call_args_list if call.args) + self.assertIn("✗", busted) + + +class LaneBoardTests(unittest.TestCase): + def test_available(self): + loud = mock.Mock(quiet=False) + hushed = mock.Mock(quiet=True) + with mock.patch.object(sys.stdout, "isatty", return_value=False): + self.assertFalse(tai.LaneBoard.available(loud)) + with mock.patch.object(sys.stdout, "isatty", return_value=True): + with mock.patch.dict(os.environ, {"NO_COLOR": ""}): + self.assertTrue(tai.LaneBoard.available(loud)) + self.assertFalse(tai.LaneBoard.available(hushed)) + with mock.patch.dict(os.environ, {"NO_COLOR": "1"}): + self.assertFalse(tai.LaneBoard.available(loud)) + + def test_paint_rewrites_in_place(self): + records = [{"id": 1, "task": "translate p1", "status": "running", "elapsed": 5, "collected": False}] + buf = io.StringIO() + board = tai.LaneBoard("t", width=100, file=buf) + with mock.patch("tai.list_agents", return_value=records): + board._paint("") + first = buf.getvalue() + self.assertNotIn("\033[2A", first) + self.assertIn("1 running", tai.strip_ansi(first)) + board._paint(tai.Ansi.BOLD) + second = buf.getvalue() + self.assertIn("\033[2A", second) + self.assertIn(tai.Ansi.BOLD + tai.Ansi.YELLOW + "●", second) + self.assertEqual(board.painted, 2) + + def test_count_change_degrades_gracefully(self): + one = [{"id": 1, "task": "t1", "status": "running", "elapsed": 5, "collected": False}] + two = one + [{"id": 2, "task": "t2", "status": "running", "elapsed": 5, "collected": False}] + buf = io.StringIO() + board = tai.LaneBoard("t", width=100, file=buf) + with mock.patch("tai.list_agents", return_value=one): + board._paint("") + mark = len(buf.getvalue()) + with mock.patch("tai.list_agents", return_value=two): + board._paint("") + fresh = buf.getvalue()[mark:] + self.assertNotIn("[A", fresh) + self.assertIn("#2", fresh) + self.assertEqual(board.painted, 3) + + def test_start_stop_lifecycle(self): + records = [{"id": 1, "task": "translate p1", "status": "running", "elapsed": 5, "collected": False}] + buf = io.StringIO() + board = tai.LaneBoard("t", width=100, file=buf) + with mock.patch("tai.list_agents", return_value=records): + board.start() + time.sleep(0.5) + board.stop() + self.assertIsNone(board.worker) + self.assertGreaterEqual(buf.getvalue().count("lanes · 1 running"), 2) + + class ResumeTests(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() @@ -3621,6 +4004,39 @@ class OrchestrationTests(unittest.TestCase): self.assertEqual(status, "done") self.assertIn("stub-done", text) + def test_poll_tool_marks_collected(self): + tools = tai.Tools(FakePrincipal()) + first = tai.spawn_agent("quick", "t", 60, None, None, 0, lambda task, profile, timeout: "r") + self.assertIn("[done]", tools.dispatch("poll", {"id": first, "wait": 5})) + + def broken(task, profile, timeout): + raise RuntimeError("boom") + + bad = tai.spawn_agent("bad", "t", 60, None, None, 0, broken) + self.assertIn("[error]", tools.dispatch("poll", {"id": bad, "wait": 5})) + with tai.AGENTS_LOCK: + self.assertTrue(tai.AGENTS[first]["collected"]) + self.assertTrue(tai.AGENTS[bad]["collected"]) + + def test_poll_running_marks_nothing(self): + tools = tai.Tools(FakePrincipal()) + agent_id = tai.spawn_agent("slow", "t", 60, None, None, 0, lambda task, profile, timeout: time.sleep(2) or "slow-done") + self.assertIn("[running]", tools.dispatch("poll", {"id": agent_id})) + with tai.AGENTS_LOCK: + self.assertFalse(tai.AGENTS[agent_id]["collected"]) + + def test_poll_wait_runs_lane_board(self): + tools = tai.Tools(FakePrincipal()) + agent_id = tai.spawn_agent("slow", "t", 60, None, None, 0, lambda task, profile, timeout: time.sleep(2) or "slow-done") + with mock.patch.object(sys.stdout, "isatty", return_value=True): + with mock.patch.dict(os.environ, {"NO_COLOR": ""}): + with mock.patch("tai.LaneBoard") as board_cls: + result = tools.dispatch("poll", {"id": agent_id, "wait": 1}) + self.assertIn("[running]", result) + board_cls.return_value.start.assert_called_once_with() + board_cls.return_value.stop.assert_called_once_with() + self.assertTrue(tools.board_shown) + def test_poll_running_and_missing(self): def slow(task, profile, timeout): time.sleep(2)