1#!/usr/bin/env python3
  2"""
  3s05: TodoWrite — add a planning tool on top of s04 hooks.
  4
  5  +---------+      +-------+      +------------------+
  6  |  User   | ---> |  LLM  | ---> | TOOL_HANDLERS    |
  7  | prompt  |      |       |      |  bash            |
  8  +---------+      +---+---+      |  read_file       |
  9                        ^         |  write_file      |
 10                        | result  |  edit_file       |
 11                        +---------+  glob            |
 12                                      todo_write ← NEW
 13                                   +------------------+
 14                                        |
 15                         in-memory current_todos
 16                                        |
 17                        if rounds_since_todo >= 3:
 18                          inject <reminder>
 19
 20Changes from s04:
 21  + todo_write tool + run_todo_write() implementation
 22  + Nag reminder (inject reminder after 3 rounds without todo update)
 23  + SYSTEM prompt includes "plan before execute" guidance
 24  + rounds_since_todo counter in agent_loop
 25  Loop unchanged: new tool auto-dispatches via TOOL_HANDLERS.
 26
 27Run: python s05_todo_write/code.py
 28Needs: pip install anthropic python-dotenv + ANTHROPIC_API_KEY in .env
 29"""
 30
 31import os, subprocess
 32from pathlib import Path
 33
 34try:
 35    import readline
 36    readline.parse_and_bind('set bind-tty-special-chars off')
 37except ImportError:
 38    pass
 39
 40from anthropic import Anthropic
 41from dotenv import load_dotenv
 42
 43load_dotenv(override=True)
 44if os.getenv("ANTHROPIC_BASE_URL"):
 45    os.environ.pop("ANTHROPIC_AUTH_TOKEN", None)
 46
 47WORKDIR = Path.cwd()
 48client = Anthropic(base_url=os.getenv("ANTHROPIC_BASE_URL"))
 49MODEL = os.environ["MODEL_ID"]
 50CURRENT_TODOS: list[dict] = []
 51
 52# s05 change: SYSTEM prompt adds planning guidance
 53SYSTEM = (
 54    f"You are a coding agent at {WORKDIR}. "
 55    "Before starting any multi-step task, use todo_write to plan your steps. "
 56    "Update status as you go."
 57)
 58
 59
 60# ═══════════════════════════════════════════════════════════
 61#  FROM s02-s04 (unchanged): Tool Implementations
 62# ═══════════════════════════════════════════════════════════
 63
 64def safe_path(p: str) -> Path:
 65    path = (WORKDIR / p).resolve()
 66    if not path.is_relative_to(WORKDIR):
 67        raise ValueError(f"Path escapes workspace: {p}")
 68    return path
 69
 70def run_bash(command: str) -> str:
 71    try:
 72        r = subprocess.run(command, shell=True, cwd=WORKDIR,
 73                           capture_output=True, text=True, timeout=120)
 74        out = (r.stdout + r.stderr).strip()
 75        return out[:50000] if out else "(no output)"
 76    except subprocess.TimeoutExpired:
 77        return "Error: Timeout (120s)"
 78
 79def run_read(path: str, limit: int | None = None) -> str:
 80    try:
 81        lines = safe_path(path).read_text().splitlines()
 82        if limit and limit < len(lines):
 83            lines = lines[:limit] + [f"... ({len(lines) - limit} more lines)"]
 84        return "\n".join(lines)
 85    except Exception as e:
 86        return f"Error: {e}"
 87
 88def run_write(path: str, content: str) -> str:
 89    try:
 90        file_path = safe_path(path)
 91        file_path.parent.mkdir(parents=True, exist_ok=True)
 92        file_path.write_text(content)
 93        return f"Wrote {len(content)} bytes to {path}"
 94    except Exception as e:
 95        return f"Error: {e}"
 96
 97def run_edit(path: str, old_text: str, new_text: str) -> str:
 98    try:
 99        file_path = safe_path(path)
