liusijin / nano.py

#!/usr/bin/env python3 """nanocode - minimal claude code alternative""" import glob as globlib import hashlib import json import os import random import re import readline import select import ssl import subprocess import sys import termios import time import tty import urllib.request import urllib.parse from datetime import datetime OPENROUTER_KEY = os.environ.get("OPENROUTER_API_KEY") LOCAL_API_KEY = os.environ.get("LOCAL_API_KEY") API_URL = ( "http://127.0.0.1:8990/v1/messages" if LOCAL_API_KEY else "https://openrouter.ai/api/v1/messages" if OPENROUTER_KEY else "https://api.anthropic.com/v1/messages" ) MODEL = os.environ.get("MODEL", "anthropic/claude-sonnet-4.5" if LOCAL_API_KEY else "anthropic/claude-opus-4.5" if OPENROUTER_KEY else "claude-opus-4-5" ) # ANSI colors RESET, BOLD, DIM = "\033[0m", "\033[1m", "\033[2m" BLUE, CYAN, GREEN, YELLOW, RED = "\033[34m", "\033[36m", "\033[32m", "\033[33m", "\033[31m" stop_flag = False def create_opener(): """Create URL opener with SSL and proxy support""" proxy = os.environ.get("http_proxy") or os.environ.get("https_proxy") ssl_ctx = ssl.create_default_context() ssl_ctx.check_hostname = False ssl_ctx.verify_mode = ssl.CERT_NONE handlers = [urllib.request.HTTPSHandler(context=ssl_ctx)] if proxy: handlers.insert(0, urllib.request.ProxyHandler({"http": proxy, "https": proxy})) return urllib.request.build_opener(*handlers) def register_tool(name, desc, params): """Register a tool from extension code""" def decorator(func): TOOLS[name] = (desc, params, func) return func return decorator def search_extension(args): """Search extensions from gist.kitchain.cn""" query = args.get("query", "") if not query: return "error: query required" try: # Split query into keywords keywords = query.lower().split() gist_info = {} # {gist_path: {"hits": count, "title": str, "desc": str, "topics": []}} opener = create_opener() # Search each keyword as a topic for keyword in keywords: url = f"https://gist.kitchain.cn/topics/{urllib.parse.quote(keyword)}" html = opener.open(urllib.request.Request(url), timeout=10).read().decode() # Extract gist URLs and titles gist_matches = re.findall( r'<a class="font-bold" href="https://gist\.kitchain\.cn/([^/]+/[a-f0-9]+)">([^<]+)</a>', html ) for gist_path, title in gist_matches: if gist_path not in gist_info: # Extract description and topics for this gist gist_section = re.search( rf'{re.escape(gist_path)}.*?' r'<h6 class="text-xs[^"]*">([^<]+)</h6>(.*?)</div>\s*</div>', html, re.DOTALL ) desc = "" topics = [] if gist_section: desc = gist_section.group(1).strip() topics_section = gist_section.group(2) topics = re.findall(r'topics/([^"]+)"[^>]*>([^<]+)<', topics_section) topics = [t[1] for t in topics] # Extract topic names gist_info[gist_path] = { "hits": 0, "title": title.strip(), "desc": desc, "topics": topics, "filename": title.strip() } gist_info[gist_path]["hits"] += 1 if not gist_info: return f"No extensions found: {query}" # Sort by hit count (descending) sorted_gists = sorted(gist_info.items(), key=lambda x: x[1]["hits"], reverse=True)[:10] result = f"Found {len(sorted_gists)} extensions:\n\n" for gist_path, info in sorted_gists: result += f"• {info['title']}\n" if info['desc']: result += f" {info['desc']}\n" if info['topics']: result += f" Topics: {', '.join(info['topics'])}\n" result += f" Matched: {info['hits']} keyword(s)\n\n" # Return first gist's load URL first_gist = sorted_gists[0][0] first_filename = sorted_gists[0][1]['filename'] result += f"To load the top result:\nload({{\"url\": \"https://gist.kitchain.cn/{first_gist}/raw/HEAD/{first_filename}\"}})" return result except Exception as e: return f"error: {e}" def load(args): """Load extension from URL""" url = args.get("url") if not url: return "error: url required" try: opener = create_opener() code = opener.open(urllib.request.Request(url), timeout=10).read().decode() exec(code, {"register_tool": register_tool, "TOOLS": TOOLS, "urllib": urllib, "json": json, "re": re, "subprocess": subprocess}) new = [k for k in TOOLS if k not in ["read","write","edit","glob","grep","bash","web_search","search_extension","load"]] return f"Loaded. New tools: {', '.join(new)}" except Exception as e: return f"error: {e}" # --- Tools --- def read(args): lines = open(args["path"]).readlines() offset, limit = args.get("offset", 0), args.get("limit", len(lines)) return "".join(f"{offset+i+1:4}| {l}" for i, l in enumerate(lines[offset:offset+limit])) def write(args): filepath = args["path"] content = args["content"] print(f"{DIM}[LOG] write: {filepath} ({len(content)} bytes){RESET}", flush=True) open(filepath, "w").write(content) print(f"{DIM}[LOG] write completed: {filepath}{RESET}", flush=True) return "ok" def edit(args): filepath = args["path"] print(f"{DIM}[LOG] edit: {filepath}{RESET}", flush=True) text = open(filepath).read() print(f"{DIM}[LOG] edit read: {len(text)} bytes{RESET}", flush=True) old, new = args["old"], args["new"] if old not in text: return "error: old_string not found" count = text.count(old) if not args.get("all") and count > 1: return f"error: old_string appears {count} times (use all=true)" result = text.replace(old, new) if args.get("all") else text.replace(old, new, 1) print(f"{DIM}[LOG] edit writing: {len(result)} bytes{RESET}", flush=True) open(filepath, "w").write(result) print(f"{DIM}[LOG] edit completed: {filepath}{RESET}", flush=True) return "ok" def glob(args): pattern = (args.get("path", ".") + "/" + args["pat"]).replace("//", "/") files = sorted(globlib.glob(pattern, recursive=True), key=lambda f: os.path.getmtime(f) if os.path.isfile(f) else 0, reverse=True) return "\n".join(files) or "none" def grep(args): pattern, hits = re.compile(args["pat"]), [] for fp in globlib.glob(args.get("path", ".") + "/**", recursive=True): try: for n, l in enumerate(open(fp), 1): if pattern.search(l): hits.append(f"{fp}:{n}:{l.rstrip()}") except: pass return "\n".join(hits[:50]) or "none" def bash(args): global stop_flag proc = subprocess.Popen(args["cmd"], shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) lines = [] old_settings = termios.tcgetattr(sys.stdin) try: tty.setcbreak(sys.stdin.fileno()) if proc.stdout: import fcntl fd = proc.stdout.fileno() fcntl.fcntl(fd, fcntl.F_SETFL, fcntl.fcntl(fd, fcntl.F_GETFL) | os.O_NONBLOCK) while True: # Check ESC key if select.select([sys.stdin], [], [], 0)[0]: if sys.stdin.read(1) == '\x1b': stop_flag = True proc.kill() lines.append("\n(stopped)") print(f"\n{YELLOW}⏸ Stopped{RESET}") break # Read output if select.select([proc.stdout], [], [], 0.1)[0]: line = proc.stdout.readline() if line: print(f" {DIM}│ {line.rstrip()}{RESET}", flush=True) lines.append(line) # Check if done if proc.poll() is not None: remaining = proc.stdout.read() if remaining: for line in remaining.split('\n'): if line: print(f" {DIM}│ {line.rstrip()}{RESET}", flush=True) lines.append(line + '\n') break if not stop_flag: proc.wait(timeout=30) except subprocess.TimeoutExpired: proc.kill() lines.append("\n(timeout)") finally: termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings) return "".join(lines).strip() or "(empty)" def web_search(args): """Search web using DuckDuckGo""" query, max_results = args["query"], args.get("max_results", 5) try: url = f"https://html.duckduckgo.com/html/?q={urllib.parse.quote_plus(query)}" headers = {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36"} opener = create_opener() html = opener.open(urllib.request.Request(url, headers=headers), timeout=30).read().decode() # Extract titles and URLs links = re.findall(r'class="result__a"[^>]+href="([^"]+)"[^>]*>([^<]+)<', html) # Extract snippets snippets = re.findall(r'class="result__snippet"[^>]*>([^<]*)<', html) if not links: return "No results found" results = [] for i, ((link, title), snippet) in enumerate(zip(links[:max_results], snippets[:max_results] + [""] * max_results), 1): results.append(f"{i}. {title.strip()}\n URL: {link}\n {snippet.strip()}\n") return "\n".join(results) except Exception as e: return f"error: {e}" TOOLS = { "read": ("Read file with line numbers", {"path": "string", "offset": "number?", "limit": "number?"}, read), "write": ("Write content to file", {"path": "string", "content": "string"}, write), "edit": ("Replace old with new in file", {"path": "string", "old": "string", "new": "string", "all": "boolean?"}, edit), "glob": ("Find files by pattern", {"pat": "string", "path": "string?"}, glob), "grep": ("Search files for regex", {"pat": "string", "path": "string?"}, grep), "bash": ("Run shell command", {"cmd": "string"}, bash), "web_search": ("Search the web using DuckDuckGo", {"query": "string", "max_results": "number?"}, web_search), "search_extension": ("Search for extensions to add new capabilities (GitHub docs, web scraping, APIs, etc)", {"query": "string"}, search_extension), "load": ("Load extension from URL to add new tools", {"url": "string"}, load), } def run_tool(name, args): try: return TOOLS[name][2](args) except Exception as e: return f"error: {e}" def make_schema(): result = [] for name, (desc, params, _) in TOOLS.items(): props, req = {}, [] for pname, ptype in params.items(): opt = ptype.endswith("?") props[pname] = {"type": "integer" if ptype.rstrip("?") == "number" else ptype.rstrip("?")} if not opt: req.append(pname) result.append({"name": name, "description": desc, "input_schema": {"type": "object", "properties": props, "required": req}}) return result def call_api(messages, system_prompt, stream=True, enable_thinking=True, use_tools=True): headers = {"Content-Type": "application/json", "anthropic-version": "2023-06-01"} if LOCAL_API_KEY: headers["Authorization"] = f"Bearer {LOCAL_API_KEY}" elif OPENROUTER_KEY: headers["Authorization"] = f"Bearer {OPENROUTER_KEY}" else: headers["x-api-key"] = os.environ.get("ANTHROPIC_API_KEY", "") data = {"model": MODEL, "max_tokens": 8192, "system": system_prompt, "messages": messages, "stream": stream} if use_tools: data["tools"] = make_schema() if enable_thinking and os.environ.get("THINKING"): data["thinking"] = {"type": "enabled", "budget_tokens": int(os.environ.get("THINKING_BUDGET", "10000"))} req = urllib.request.Request(API_URL, json.dumps(data).encode(), headers, method="POST") return create_opener().open(req) def summarize_changes(user_input, files_modified, checkpoint_manager, checkpoint_id): """Use LLM to summarize the changes made in this turn Args: user_input: User's request files_modified: Set of modified file paths checkpoint_manager: CheckpointManager instance checkpoint_id: Checkpoint hash to get diff from Returns: str: One-line summary of changes """ if not files_modified or not checkpoint_id: return user_input[:50] try: # Get diff from git diff_output = checkpoint_manager._git_command( "--git-dir", checkpoint_manager.bare_repo, "show", "--format=", checkpoint_id ) # Check if diff is empty or error - no actual changes if not diff_output or diff_output.startswith("error") or len(diff_output.strip()) == 0: # No diff available, just use user input return user_input[:50] # Limit diff size to avoid token overflow (max ~3000 chars) if len(diff_output) > 3000: diff_output = diff_output[:3000] + "\n... (truncated)" summary_prompt = f"""Based on the actual code changes (diff), generate a brief Chinese summary (max 30 Chinese characters). IMPORTANT: Must be based on the actual code changes, not the user's description. Code changes (diff): {diff_output} User description (for reference only): {user_input} Requirements: 1. Describe what code/functionality was actually modified 2. Reply in Chinese only, no explanation 3. No quotes 4. Max 30 Chinese characters Good examples: - 在 auth.py 添加 JWT 验证 - 修复 parser.py 空指针异常 - 重构 database.py 连接池 - 更新 README 添加安装说明 """ messages = [{"role": "user", "content": summary_prompt}] response = call_api(messages, "You are a code change analyzer, skilled at extracting key information from diffs. Reply in Chinese.", stream=False, enable_thinking=False, use_tools=False) # Parse non-streaming response data = json.loads(response.read().decode()) blocks = data.get("content", []) for block in blocks: if block.get("type") == "text": summary = block.get("text", "").strip() # Remove thinking tags if present if "<thinking>" in summary: # Extract content after </thinking> parts = summary.split("</thinking>") if len(parts) > 1: summary = parts[-1].strip() # Clean up and limit length summary = summary.replace('"', '').replace("'", "") if summary and len(summary) <= 80: return summary # Fallback to user input return user_input[:50] except Exception as e: # On error, fallback to user input return user_input[:50] def process_stream(response): """简化的流式处理，支持ESC中断""" global stop_flag blocks, current, text_buf, json_buf, think_buf = [], None, "", "", "" # Save terminal settings old_settings = termios.tcgetattr(sys.stdin) try: tty.setcbreak(sys.stdin.fileno()) for line in response: if select.select([sys.stdin], [], [], 0)[0]: ch = sys.stdin.read(1) if ch == '\x1b': # ESC key stop_flag = True print(f"\n{YELLOW}⏸ Stopped{RESET}") break line = line.decode("utf-8").strip() if not line.startswith("data: "): continue if line == "data: [DONE]": continue try: data = json.loads(line[6:]) etype = data.get("type") if etype == "content_block_start": block = data.get("content_block", {}) current = {"type": block.get("type"), "id": block.get("id")} if current["type"] == "text": text_buf = "" print(f"\n{CYAN}⏺{RESET} ", end="", flush=True) elif current["type"] == "thinking": think_buf = "" print(f"\n{YELLOW}💭{RESET} {DIM}", end="", flush=True) elif current["type"] == "tool_use": current["name"] = block.get("name") json_buf = "" elif etype == "content_block_delta": delta = data.get("delta", {}) dtype = delta.get("type") if dtype == "text_delta": text = delta.get("text", "") text_buf += text print(text, end="", flush=True) elif dtype == "thinking_delta": text = delta.get("thinking", "") think_buf += text print(text, end="", flush=True) elif dtype == "input_json_delta" and current: json_buf += delta.get("partial_json", "") elif etype == "content_block_stop" and current: if current["type"] == "text": current["text"] = text_buf print() elif current["type"] == "thinking": print(RESET) elif current["type"] == "tool_use": try: current["input"] = json.loads(json_buf) except: current["input"] = {} blocks.append(current) current = None except: pass finally: termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings) return blocks def is_file_in_project(filepath, project_path): """Check if file is within project directory""" try: abs_file = os.path.abspath(filepath) abs_project = os.path.abspath(project_path) # Check if file is under project directory return abs_file.startswith(abs_project + os.sep) or abs_file == abs_project except: return False def read_multiline_input(): """Read multiline input. Enter to submit, Alt+Enter for newline.""" lines = [] current = "" cursor_pos = 0 # Cursor position in current line # Enable bracketed paste mode print("\033[?2004h", end="", flush=True) old_settings = termios.tcgetattr(sys.stdin) try: tty.setcbreak(sys.stdin.fileno()) print(f"{BOLD}{BLUE}❯{RESET} ", end="", flush=True) while True: ch = sys.stdin.read(1) if ch == '\x03': # Ctrl+C - clear input lines.clear() current = "" cursor_pos = 0 print("\r\033[K", end="", flush=True) print(f"{BOLD}{BLUE}❯{RESET} ", end="", flush=True) continue if ch == '\x04': # Ctrl+D raise EOFError if ch == '\x1b': # Escape sequence next_ch = sys.stdin.read(1) if next_ch in ('\r', '\n'): # Alt+Enter lines.append(current) current = "" cursor_pos = 0 print(f"\n{BOLD}{BLUE}│{RESET} ", end="", flush=True) elif next_ch == '[': # Escape sequence seq = sys.stdin.read(1) if seq == 'C': # Right arrow if cursor_pos < len(current): cursor_pos += 1 print("\033[C", end="", flush=True) elif seq == 'D': # Left arrow if cursor_pos > 0: cursor_pos -= 1 print("\033[D", end="", flush=True) elif seq == '2': # Bracketed paste start: ESC[200~ rest = sys.stdin.read(3) # Read "00~" if rest == '00~': # Read pasted content until ESC[201~ paste_buf = "" while True: c = sys.stdin.read(1) if c == '\x1b': # Check for [201~ peek = sys.stdin.read(5) if peek == '[201~': break else: paste_buf += c + peek else: paste_buf += c # Process pasted content paste_lines = paste_buf.split('\n') if len(paste_lines) == 1: # Single line paste current = current[:cursor_pos] + paste_lines[0] + current[cursor_pos:] cursor_pos += len(paste_lines[0]) prefix = f"{BOLD}{BLUE}{'│' if lines else '❯'}{RESET} " print(f"\r\033[K{prefix}{current}", end="", flush=True) else: # Multi-line paste # First line appends to current first_line = current[:cursor_pos] + paste_lines[0] print(paste_lines[0], end="", flush=True) if first_line: lines.append(first_line) # Middle lines for line in paste_lines[1:-1]: print(f"\n{BOLD}{BLUE}│{RESET} {line}", end="", flush=True) lines.append(line) # Last line becomes new current current = paste_lines[-1] cursor_pos = len(current) print(f"\n{BOLD}{BLUE}│{RESET} {current}", end="", flush=True) continue if ch in ('\r', '\n'): # Enter - submit if current: lines.append(current) print() break if ch in ('\x7f', '\x08'): # Backspace if cursor_pos > 0: # Delete character before cursor current = current[:cursor_pos-1] + current[cursor_pos:] cursor_pos -= 1 # Redraw current line prefix = f"{BOLD}{BLUE}{'│' if lines else '❯'}{RESET} " print(f"\r\033[K{prefix}{current}", end="", flush=True) # Move cursor back to position if cursor_pos < len(current): print(f"\033[{len(current) - cursor_pos}D", end="", flush=True) elif lines: # Merge with previous line prev_line = lines.pop() cursor_pos = len(prev_line) # Cursor at end of previous line current = prev_line + current # Move up and redraw print("\033[A\033[K", end="", flush=True) prefix = f"{BOLD}{BLUE}{'│' if lines else '❯'}{RESET} " print(f"\r{prefix}{current}", end="", flush=True) if cursor_pos < len(current): print(f"\033[{len(current) - cursor_pos}D", end="", flush=True) continue if ch.isprintable() or ch == '\t': # Insert character at cursor position current = current[:cursor_pos] + ch + current[cursor_pos:] cursor_pos += 1 # Redraw from cursor position print(f"{ch}{current[cursor_pos:]}", end="", flush=True) # Move cursor back if needed if cursor_pos < len(current): print(f"\033[{len(current) - cursor_pos}D", end="", flush=True) finally: # Disable bracketed paste mode print("\033[?2004l", end="", flush=True) termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings) return "\n".join(lines).strip() def main(): global stop_flag # Parse command line arguments continue_session = "-c" in sys.argv or "--continue" in sys.argv list_sessions = "-l" in sys.argv or "--list" in sys.argv # Disable Ctrl+C signal old_settings = termios.tcgetattr(sys.stdin) new_settings = termios.tcgetattr(sys.stdin) new_settings[3] = new_settings[3] & ~termios.ISIG # Disable signal generation termios.tcsetattr(sys.stdin, termios.TCSADRAIN, new_settings) try: proxy = os.environ.get("http_proxy") or os.environ.get("https_proxy") proxy_info = f" | {DIM}🌐 {proxy}{RESET}" if proxy else "" thinking_info = f" | {YELLOW}💭{RESET}" if os.environ.get("THINKING") else "" if list_sessions: session_mode = f" | {YELLOW}Select{RESET}" elif continue_session: session_mode = f" | {GREEN}Continue{RESET}" else: session_mode = f" | {CYAN}New{RESET}" print(f"{BOLD}nanocode{RESET} | {DIM}{MODEL} | {os.getcwd()}{proxy_info}{thinking_info}{session_mode}{RESET}") print(f"{DIM}Shortcuts: Enter=submit | Alt+Enter=newline | Ctrl+C=clear input | Ctrl+D=exit | ESC=stop{RESET}") print(f"{DIM}Commands: /c [all|baseline|<id>] | /ca | /clear{RESET}") print(f"{DIM}Usage: nanocode (new) | nanocode -c (continue) | nanocode -l (select){RESET}\n") selected_session_id = None if list_sessions: selected_session_id = select_session_interactive() if not selected_session_id: print(f"{DIM}Exiting...{RESET}") return run_main_loop(continue_session, selected_session_id) finally: termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings) def select_session_interactive(): """Display sessions and let user select one Returns: session_id: Selected session ID, or None if cancelled """ session_manager = SessionManager(os.getcwd()) sessions = session_manager.list_sessions()[:10] # Limit to 10 most recent if not sessions: print(f"{YELLOW}⚠ No previous sessions found{RESET}") print(f"{DIM}Starting new session...{RESET}\n") return None print(f"{BOLD}📂 Recent Sessions:{RESET}\n") for i, sess_info in enumerate(sessions, 1): created = datetime.fromtimestamp(sess_info['metadata']['created_at']).strftime('%Y-%m-%d %H:%M') last_active = datetime.fromtimestamp(sess_info['metadata']['last_active']).strftime('%Y-%m-%d %H:%M') desc = sess_info['metadata'].get('description', '(no description)') # Git info git_commit = sess_info['metadata'].get('git_commit') git_branch = sess_info['metadata'].get('git_branch') git_dirty = sess_info['metadata'].get('git_dirty', False) print(f"{CYAN}{i}.{RESET} {BOLD}{sess_info['session_id']}{RESET}") print(f" {desc}") git_info = "" if git_commit and git_branch: dirty_mark = f"{YELLOW}*{RESET}" if git_dirty else "" git_info = f" | Git: {git_branch}@{git_commit}{dirty_mark}" print(f" Created: {created} | Last: {last_active} | {sess_info['message_count']} messages{git_info}\n") print(f"{DIM}Enter session number (1-{len(sessions)}), or press Enter for new session:{RESET}") try: choice = input(f"{BOLD}{BLUE}❯{RESET} ").strip() if not choice: # Empty input = new session return None try: idx = int(choice) - 1 if 0 <= idx < len(sessions): return sessions[idx]['session_id'] else: print(f"{RED}✗ Invalid number{RESET}") return None except ValueError: print(f"{RED}✗ Invalid input{RESET}") return None except (EOFError, KeyboardInterrupt): return None def run_main_loop(continue_session=False, selected_session_id=None): # Initialize session manager session_manager = SessionManager(os.getcwd()) # Load or create session based on parameters if selected_session_id: # Load specific session selected by user session = session_manager.load_session(selected_session_id) if session: git_info = "" git_commit = session.metadata.get('git_commit') git_branch = session.metadata.get('git_branch') if git_commit and git_branch: git_dirty = session.metadata.get('git_dirty', False) dirty_mark = f"{YELLOW}*{RESET}" if git_dirty else "" git_info = f" | Git: {git_branch}@{git_commit}{dirty_mark}" print(f"{GREEN}✓ Loaded session: {session.session_id}{RESET}") print(f"{DIM} └─ {len(session.messages)} messages{git_info}{RESET}") # Check for conflicts conflicts = session.detect_conflicts() if conflicts: print(f"\n{YELLOW}⚠ File conflicts detected:{RESET}") for filepath in conflicts[:5]: print(f" - {filepath}") if len(conflicts) > 5: print(f" ... and {len(conflicts)-5} more") print(f"\n{DIM}These files have been modified outside this session.{RESET}") confirm = input(f"{BOLD}Continue anyway? (y/N/u=update): {RESET}").strip().lower() if confirm == 'u': session.update_file_states() session_manager.save_session() print(f"{GREEN}✓ Updated file states{RESET}\n") elif confirm != 'y': print(f"{DIM}Creating new session instead...{RESET}\n") session_manager.create_session() else: print() else: print() else: print(f"{RED}✗ Failed to load session{RESET}") print(f"{GREEN}✓ Creating new session instead{RESET}\n") session_manager.create_session() elif continue_session: # Continue last session last_session = session_manager.load_last_session() if last_session: git_info = "" git_commit = last_session.metadata.get('git_commit') git_branch = last_session.metadata.get('git_branch') if git_commit and git_branch: git_dirty = last_session.metadata.get('git_dirty', False) dirty_mark = f"{YELLOW}*{RESET}" if git_dirty else "" git_info = f" | Git: {git_branch}@{git_commit}{dirty_mark}" print(f"{GREEN}✓ Continued session: {last_session.session_id}{RESET}") print(f"{DIM} └─ {len(last_session.messages)} messages{git_info}{RESET}") # Check for conflicts conflicts = last_session.detect_conflicts() if conflicts: print(f"\n{YELLOW}⚠ File conflicts detected:{RESET}") for filepath in conflicts[:5]: print(f" - {filepath}") if len(conflicts) > 5: print(f" ... and {len(conflicts)-5} more") print(f"\n{DIM}These files have been modified outside this session.{RESET}") confirm = input(f"{BOLD}Continue anyway? (y/N/u=update): {RESET}").strip().lower() if confirm == 'u': last_session.update_file_states() session_manager.save_session() print(f"{GREEN}✓ Updated file states{RESET}\n") elif confirm != 'y': print(f"{DIM}Creating new session instead...{RESET}\n") session_manager.create_session() else: print() else: print() else: # No previous session, create new one session_manager.create_session() print(f"{YELLOW}⚠ No previous session found{RESET}") print(f"{GREEN}✓ Created new session: {session_manager.current_session.session_id}{RESET}\n") else: # Always create new session by default # Try to detect parent from last session's latest checkpoint parent_checkpoint = None parent_session = None last_session = session_manager.load_last_session() if last_session: # Get the latest checkpoint from last session checkpoints = session_manager.checkpoint_manager.list_checkpoints(show_all=False) if checkpoints: parent_checkpoint = checkpoints[0][0] # Latest checkpoint hash parent_session = last_session.session_id session_manager.create_session( parent_checkpoint=parent_checkpoint, parent_session=parent_session ) git_info = "" git_commit = session_manager.current_session.metadata.get('git_commit') git_branch = session_manager.current_session.metadata.get('git_branch') if git_commit and git_branch: git_dirty = session_manager.current_session.metadata.get('git_dirty', False) dirty_mark = f"{YELLOW}*{RESET}" if git_dirty else "" git_info = f" | Git: {git_branch}@{git_commit}{dirty_mark}" if parent_checkpoint: print(f"{GREEN}✓ Created new session: {session_manager.current_session.session_id}{RESET}") print(f"{DIM} └─ Branched from {parent_session[:8]}... @ {parent_checkpoint}{git_info}{RESET}\n") else: print(f"{GREEN}✓ Created new session: {session_manager.current_session.session_id}{RESET}") if git_info: print(f"{DIM} └─{git_info}{RESET}\n") else: print() files_modified = set() auto_checkpoint = True current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") system_prompt = f"""Concise coding assistant. cwd: {os.getcwd()} Current time: {current_time} IMPORTANT: When you don't have a tool for the task, ALWAYS try search_extension first before saying you can't do it. Examples: - User asks about GitHub repo → search_extension({{"query": "github documentation"}}) - User needs web data → search_extension({{"query": "web scraping"}}) - User needs API → search_extension({{"query": "api client"}})""" while True: try: print(f"{DIM}{'─'*80}{RESET}") user_input = read_multiline_input() print(f"{DIM}{'─'*80}{RESET}") if not user_input: continue if user_input in ("/q", "exit"): session_manager.save_session() break # Handle /clear command first (before /c to avoid conflict) if user_input == "/clear": # Save current session session_manager.save_session() # Get latest checkpoint from current session (if any) checkpoints = session_manager.checkpoint_manager.list_checkpoints(show_all=False) parent_checkpoint = checkpoints[0][0] if checkpoints else None parent_session = session_manager.current_session.session_id # Create new session branched from current session_manager.create_session( parent_checkpoint=parent_checkpoint, parent_session=parent_session ) # Reset state files_modified.clear() print(f"{GREEN}✓ Started new session: {session_manager.current_session.session_id}{RESET}") if parent_checkpoint: print(f"{DIM} └─ Branched from {parent_session[:8]}... @ {parent_checkpoint}{RESET}") continue # Handle checkpoint commands if user_input.startswith("/checkpoint") or user_input.startswith("/c") or user_input == "/ca": parts = user_input.split() # /ca is shortcut for /c all if parts[0] == "/ca": parts = ["/c", "all"] # /c without args defaults to list if len(parts) == 1 and parts[0] in ["/c", "/checkpoint"]: parts.append("list") restored_messages = handle_checkpoint_command(parts, session_manager, files_modified) if restored_messages is not None: # Restore conversation by replacing session messages session_manager.current_session.messages = restored_messages session_manager.save_session() continue # Add user message to current session session_manager.current_session.messages.append({"role": "user", "content": user_input}) # Reset stop flag for new turn stop_flag = False # Track files modified in this turn files_modified_this_turn = set() while True: response = call_api(session_manager.current_session.messages, system_prompt) blocks = process_stream(response) if stop_flag: break tool_results = [] for block in blocks: if block["type"] == "tool_use": name, args = block["name"], block["input"] # Save baseline BEFORE executing write/edit if name in ['write', 'edit']: filepath = args.get('path') if filepath and is_file_in_project(filepath, session_manager.project_path): session_manager.save_baseline_if_needed(filepath) preview = str(list(args.values())[0])[:50] if args else "" print(f"\n{GREEN}⏺ {name}{RESET}({DIM}{preview}{RESET})") result = run_tool(name, args) lines = result.split("\n") prev = lines[0][:60] + ("..." if len(lines[0]) > 60 else "") if len(lines) > 1: prev += f" +{len(lines)-1}" print(f" {DIM}⎿ {prev}{RESET}") # Track file modifications (only project files) if name in ['write', 'edit']: filepath = args.get('path') if filepath and is_file_in_project(filepath, session_manager.project_path): files_modified.add(filepath) files_modified_this_turn.add(filepath) session_manager.current_session.track_file_state(filepath) tool_results.append({"type": "tool_result", "tool_use_id": block["id"], "content": result}) # Check stop_flag after each tool execution if stop_flag: print(f"{YELLOW}⚠ Tool execution stopped{RESET}") break session_manager.current_session.messages.append({"role": "assistant", "content": blocks}) if not tool_results or stop_flag: break session_manager.current_session.messages.append({"role": "user", "content": tool_results}) # Auto checkpoint after AI work (if project files were modified) if auto_checkpoint and files_modified_this_turn: # files_modified_this_turn already filtered to project files only # Use parent_commit for first checkpoint of new session parent_commit = session_manager.parent_commit_for_next_checkpoint checkpoint_id = session_manager.checkpoint_manager.create_checkpoint( f"Auto: {user_input[:50]}", list(files_modified_this_turn), conversation_snapshot=session_manager.current_session.messages.copy(), parent_commit=parent_commit ) # Clear parent after first checkpoint if parent_commit: session_manager.parent_commit_for_next_checkpoint = None if checkpoint_id: # Generate summary using LLM with actual diff print(f"{DIM}Generating checkpoint summary...{RESET}", end="", flush=True) summary = summarize_changes( user_input, files_modified_this_turn, session_manager.checkpoint_manager, checkpoint_id ) print(f"\r{' ' * 40}\r", end="", flush=True) # Clear the line # Update commit message with better summary (only if different from temp message) temp_message = f"Auto: {user_input[:50]}" if summary != user_input[:50] and summary != temp_message: session_manager.checkpoint_manager._git_command( "--git-dir", session_manager.checkpoint_manager.bare_repo, "commit", "--amend", "-m", summary ) print(f"\n{YELLOW}📍 {checkpoint_id}: {summary}{RESET}") else: # Checkpoint creation failed (e.g., no actual diff) print(f"\n{DIM}(No project file changes to checkpoint){RESET}") # Auto-save session after each interaction session_manager.save_session() print() except EOFError: session_manager.save_session() break except Exception as e: print(f"{RED}⏺ Error: {e}{RESET}") # ============================================================================ # Checkpoint & Session Management (Phase 1+2) # ============================================================================ class CheckpointManager: """Manage checkpoints using shadow bare git repository with session isolation""" def __init__(self, project_path, session_id=None): self.project_path = project_path self.session_id = session_id self.nanocode_dir = os.path.join(project_path, ".nanocode") self.bare_repo = os.path.join(self.nanocode_dir, "checkpoint.git") self._init_bare_repo() def set_session(self, session_id): """Set current session for checkpoint operations""" self.session_id = session_id def _get_branch_name(self): """Get git branch name for current session""" if not self.session_id: return "main" return f"session_{self.session_id}" def _init_bare_repo(self): """Initialize shadow bare repository""" if not os.path.exists(self.bare_repo): os.makedirs(self.bare_repo, exist_ok=True) try: subprocess.run( ["git", "init", "--bare", self.bare_repo], capture_output=True, check=True ) except (subprocess.CalledProcessError, FileNotFoundError): # Git not available, will handle gracefully pass def _git_command(self, *args, cwd=None): """Execute git command""" try: result = subprocess.run( ["git"] + list(args), cwd=cwd or self.project_path, capture_output=True, text=True, check=True ) return result.stdout.strip() except (subprocess.CalledProcessError, FileNotFoundError) as e: return f"error: {e}" def save_file_to_blob(self, filepath): """Save file to git blob storage Returns: str: blob hash """ try: result = subprocess.run( ["git", "--git-dir", self.bare_repo, "hash-object", "-w", filepath], capture_output=True, text=True, check=True ) return result.stdout.strip() except Exception as e: return None def restore_file_from_blob(self, blob_hash, filepath): """Restore file from git blob storage""" try: content = subprocess.run( ["git", "--git-dir", self.bare_repo, "cat-file", "-p", blob_hash], capture_output=True, check=True ).stdout os.makedirs(os.path.dirname(filepath) or ".", exist_ok=True) with open(filepath, 'wb') as f: f.write(content) return True except Exception as e: return False def get_file_git_info(self, filepath): """Get file's git info from project .git Returns: dict: {"commit": "abc123", "has_changes": True/False} or None """ try: # Check if file is tracked result = subprocess.run( ["git", "ls-files", "--error-unmatch", filepath], cwd=self.project_path, capture_output=True, text=True ) if result.returncode != 0: return None # Not tracked # Get last commit for this file commit = subprocess.run( ["git", "log", "-1", "--format=%H", "--", filepath], cwd=self.project_path, capture_output=True, text=True, check=True ).stdout.strip() # Check if file has local changes diff = subprocess.run( ["git", "diff", "HEAD", "--", filepath], cwd=self.project_path, capture_output=True, text=True, check=True ).stdout.strip() return { "commit": commit, "has_changes": bool(diff) } except Exception as e: return None def create_checkpoint(self, message, files_changed, conversation_snapshot=None, parent_commit=None): """Create a checkpoint on current session's branch Args: message: Commit message files_changed: List of modified files conversation_snapshot: Conversation state to save parent_commit: Parent commit hash to branch from (for new sessions) """ print(f"{DIM}[LOG] create_checkpoint: files_changed={files_changed}{RESET}", flush=True) if not files_changed or not self.session_id: return None branch_name = self._get_branch_name() # Save conversation snapshot if conversation_snapshot: snapshot_file = os.path.join(self.nanocode_dir, "conversation_snapshots.json") snapshots = {} if os.path.exists(snapshot_file): with open(snapshot_file, 'r') as f: snapshots = json.load(f) # Create temp worktree for this session temp_worktree = os.path.join(self.nanocode_dir, f"temp_worktree_{self.session_id}") try: # Check if branch exists branch_exists = self._git_command("--git-dir", self.bare_repo, "rev-parse", "--verify", branch_name) if not branch_exists or branch_exists.startswith("error"): # Create new branch os.makedirs(temp_worktree, exist_ok=True) self._git_command("--git-dir", self.bare_repo, "--work-tree", temp_worktree, "config", "core.bare", "false") # If parent_commit specified, branch from it if parent_commit: # Create branch from parent commit self._git_command("--git-dir", self.bare_repo, "branch", branch_name, parent_commit) self._git_command("--git-dir", self.bare_repo, "--work-tree", temp_worktree, "checkout", branch_name, "-f") else: # Create orphan branch (no parent) self._git_command("--git-dir", self.bare_repo, "--work-tree", temp_worktree, "checkout", "--orphan", branch_name) # Copy files to temp worktree for filepath in files_changed: print(f"{DIM}[LOG] checkpoint copying: {filepath}{RESET}", flush=True) if os.path.exists(filepath): file_size = os.path.getsize(filepath) print(f"{DIM}[LOG] source file exists: {filepath} ({file_size} bytes){RESET}", flush=True) # Convert absolute path to relative path if os.path.isabs(filepath): rel_filepath = os.path.relpath(filepath, self.project_path) else: rel_filepath = filepath dest = os.path.join(temp_worktree, rel_filepath) os.makedirs(os.path.dirname(dest), exist_ok=True) with open(filepath, 'rb') as src, open(dest, 'wb') as dst: content = src.read() dst.write(content) print(f"{DIM}[LOG] copied to temp_worktree: {dest} ({len(content)} bytes){RESET}", flush=True) else: print(f"{DIM}[LOG] source file NOT exists: {filepath}{RESET}", flush=True) # Commit self._git_command("--git-dir", self.bare_repo, "--work-tree", temp_worktree, "add", "-A") self._git_command("--git-dir", self.bare_repo, "--work-tree", temp_worktree, "commit", "-m", message, "--allow-empty") commit_hash = self._git_command("--git-dir", self.bare_repo, "rev-parse", "HEAD") checkpoint_id = commit_hash[:8] if commit_hash and not commit_hash.startswith("error") else None # Save conversation snapshot with checkpoint_id if checkpoint_id and conversation_snapshot: snapshots[checkpoint_id] = conversation_snapshot with open(snapshot_file, 'w') as f: json.dump(snapshots, f, indent=2) return checkpoint_id else: # Branch exists, checkout and commit self._git_command("--git-dir", self.bare_repo, "--work-tree", temp_worktree, "checkout", branch_name, "-f") # Update temp worktree for filepath in files_changed: print(f"{DIM}[LOG] checkpoint updating: {filepath}{RESET}", flush=True) if os.path.exists(filepath): file_size = os.path.getsize(filepath) print(f"{DIM}[LOG] source file exists: {filepath} ({file_size} bytes){RESET}", flush=True) # Convert absolute path to relative path if os.path.isabs(filepath): rel_filepath = os.path.relpath(filepath, self.project_path) else: rel_filepath = filepath dest = os.path.join(temp_worktree, rel_filepath) os.makedirs(os.path.dirname(dest), exist_ok=True) with open(filepath, 'rb') as src, open(dest, 'wb') as dst: content = src.read() dst.write(content) print(f"{DIM}[LOG] copied to temp_worktree: {dest} ({len(content)} bytes){RESET}", flush=True) else: print(f"{DIM}[LOG] source file NOT exists: {filepath}{RESET}", flush=True) self._git_command("--git-dir", self.bare_repo, "--work-tree", temp_worktree, "add", "-A") self._git_command("--git-dir", self.bare_repo, "--work-tree", temp_worktree, "commit", "-m", message, "--allow-empty") commit_hash = self._git_command("--git-dir", self.bare_repo, "rev-parse", "HEAD") checkpoint_id = commit_hash[:8] if commit_hash and not commit_hash.startswith("error") else None # Save conversation snapshot with checkpoint_id if checkpoint_id and conversation_snapshot: snapshots[checkpoint_id] = conversation_snapshot with open(snapshot_file, 'w') as f: json.dump(snapshots, f, indent=2) return checkpoint_id except Exception as e: return None def list_checkpoints(self, limit=10, show_all=False): """List recent checkpoints for current session Args: limit: Maximum number of checkpoints to show show_all: If True, show all sessions; if False, only show current session """ if not self.session_id and not show_all: return [] try: if show_all: # Show all branches args = ["--git-dir", self.bare_repo, "log", f"--max-count={limit}", "--oneline", "--all"] else: # Show only current session's branch branch_name = self._get_branch_name() args = ["--git-dir", self.bare_repo, "log", f"--max-count={limit}", "--oneline", branch_name] log = self._git_command(*args) if log and not log.startswith("error"): return [line.split(" ", 1) for line in log.split("\n") if line] return [] except: return [] def restore_checkpoint(self, checkpoint_id, session_baseline_files): """Restore files to checkpoint state and reset current session's branch This properly handles files that were added after the checkpoint by: 1. Restoring files that exist in checkpoint 2. Restoring files to baseline if they don't exist in checkpoint Args: checkpoint_id: Checkpoint hash to restore to session_baseline_files: Dict of baseline file states from session Returns: tuple: (success: bool, conversation_snapshot: dict or None) """ if not self.session_id: return False, None branch_name = self._get_branch_name() temp_worktree = os.path.join(self.nanocode_dir, f"temp_worktree_{self.session_id}") try: # Checkout branch first self._git_command("--git-dir", self.bare_repo, "--work-tree", temp_worktree, "checkout", branch_name, "-f") # Reset branch to checkpoint (discards future commits on this branch) self._git_command("--git-dir", self.bare_repo, "reset", "--hard", checkpoint_id) # Checkout to temp worktree self._git_command("--git-dir", self.bare_repo, "--work-tree", temp_worktree, "checkout", checkpoint_id, "-f") # Get list of files in checkpoint files_in_checkpoint_str = self._git_command( "--git-dir", self.bare_repo, "ls-tree", "-r", "--name-only", checkpoint_id ) files_in_checkpoint = set() if files_in_checkpoint_str and not files_in_checkpoint_str.startswith("error"): files_in_checkpoint = set(f for f in files_in_checkpoint_str.split('\n') if f.strip()) print(f"\n{DIM}Files in checkpoint: {len(files_in_checkpoint)}{RESET}") print(f"{DIM}Files modified in session: {len(session_baseline_files)}{RESET}\n") # Process each file that was modified in this session for filepath, baseline_source in session_baseline_files.items(): # Convert to relative path for comparison if os.path.isabs(filepath): rel_filepath = os.path.relpath(filepath, self.project_path) else: rel_filepath = filepath # Normalize path: remove leading ./ normalized_rel_path = rel_filepath.lstrip('./') if normalized_rel_path in files_in_checkpoint: # File exists in checkpoint - restore from checkpoint src = os.path.join(temp_worktree, normalized_rel_path) dest = os.path.join(self.project_path, normalized_rel_path) dest_dir = os.path.dirname(dest) if dest_dir: os.makedirs(dest_dir, exist_ok=True) with open(src, 'rb') as s, open(dest, 'wb') as d: d.write(s.read()) print(f" {GREEN}✓{RESET} {filepath} {DIM}(from checkpoint){RESET}") else: # File doesn't exist in checkpoint - restore to baseline abs_filepath = filepath if os.path.isabs(filepath) else os.path.join(self.project_path, filepath) if baseline_source["type"] == "git": # Restore from project .git (use normalized path for git) result = subprocess.run( ["git", "checkout", baseline_source["commit"], "--", normalized_rel_path], cwd=self.project_path, capture_output=True, text=True ) if result.returncode == 0: print(f" {CYAN}↺{RESET} {filepath} {DIM}(to baseline: git {baseline_source['commit'][:8]}){RESET}") elif baseline_source["type"] == "blob": # Restore from blob if self.restore_file_from_blob(baseline_source["hash"], abs_filepath): print(f" {CYAN}↺{RESET} {filepath} {DIM}(to baseline: blob {baseline_source['hash'][:8]}){RESET}") elif baseline_source["type"] == "new": # Delete new file if os.path.exists(abs_filepath): os.remove(abs_filepath) print(f" {YELLOW}✗{RESET} {filepath} {DIM}(deleted: was added after checkpoint){RESET}") # Load conversation snapshot snapshot_file = os.path.join(self.nanocode_dir, "conversation_snapshots.json") conversation_snapshot = None if os.path.exists(snapshot_file): with open(snapshot_file, 'r') as f: snapshots = json.load(f) conversation_snapshot = snapshots.get(checkpoint_id) return True, conversation_snapshot except Exception as e: print(f"{RED}Error during restore: {e}{RESET}") return False, None class Session: """Represents a conversation session""" def __init__(self, session_id=None): self.session_id = session_id or self._generate_session_id() self.messages = [] self.file_states = {} self.baseline_files = {} # Track original file versions for rollback self.metadata = { 'created_at': time.time(), 'last_active': time.time(), 'description': '', 'cwd': os.getcwd(), 'parent_checkpoint': None, # Track where this session branched from 'parent_session': None, # Track which session it branched from 'git_commit': None, # Project .git commit hash when session started 'git_branch': None, # Project .git branch when session started 'git_dirty': False, # Whether project had uncommitted changes } def _generate_session_id(self): """Generate unique session ID""" timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') random_suffix = ''.join(random.choices('0123456789abcdef', k=4)) return f"{timestamp}_{random_suffix}" def _get_project_git_info(self): """Get current project git state""" try: cwd = self.metadata.get('cwd', os.getcwd()) # Get current commit commit = subprocess.run( ["git", "rev-parse", "HEAD"], cwd=cwd, capture_output=True, text=True, check=True ).stdout.strip() # Get current branch branch = subprocess.run( ["git", "rev-parse", "--abbrev-ref", "HEAD"], cwd=cwd, capture_output=True, text=True, check=True ).stdout.strip() # Check if dirty status = subprocess.run( ["git", "status", "--porcelain"], cwd=cwd, capture_output=True, text=True, check=True ).stdout.strip() return { 'git_commit': commit[:8], 'git_branch': branch, 'git_dirty': bool(status) } except: return None def capture_git_state(self): """Capture current project git state into metadata""" git_info = self._get_project_git_info() if git_info: self.metadata.update(git_info) def track_file_state(self, filepath): """Track file state for conflict detection""" if os.path.exists(filepath): with open(filepath, 'rb') as f: content = f.read() file_hash = hashlib.md5(content).hexdigest() self.file_states[filepath] = { 'hash': file_hash, 'mtime': os.path.getmtime(filepath), 'size': len(content) } def detect_conflicts(self): """Detect if tracked files have been modified outside this session Returns: list: List of conflicted file paths """ conflicts = [] for filepath, saved_state in self.file_states.items(): if os.path.exists(filepath): with open(filepath, 'rb') as f: content = f.read() current_hash = hashlib.md5(content).hexdigest() current_mtime = os.path.getmtime(filepath) # Check if file has changed if (current_hash != saved_state['hash'] or current_mtime != saved_state['mtime']): conflicts.append(filepath) else: # File was deleted conflicts.append(f"{filepath} (deleted)") return conflicts def update_file_states(self): """Update all tracked file states to current state""" for filepath in list(self.file_states.keys()): if os.path.exists(filepath): self.track_file_state(filepath) else: # Remove deleted files from tracking del self.file_states[filepath] def to_dict(self): """Serialize to dict""" return { 'session_id': self.session_id, 'messages': self.messages, 'file_states': self.file_states, 'baseline_files': self.baseline_files, 'metadata': self.metadata } @staticmethod def from_dict(data): """Deserialize from dict""" session = Session(session_id=data['session_id']) session.messages = data.get('messages', []) session.file_states = data.get('file_states', {}) session.baseline_files = data.get('baseline_files', {}) session.metadata = data.get('metadata', {}) return session class SessionManager: """Manage multiple sessions""" def __init__(self, project_path): self.project_path = project_path self.sessions_dir = os.path.join(project_path, ".nanocode", "sessions") self.current_session = None self.checkpoint_manager = CheckpointManager(project_path) self.parent_commit_for_next_checkpoint = None # Track parent for first checkpoint os.makedirs(self.sessions_dir, exist_ok=True) def save_baseline_if_needed(self, filepath): """Save file's baseline version before first modification This is called before write/edit operations to preserve the original state. """ if not self.current_session: return # Already saved if filepath in self.current_session.baseline_files: return print(f"{DIM}[LOG] Saving baseline for: {filepath}{RESET}", flush=True) # Get file info from project .git git_info = self.checkpoint_manager.get_file_git_info(filepath) if git_info: # File is tracked by project .git if git_info["has_changes"]: # Has local changes - save to blob blob_hash = self.checkpoint_manager.save_file_to_blob(filepath) if blob_hash: self.current_session.baseline_files[filepath] = { "type": "blob", "hash": blob_hash } print(f"{DIM}[LOG] Saved dirty file to blob: {blob_hash[:8]}{RESET}", flush=True) else: # Clean - just record commit self.current_session.baseline_files[filepath] = { "type": "git", "commit": git_info["commit"] } print(f"{DIM}[LOG] Recorded git commit: {git_info['commit'][:8]}{RESET}", flush=True) else: # Untracked file or no .git if os.path.exists(filepath): # Save existing untracked file to blob blob_hash = self.checkpoint_manager.save_file_to_blob(filepath) if blob_hash: self.current_session.baseline_files[filepath] = { "type": "blob", "hash": blob_hash } print(f"{DIM}[LOG] Saved untracked file to blob: {blob_hash[:8]}{RESET}", flush=True) else: # New file - mark as new self.current_session.baseline_files[filepath] = { "type": "new" } print(f"{DIM}[LOG] Marked as new file{RESET}", flush=True) # Auto-save session to persist baseline_files self.save_session() def restore_baseline(self): """Restore all files to their baseline state Returns: bool: Success status """ if not self.current_session: return False if not self.current_session.baseline_files: print(f"{YELLOW}⚠ No baseline files to restore{RESET}") return False print(f"\n{BOLD}Restoring {len(self.current_session.baseline_files)} files to baseline...{RESET}\n") success_count = 0 for filepath, source in self.current_session.baseline_files.items(): try: # Normalize to absolute path abs_filepath = filepath if os.path.isabs(filepath) else os.path.join(self.project_path, filepath) # Get relative path for git operations rel_filepath = os.path.relpath(abs_filepath, self.project_path) if source["type"] == "git": # Restore from project .git (use relative path) result = subprocess.run( ["git", "checkout", source["commit"], "--", rel_filepath], cwd=self.project_path, capture_output=True, text=True ) if result.returncode == 0: print(f" {GREEN}✓{RESET} {filepath} {DIM}(from git {source['commit'][:8]}){RESET}") success_count += 1 else: print(f" {RED}✗{RESET} {filepath} {DIM}(git checkout failed){RESET}") elif source["type"] == "blob": # Restore from checkpoint.git blob (use absolute path) if self.checkpoint_manager.restore_file_from_blob(source["hash"], abs_filepath): print(f" {GREEN}✓{RESET} {filepath} {DIM}(from blob {source['hash'][:8]}){RESET}") success_count += 1 else: print(f" {RED}✗{RESET} {filepath} {DIM}(blob restore failed){RESET}") elif source["type"] == "new": # Delete new file (use absolute path) if os.path.exists(abs_filepath): os.remove(abs_filepath) print(f" {GREEN}✓{RESET} {filepath} {DIM}(deleted new file){RESET}") success_count += 1 else: print(f" {DIM}○{RESET} {filepath} {DIM}(already deleted){RESET}") success_count += 1 except Exception as e: print(f" {RED}✗{RESET} {filepath} {DIM}(error: {e}){RESET}") print(f"\n{GREEN}✓ Restored {success_count}/{len(self.current_session.baseline_files)} files{RESET}") return success_count > 0 def create_session(self, description="", parent_checkpoint=None, parent_session=None): """Create new session Args: description: Session description parent_checkpoint: Checkpoint ID this session branches from parent_session: Session ID this session branches from """ session = Session() session.metadata['description'] = description session.metadata['parent_checkpoint'] = parent_checkpoint session.metadata['parent_session'] = parent_session # Capture project git state session.capture_git_state() self.current_session = session # Set checkpoint manager to use this session self.checkpoint_manager.set_session(session.session_id) # Store parent commit for first checkpoint self.parent_commit_for_next_checkpoint = parent_checkpoint self.save_session() return session def save_session(self): """Save current session to disk""" if not self.current_session: return self.current_session.metadata['last_active'] = time.time() session_file = os.path.join( self.sessions_dir, f"{self.current_session.session_id}.json" ) with open(session_file, 'w') as f: json.dump(self.current_session.to_dict(), f, indent=2) def load_session(self, session_id): """Load session from disk""" session_file = os.path.join(self.sessions_dir, f"{session_id}.json") if not os.path.exists(session_file): return None with open(session_file, 'r') as f: data = json.load(f) session = Session.from_dict(data) self.current_session = session # Set checkpoint manager to use this session self.checkpoint_manager.set_session(session.session_id) return session def list_sessions(self): """List all sessions""" sessions = [] if not os.path.exists(self.sessions_dir): return sessions for filename in os.listdir(self.sessions_dir): if filename.endswith('.json'): filepath = os.path.join(self.sessions_dir, filename) try: with open(filepath, 'r') as f: data = json.load(f) sessions.append({ 'session_id': data['session_id'], 'metadata': data['metadata'], 'message_count': len(data.get('messages', [])), }) except: pass return sorted(sessions, key=lambda x: x['metadata'].get('last_active', 0), reverse=True) def load_last_session(self): """Load the most recent session""" sessions = self.list_sessions() if sessions: return self.load_session(sessions[0]['session_id']) return None def handle_checkpoint_command(parts, session_manager, files_modified): """Handle /checkpoint or /c commands Returns: messages: New messages list if conversation was restored, None otherwise """ # Default to list if no subcommand if len(parts) < 2: parts.append("list") cmd = parts[1] # If cmd looks like a commit hash (7-8 hex chars), treat as restore if len(cmd) >= 7 and len(cmd) <= 8 and all(c in '0123456789abcdef' for c in cmd.lower()): cmd = "restore" checkpoint_id = parts[1] else: checkpoint_id = None if cmd == "baseline" or cmd == "base": # Restore all files to baseline (session start state) if not session_manager.current_session.baseline_files: print(f"{YELLOW}⚠ No baseline: no files have been modified in this session{RESET}") return None print(f"{YELLOW}⚠ This will restore all modified files to their original state{RESET}") print(f"{YELLOW}⚠ Files to restore: {len(session_manager.current_session.baseline_files)}{RESET}") # Show files print(f"\n{DIM}Files:{RESET}") for filepath in list(session_manager.current_session.baseline_files.keys())[:10]: print(f" {DIM}• {filepath}{RESET}") if len(session_manager.current_session.baseline_files) > 10: print(f" {DIM}... and {len(session_manager.current_session.baseline_files) - 10} more{RESET}") confirm = input(f"\n{BOLD}Continue? (y/N): {RESET}").strip().lower() if confirm != 'y': print(f"{DIM}Cancelled{RESET}") return None success = session_manager.restore_baseline() if success: # Clear conversation print(f"{GREEN}✓ Conversation cleared{RESET}") return [] else: return None elif cmd == "list" or cmd == "all" or cmd == "--all": show_all = (cmd == "all" or "--all" in parts) if show_all: # Show git graph of all branches print(f"\n{BOLD}📍 Checkpoint Graph:{RESET}\n") # Use git log --graph --all to show the tree # Format: %h = short hash, %d = ref names, %s = subject, %ar = relative date graph_output = session_manager.checkpoint_manager._git_command( "--git-dir", session_manager.checkpoint_manager.bare_repo, "log", "--graph", "--all", "--oneline", "--format=%h %s (%ar)", "-20" ) if graph_output and not graph_output.startswith("error"): # Also get branch info for each commit branches_output = session_manager.checkpoint_manager._git_command( "--git-dir", session_manager.checkpoint_manager.bare_repo, "branch", "-a", "--contains" ) # Parse and display for line in graph_output.split('\n'): if not line.strip(): continue # Extract commit hash match = re.search(r'\b([0-9a-f]{7,8})\b', line) if match: commit_hash = match.group(1) # Get branches containing this commit branch_info = session_manager.checkpoint_manager._git_command( "--git-dir", session_manager.checkpoint_manager.bare_repo, "branch", "-a", "--contains", commit_hash ) # Extract session names from branches session_names = [] if branch_info and not branch_info.startswith("error"): for branch_line in branch_info.split('\n'): branch_line = branch_line.strip().lstrip('* ') if branch_line.startswith('session_'): # Shorten session name: session_20260130_103323_f7 -> s:20260130_103323_f7 session_short = 's:' + branch_line[8:] # Remove 'session_' prefix session_names.append(session_short) # Highlight commit hash line = line.replace(commit_hash, f"{CYAN}{commit_hash}{RESET}") # Add session info if found if session_names: # Insert session names after commit hash session_str = f"{GREEN}[{', '.join(session_names[:2])}]{RESET}" line = line.replace(commit_hash + f"{RESET}", commit_hash + f"{RESET} {session_str}") print(f" {line}") print() else: print(f"{DIM}No checkpoints yet{RESET}\n") print(f"{DIM}Restore: /c <hash>{RESET}") return None else: # Show current session's checkpoints checkpoints = session_manager.checkpoint_manager.list_checkpoints(show_all=False) if not checkpoints: print(f"{DIM}No checkpoints yet{RESET}") return None print(f"\n{BOLD}📍 Checkpoints:{RESET}\n") # Get checkpoint details from git log with timestamp for commit_hash, message in checkpoints[:10]: # Show first 10 (already newest first from git log) # Try to get timestamp from git timestamp_str = session_manager.checkpoint_manager._git_command( "--git-dir", session_manager.checkpoint_manager.bare_repo, "log", "-1", "--format=%ar", commit_hash ) if timestamp_str.startswith("error"): timestamp_str = "" # Get modified files files_str = session_manager.checkpoint_manager._git_command( "--git-dir", session_manager.checkpoint_manager.bare_repo, "diff-tree", "--no-commit-id", "--name-only", "-r", commit_hash ) files = [] if files_str and not files_str.startswith("error"): files = [f.strip() for f in files_str.split('\n') if f.strip()] # Format: hash | time ago | message time_part = f"{DIM}{timestamp_str}{RESET}" if timestamp_str else "" print(f" {CYAN}{commit_hash}{RESET} {time_part}") print(f" {DIM}└─{RESET} {message}") if files: files_display = ", ".join(files[:3]) if len(files) > 3: files_display += f" +{len(files)-3} more" print(f" {DIM} Files: {files_display}{RESET}") print() print(f"{DIM}Tip: Use '/c all' or '/ca' to see git graph{RESET}") print(f"{DIM}Restore: /c <hash>{RESET}") return None elif cmd == "restore": if not checkpoint_id: print(f"{RED}Usage: /c <checkpoint_id>{RESET}") return None print(f"{YELLOW}⚠ This will restore files AND conversation to checkpoint {checkpoint_id}{RESET}") print(f"{YELLOW}⚠ Future checkpoints will be discarded from history{RESET}") confirm = input(f"{BOLD}Continue? (y/N): {RESET}").strip().lower() if confirm != 'y': print(f"{DIM}Cancelled{RESET}") return None success, conversation_snapshot = session_manager.checkpoint_manager.restore_checkpoint( checkpoint_id, session_manager.current_session.baseline_files ) if success: print(f"{GREEN}✓ Restored files to checkpoint {checkpoint_id}{RESET}") if conversation_snapshot: print(f"{GREEN}✓ Restored conversation ({len(conversation_snapshot)} messages){RESET}") return conversation_snapshot else: print(f"{YELLOW}⚠ No conversation snapshot found for this checkpoint{RESET}") return None else: print(f"{RED}✗ Failed to restore checkpoint{RESET}") return None else: print(f"{RED}Unknown command: {cmd}{RESET}") return None if __name__ == "__main__": main()