| 1 | from __future__ import annotations |
| 2 | |
| 3 | import argparse |
| 4 | import asyncio |
| 5 | import json |
| 6 | import sys |
| 7 | from typing import Any, Iterable |
| 8 | from uuid import uuid4 |
| 9 | |
| 10 | ORIGINAL_STDOUT = sys.stdout |
| 11 | |
| 12 | |
| 13 | def event_stdout(): |
| 14 | if sys.stdout.__class__.__name__ == "_DiscardStream": |
| 15 | return ORIGINAL_STDOUT |
| 16 | return sys.stdout |
| 17 | |
| 18 | |
| 19 | def parse_args(argv: list[str] | None = None) -> argparse.Namespace: |
| 20 | parser = argparse.ArgumentParser(description="Run the ViMax agent loop.") |
| 21 | parser.add_argument("--session", default="", help="Existing session id to activate before the run starts.") |
| 22 | parser.add_argument("--new-session", action="store_true", help="Create and activate a new empty session before the run starts.") |
| 23 | parser.add_argument("--new-session-name", default="", help="Display name for a newly created session.") |
| 24 | parser.add_argument("--jsonl", action="store_true", help="Print one JSON event per line.") |
| 25 | parser.add_argument("--once", default="", help="Run a single prompt and exit. If omitted and stdin is not a TTY, stdin is consumed as one prompt.") |
| 26 | parser.add_argument("--stdin-repl", action="store_true", help=argparse.SUPPRESS) |
| 27 | return parser.parse_args(argv) |
| 28 | |
| 29 | |
| 30 | def load_runtime(): |
| 31 | from agent_runtime import build_runtime |
| 32 | |
| 33 | return build_runtime(".") |
| 34 | |
| 35 | |
| 36 | def load_session_index(): |
| 37 | from agent_runtime.session_index import SessionIndex |
| 38 | |
| 39 | return SessionIndex(".") |
| 40 | |
| 41 | |
| 42 | def print_event(event: dict[str, Any], *, jsonl: bool) -> None: |
| 43 | out = event_stdout() |
| 44 | if jsonl: |
| 45 | print(json.dumps(event, ensure_ascii=False, default=str), file=out, flush=True) |
| 46 | return |
| 47 | event_type = event.get("type") |
| 48 | if event_type == "turn": |
| 49 | print(f"· turn: {event.get('turn_id', '')}", file=out, flush=True) |
| 50 | elif event_type == "token": |
| 51 | print(event.get("delta", ""), end="", file=out, flush=True) |
| 52 | elif event_type == "tool_start": |
| 53 | tool = event.get("tool", {}) |
| 54 | print(f"\n· tool: {tool.get('name')} started", file=out, flush=True) |
| 55 | elif event_type == "tool_progress": |
| 56 | progress = event.get("progress", {}) |
| 57 | tool = event.get("tool", {}) |
| 58 | print(f"· tool: {tool.get('name')} {progress.get('stage', 'running')}: {progress.get('message', '')}", file=out, flush=True) |
| 59 | elif event_type == "tool_result": |
| 60 | result = event["tool_result"] |
| 61 | status = "done" if result.get("ok") else "error" |
| 62 | print(f"· tool: {result.get('name')} {status}", file=out, flush=True) |
| 63 | elif event_type == "terminal": |
| 64 | stream = event.get("stream", "stdout") |
| 65 | print(f"· terminal[{stream}]: {event.get('line', '')}", file=out, flush=True) |
| 66 | elif event_type == "status": |
| 67 | print(f"· status: {event.get('phase')}: {event.get('message', '')}", file=out, flush=True) |
| 68 | elif event_type == "session": |
| 69 | session = (event.get("session") or {}).get("session") or {} |
| 70 | if session: |
| 71 | print(f"· session: {session.get('session_id')} {session.get('stage', '')}", file=out, flush=True) |
| 72 | elif event_type == "done": |
| 73 | print("", file=out, flush=True) |
| 74 | elif event_type == "error": |
| 75 | print(f"\nerror: {event.get('message', '')}", file=out, flush=True) |
| 76 | |
| 77 | |
| 78 | def prompt_inputs(args: argparse.Namespace) -> Iterable[str]: |
| 79 | if args.once: |
| 80 | yield args.once |
| 81 | return |
| 82 | if args.stdin_repl: |
| 83 | for line in sys.stdin: |
| 84 | user_input = line.strip() |
| 85 | if user_input: |
| 86 | yield user_input |
| 87 | return |
| 88 | if not sys.stdin.isatty(): |
| 89 | payload = sys.stdin.read().strip() |
| 90 | if payload: |
| 91 | yield payload |
| 92 | return |
| 93 | while True: |
| 94 | try: |
| 95 | user_input = input("› " if not args.jsonl else "") |
| 96 | except EOFError: |
| 97 | break |
| 98 | if user_input.strip(): |
| 99 | yield user_input.strip() |
| 100 | |
| 101 | |
| 102 | async def amain(argv: list[str] | None = None) -> int: |
| 103 | args = parse_args(argv) |
| 104 | if args.session and args.new_session: |
| 105 | print("error: --session and --new-session cannot be used together", file=sys.stderr) |
| 106 | return 2 |
| 107 | if args.new_session_name and not args.new_session: |
| 108 | print("error: --new-session-name requires --new-session", file=sys.stderr) |
| 109 | return 2 |
| 110 | if args.session or args.new_session: |
| 111 | try: |
| 112 | session_index = load_session_index() |
| 113 | if args.new_session: |
| 114 | if args.new_session_name: |
| 115 | session_index.create(project_name=args.new_session_name) |
| 116 | else: |
| 117 | session_index.create() |
| 118 | else: |
| 119 | session_index.set_active(args.session) |
| 120 | except KeyError: |
| 121 | print(f"error: unknown session id: {args.session}", file=sys.stderr) |
| 122 | return 2 |
| 123 | except ValueError as exc: |
| 124 | print(f"error: invalid session id: {exc}", file=sys.stderr) |
| 125 | return 2 |
| 126 | runtime = load_runtime() |
| 127 | interactive = sys.stdin.isatty() and not args.once |
| 128 | if interactive and not args.jsonl: |
| 129 | print("ViMax agent ready. Ctrl+C to exit.") |
| 130 | for user_input in prompt_inputs(args): |
| 131 | if user_input.strip() == "/compact": |
| 132 | turn_id = f"turn-{uuid4().hex[:12]}" |
| 133 | print_event({"type": "turn", "turn_id": turn_id, "turn": {"id": turn_id}}, jsonl=args.jsonl) |
| 134 | print_event({"type": "status", "turn_id": turn_id, "phase": "compact", "message": "Compacting context"}, jsonl=args.jsonl) |
| 135 | message = await runtime.compact_history(reason="manual") |
| 136 | print_event({"type": "token", "turn_id": turn_id, "delta": message}, jsonl=args.jsonl) |
| 137 | print_event({"type": "done", "turn_id": turn_id, "assistant": message, "tool_results": []}, jsonl=args.jsonl) |
| 138 | print_event({"type": "session", "turn_id": turn_id, "session": runtime.session_index.snapshot()}, jsonl=args.jsonl) |
| 139 | continue |
| 140 | try: |
| 141 | async for event in runtime.stream_events(user_input): |
| 142 | print_event(event, jsonl=args.jsonl) |
| 143 | except Exception as exc: |
| 144 | # Keep the REPL alive: one failed turn must not kill the process |
| 145 | # (and with it the TUI driving us over stdio). |
| 146 | turn_id = f"turn-{uuid4().hex[:12]}" |
| 147 | print_event({"type": "error", "turn_id": turn_id, "message": f"turn failed: {exc}"}, jsonl=args.jsonl) |
| 148 | print_event({"type": "done", "turn_id": turn_id, "assistant": "", "tool_results": []}, jsonl=args.jsonl) |
| 149 | return 0 |
| 150 | |
| 151 | |
| 152 | def main() -> None: |
| 153 | try: |
| 154 | raise SystemExit(asyncio.run(amain())) |
| 155 | except KeyboardInterrupt: |
| 156 | print("", file=sys.stderr) |
| 157 | raise SystemExit(130) |
| 158 | |
| 159 | |
| 160 | if __name__ == "__main__": |
| 161 | main() |
| 162 |