100        text = file_path.read_text()
101        if old_text not in text:
102            return f"Error: text not found in {path}"
103        file_path.write_text(text.replace(old_text, new_text, 1))
104        return f"Edited {path}"
105    except Exception as e:
106        return f"Error: {e}"
107
108def run_glob(pattern: str) -> str:
109    import glob as g
110    try:
111        results = []
112        for match in g.glob(pattern, root_dir=WORKDIR):
113            if (WORKDIR / match).resolve().is_relative_to(WORKDIR):
114                results.append(match)
115        return "\n".join(results) if results else "(no matches)"
116    except Exception as e:
117        return f"Error: {e}"
118
119
120# ═══════════════════════════════════════════════════════════
121#  NEW in s05: todo_write tool — plan only, no execution
122# ═══════════════════════════════════════════════════════════
123
124# 传入的todos是模型新修改or创建的任务列表
125def run_todo_write(todos: list) -> str:
126    
127    # CURRENT_TODOS 是一个列表
128    # 每个元素都是一个字典，包含 "content" 和 "status" 两个字段
129    global CURRENT_TODOS
130    
131    # validate required fields
132    for i, t in enumerate(todos):
133        # 缺少 content 或 status 字段
134        if "content" not in t or "status" not in t: 
135            return f"Error: todos[{i}] missing 'content' or 'status'"
136        
137        # status 字段值不合法
138        if t["status"] not in ("pending", "in_progress", "completed"):
139            return f"Error: todos[{i}] has invalid status '{t['status']}'"
140        
141    CURRENT_TODOS = todos
142    lines = ["\n\033[33m## Current Tasks\033[0m"]
143    for t in CURRENT_TODOS:
144        icon = {"pending": " ", "in_progress": "\033[36m▸\033[0m", "completed": "\033[32m✓\033[0m"}[t["status"]]
145        lines.append(f"  [{icon}] {t['content']}")
146    print("\n".join(lines))
147    return f"Updated {len(CURRENT_TODOS)} tasks"
148
149TOOLS = [
150    {"name": "bash", "description": "Run a shell command.",
151     "input_schema": {"type": "object", "properties": {"command": {"type": "string"}}, "required": ["command"]}},
152    {"name": "read_file", "description": "Read file contents.",
153     "input_schema": {"type": "object", "properties": {"path": {"type": "string"}, "limit": {"type": "integer"}}, "required": ["path"]}},
154    {"name": "write_file", "description": "Write content to a file.",
155     "input_schema": {"type": "object", "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, "required": ["path", "content"]}},
156    {"name": "edit_file", "description": "Replace exact text in a file once.",
157     "input_schema": {"type": "object", "properties": {"path": {"type": "string"}, "old_text": {"type": "string"}, "new_text": {"type": "string"}}, "required": ["path", "old_text", "new_text"]}},
158    {"name": "glob", "description": "Find files matching a glob pattern.",
159     "input_schema": {"type": "object", "properties": {"pattern": {"type": "string"}}, "required": ["pattern"]}},
160    # s05: new tool
161    {"name": "todo_write", "description": "Create and manage a task list for your current coding session.",
162     "input_schema": {"type": "object", "properties": {"todos": {"type": "array", "items": {"type": "object", "properties": {"content": {"type": "string"}, "status": {"type": "string", "enum": ["pending", "in_progress", "completed"]}}, "required": ["content", "status"]}}}, "required": ["todos"]}},
163]
164
165TOOL_HANDLERS = {
166    "bash": run_bash, "read_file": run_read, "write_file": run_write,
167    "edit_file": run_edit, "glob": run_glob, "todo_write": run_todo_write,
168}
169
170
171# ═══════════════════════════════════════════════════════════
172#  FROM s04 (unchanged): Hook System
173# ═══════════════════════════════════════════════════════════
174
175HOOKS = {"UserPromptSubmit": [], "PreToolUse": [], "PostToolUse": [], "Stop": []}
176
177def register_hook(event: str, callback):
178    HOOKS[event].append(callback)
179
180def trigger_hooks(event: str, *args):
181    for callback in HOOKS[event]:
182        result = callback(*args)
183        if result is not None:
184            return result
185    return None
186
187# s04 hooks preserved
188DENY_LIST = ["rm -rf /", "sudo", "shutdown", "reboot", "mkfs", "dd if="]
189
190def permission_hook(block):
191    """PreToolUse: deny list check."""
192    if block.name == "bash":
193        for p in DENY_LIST:
194            if p in block.input.get("command", ""):
195                print(f"\n\033[31m⛔ Blocked: '{p}'\033[0m")
196                return "Permission denied"
197    return None
198
199def log_hook(block):
200    """PreToolUse: log tool calls."""
201    print(f"\033[90m[HOOK] {block.name}\033[0m" + f"{block.input if block.input else ''}")
202    return None
203
204def context_inject_hook(query: str):
205    """UserPromptSubmit: log working directory."""
206    print(f"\033[90m[HOOK] UserPromptSubmit: working in {WORKDIR}\033[0m")
207    return None
208
209def summary_hook(messages: list):
210    """Stop: print a clean summary of assistant texts and tool calls."""
211    print("\n\033[33m══════ Session Summary ══════\033[0m")
212    for msg in messages:
213        if msg["role"] != "assistant":
214            continue
215        for block in msg["content"]:
216            if getattr(block, "type", None) == "text":
217                print(f"\n\033[36m[Assistant]\033[0m {block.text}")
218            elif getattr(block, "type", None) == "tool_use":
219                print(f"\033[90m[Tool: {block.name}]\033[0m {block.input}")
220    print("\033[33m══════════════════════════════\033[0m\n")
221    return None
222
223register_hook("UserPromptSubmit", context_inject_hook)
224register_hook("PreToolUse", permission_hook)
225register_hook("PreToolUse", log_hook)
226register_hook("Stop", summary_hook)
227
228
229# ═══════════════════════════════════════════════════════════
230#  agent_loop — same as s04 + nag reminder counter
231# ═══════════════════════════════════════════════════════════
232
233rounds_since_todo = 0
234
235def agent_loop(messages: list):
236    global rounds_since_todo
237    while True:
238        # s05: nag reminder — inject if model hasn't updated todos for 3 rounds
239        # 提示词注入 提醒模型更新todos
240        if rounds_since_todo >= 3 and messages:
241            messages.append({"role": "user",
242                             "content": "<reminder>Update your todos.</reminder>"})
243            rounds_since_todo = 0
244
245        response = client.messages.create(
246            model=MODEL, system=SYSTEM, messages=messages,
247            tools=TOOLS, max_tokens=8000,
248        )
249        messages.append({"role": "assistant", "content": response.content})
250
251        if response.stop_reason != "tool_use":
252            force = trigger_hooks("Stop", messages)
253            if force:
254                messages.append({"role": "user", "content": force})
255                continue
256            return
257
258        rounds_since_todo += 1 # todo+1
259        results = []
260        for block in response.content:
261            if block.type != "tool_use":
262                continue
263
264            blocked = trigger_hooks("PreToolUse", block)
265            if blocked:
266                results.append({"type": "tool_result", "tool_use_id": block.id,
267                                "content": str(blocked)})
268                continue
269
270            handler = TOOL_HANDLERS.get(block.name)
271            output = handler(**block.input) if handler else f"Unknown: {block.name}"
272
273            trigger_hooks("PostToolUse", block, output)
274
275            # s05: reset nag counter when todo_write is called
276            if block.name == "todo_write":
277                rounds_since_todo = 0 # 清空
278
279            results.append({"type": "tool_result", "tool_use_id": block.id,
280                            "content": output})
281
282        messages.append({"role": "user", "content": results})
283
284
285if __name__ == "__main__":
286    print("s05: TodoWrite — plan before execute, nag if you forget")
287    print("Type a question, press Enter. Type q to quit.\n")
288
289    history = []
290    while True:
291        try:
292            query = input("\033[36ms05 >> \033[0m")
293        except (EOFError, KeyboardInterrupt):
294            break
295        if query.strip().lower() in ("q", "exit", ""):
296            break
297        trigger_hooks("UserPromptSubmit", query)
298        history.append({"role": "user", "content": query})
299        agent_loop(history)
300        for block in history[-1]["content"]:
301            if getattr(block, "type", None) == "text":
302                print(block.text)
303        print()