git_rebase_ai.py 36 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888
  1. import subprocess
  2. import google.generativeai as genai
  3. import os
  4. import argparse
  5. import sys
  6. import datetime
  7. import re
  8. import logging
  9. import tempfile
  10. # --- Configuration ---
  11. # Configure logging
  12. logging.basicConfig(level=logging.WARN, format="%(levelname)s: %(message)s")
  13. # Attempt to get API key from environment variable
  14. API_KEY = os.getenv("GEMINI_API_KEY")
  15. if not API_KEY:
  16. logging.error("GEMINI_API_KEY environment variable not set.")
  17. logging.error(
  18. "Please obtain an API key from Google AI Studio (https://aistudio.google.com/app/apikey)"
  19. )
  20. logging.error("and set it as an environment variable:")
  21. logging.error(" export GEMINI_API_KEY='YOUR_API_KEY' (Linux/macOS)")
  22. logging.error(" set GEMINI_API_KEY=YOUR_API_KEY (Windows CMD)")
  23. logging.error(" $env:GEMINI_API_KEY='YOUR_API_KEY' (Windows PowerShell)")
  24. sys.exit(1)
  25. # Configure the Gemini AI Client
  26. try:
  27. genai.configure(api_key=API_KEY)
  28. # Use a model suitable for complex reasoning like code analysis.
  29. # Adjust model name if needed (e.g., 'gemini-1.5-flash-latest').
  30. MODEL_NAME = os.getenv("GEMINI_MODEL")
  31. if not MODEL_NAME:
  32. logging.error("GEMINI_MODEL environment variable not set.")
  33. logging.error(
  34. "Please set the desired Gemini model name (e.g., 'gemini-1.5-flash-latest')."
  35. )
  36. logging.error(" export GEMINI_MODEL='gemini-1.5-flash-latest' (Linux/macOS)")
  37. logging.error(" set GEMINI_MODEL=gemini-1.5-flash-latest (Windows CMD)")
  38. logging.error(
  39. " $env:GEMINI_MODEL='gemini-1.5-flash-latest' (Windows PowerShell)"
  40. )
  41. sys.exit(1)
  42. model = genai.GenerativeModel(MODEL_NAME)
  43. logging.info(f"Using Gemini model: {MODEL_NAME}")
  44. except Exception as e:
  45. logging.error(f"Error configuring Gemini AI: {e}")
  46. sys.exit(1)
  47. # --- Git Helper Functions ---
  48. def run_git_command(command_list, check=True, capture_output=True, env=None):
  49. """
  50. Runs a Git command as a list of arguments and returns its stdout.
  51. Handles errors and returns None on failure if check=True.
  52. Allows passing environment variables.
  53. """
  54. full_command = []
  55. try:
  56. full_command = ["git"] + command_list
  57. logging.debug(f"Running command: {' '.join(full_command)}")
  58. cmd_env = os.environ.copy()
  59. if env:
  60. cmd_env.update(env)
  61. result = subprocess.run(
  62. full_command,
  63. check=check,
  64. capture_output=capture_output,
  65. text=True,
  66. encoding="utf-8",
  67. errors="replace",
  68. env=cmd_env,
  69. )
  70. logging.debug(f"Command successful. Output:\n{result.stdout[:200]}...")
  71. return result.stdout.strip() if capture_output else ""
  72. except subprocess.CalledProcessError as e:
  73. logging.error(f"Error executing Git command: {' '.join(full_command)}")
  74. stderr_safe = (
  75. e.stderr.strip().encode("utf-8", "replace").decode("utf-8")
  76. if e.stderr
  77. else ""
  78. )
  79. stdout_safe = (
  80. e.stdout.strip().encode("utf-8", "replace").decode("utf-8")
  81. if e.stdout
  82. else ""
  83. )
  84. logging.error(f"Exit Code: {e.returncode}")
  85. if stderr_safe:
  86. logging.error(f"Stderr: {stderr_safe}")
  87. if stdout_safe:
  88. logging.error(f"Stdout: {stdout_safe}")
  89. return None
  90. except FileNotFoundError:
  91. logging.error(
  92. "Error: 'git' command not found. Is Git installed and in your PATH?"
  93. )
  94. sys.exit(1)
  95. except Exception as e:
  96. logging.error(f"An unexpected error occurred running git: {e}")
  97. return None
  98. def check_git_repository():
  99. """Checks if the current directory is the root of a Git repository."""
  100. output = run_git_command(["rev-parse", "--is-inside-work-tree"])
  101. return output == "true"
  102. def get_current_branch():
  103. """Gets the current active Git branch name."""
  104. return run_git_command(["rev-parse", "--abbrev-ref", "HEAD"])
  105. def create_backup_branch(branch_name):
  106. """Creates a timestamped backup branch from the given branch name."""
  107. timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
  108. backup_branch_name = f"{branch_name}-backup-{timestamp}"
  109. logging.info(
  110. f"Attempting to create backup branch: {backup_branch_name} from {branch_name}"
  111. )
  112. output = run_git_command(["branch", backup_branch_name, branch_name])
  113. if output is not None:
  114. logging.info(f"Successfully created backup branch: {backup_branch_name}")
  115. return backup_branch_name
  116. else:
  117. logging.error("Failed to create backup branch.")
  118. return None
  119. def get_commit_range(upstream_ref, current_branch):
  120. """
  121. Determines the commit range (merge_base..current_branch).
  122. Returns the range string and the merge base hash.
  123. """
  124. logging.info(
  125. f"Finding merge base between '{upstream_ref}' and '{current_branch}'..."
  126. )
  127. merge_base = run_git_command(["merge-base", upstream_ref, current_branch])
  128. if not merge_base:
  129. logging.error(
  130. f"Could not find merge base between '{upstream_ref}' and '{current_branch}'."
  131. )
  132. logging.error(
  133. f"Ensure '{upstream_ref}' is a valid reference (branch, commit, tag)"
  134. )
  135. logging.error("and that it has been fetched (e.g., 'git fetch origin').")
  136. return None, None # Indicate failure
  137. logging.info(f"Found merge base: {merge_base}")
  138. commit_range = f"{merge_base}..{current_branch}"
  139. return commit_range, merge_base
  140. def get_commits_in_range(commit_range):
  141. """Gets a list of commit hashes and subjects in the specified range (oldest first)."""
  142. log_output = run_git_command(
  143. ["log", "--pretty=format:%h %s", "--reverse", commit_range]
  144. )
  145. if log_output is not None:
  146. commits = log_output.splitlines()
  147. logging.info(f"Found {len(commits)} commits in range {commit_range}.")
  148. return commits
  149. return [] # Return empty list on failure or no commits
  150. def get_changed_files_in_range(commit_range):
  151. """
  152. Gets a list of files changed in the specified range and generates
  153. a simple directory structure string representation.
  154. """
  155. diff_output = run_git_command(["diff", "--name-only", commit_range])
  156. if diff_output is not None:
  157. files = diff_output.splitlines()
  158. logging.info(f"Found {len(files)} changed files in range {commit_range}.")
  159. # Basic tree structure representation
  160. tree = {}
  161. for file_path in files:
  162. parts = file_path.replace("\\", "/").split("/")
  163. node = tree
  164. for i, part in enumerate(parts):
  165. if not part:
  166. continue
  167. if i == len(parts) - 1:
  168. node[part] = "file"
  169. else:
  170. if part not in node:
  171. node[part] = {}
  172. if isinstance(node[part], dict):
  173. node = node[part]
  174. else:
  175. logging.warning(
  176. f"Path conflict building file tree for: {file_path}"
  177. )
  178. break
  179. def format_tree(d, indent=0):
  180. lines = []
  181. for key, value in sorted(d.items()):
  182. prefix = " " * indent
  183. if isinstance(value, dict):
  184. lines.append(f"{prefix}📁 {key}/")
  185. lines.extend(format_tree(value, indent + 1))
  186. else:
  187. lines.append(f"{prefix}📄 {key}")
  188. return lines
  189. tree_str = "\n".join(format_tree(tree))
  190. return tree_str, files
  191. return "", []
  192. def get_diff_in_range(commit_range):
  193. """Gets the combined diffstat and patch for the specified range."""
  194. diff_output = run_git_command(["diff", "--patch-with-stat", commit_range])
  195. if diff_output is not None:
  196. logging.info(
  197. f"Generated diff for range {commit_range} (length: {len(diff_output)} chars)."
  198. )
  199. else:
  200. logging.warning(f"Could not generate diff for range {commit_range}.")
  201. return diff_output if diff_output is not None else ""
  202. def get_file_content_at_commit(commit_hash, file_path):
  203. """Gets the content of a specific file at a specific commit hash."""
  204. logging.info(f"Fetching content of '{file_path}' at commit {commit_hash[:7]}...")
  205. content = run_git_command(["show", f"{commit_hash}:{file_path}"])
  206. if content is None:
  207. logging.warning(
  208. f"Could not retrieve content for {file_path} at {commit_hash[:7]}."
  209. )
  210. return None
  211. return content
  212. # --- AI Interaction ---
  213. def generate_fixup_suggestion_prompt(
  214. commit_range, merge_base, commits, file_structure, diff
  215. ):
  216. """
  217. Creates a prompt asking the AI specifically to identify potential
  218. fixup candidates within the commit range.
  219. Returns suggestions in a parsable format.
  220. """
  221. commit_list_str = (
  222. "\n".join([f"- {c}" for c in commits]) if commits else "No commits in range."
  223. )
  224. prompt = f"""
  225. You are an expert Git assistant. Your task is to analyze the provided Git commit history and identify commits within the range `{commit_range}` that should be combined using `fixup` during an interactive rebase (`git rebase -i {merge_base}`).
  226. **Goal:** Identify commits that are minor corrections or direct continuations of the immediately preceding commit, where the commit message can be discarded.
  227. **Git Commit Message Conventions (for context):**
  228. * Subject: Imperative, < 50 chars, capitalized, no period. Use types like `feat:`, `fix:`, `refactor:`, etc.
  229. * Body: Explain 'what' and 'why', wrap at 72 chars.
  230. **Provided Context:**
  231. 1. **Commit Range:** `{commit_range}`
  232. 2. **Merge Base Hash:** `{merge_base}`
  233. 3. **Commits in Range (Oldest First - Short Hash & Subject):**
  234. ```
  235. {commit_list_str}
  236. ```
  237. 4. **Changed Files Structure in Range:**
  238. ```
  239. {file_structure if file_structure else "No files changed or unable to list."}
  240. ```
  241. 5. **Combined Diff for the Range (`git diff --patch-with-stat {commit_range}`):**
  242. ```diff
  243. {diff if diff else "No differences found or unable to get diff."}
  244. ```
  245. **Instructions:**
  246. 1. Analyze the commits, their messages, the changed files, and the diff.
  247. 2. Identify commits from the list that are strong candidates for being combined into their **immediately preceding commit** using `fixup` (combine changes, discard message). Focus on small fixes, typo corrections, or direct continuations where the commit message isn't valuable.
  248. 3. For each suggestion, output *only* a line in the following format:
  249. `FIXUP: <hash_to_fixup> INTO <preceding_hash>`
  250. Use the short commit hashes provided in the commit list.
  251. 4. Provide *only* lines in the `FIXUP:` format. Do not include explanations, introductory text, or any other formatting. If no fixups are suggested, output nothing.
  252. **Example Output:**
  253. ```text
  254. FIXUP: hash2 INTO hash1
  255. FIXUP: hash5 INTO hash4
  256. ```
  257. 5. **File Content Request:** If you absolutely need the content of specific files *at specific commits* to confidently determine if they should be fixed up, ask for them clearly ONCE. List the files using this exact format at the end of your response:
  258. `REQUEST_FILES: [commit_hash1:path/to/file1.py, commit_hash2:another/path/file2.js]`
  259. Use the short commit hashes provided in the commit list. Do *not* ask for files unless essential for *this specific task* of identifying fixup candidates.
  260. Now, analyze the provided context and generate *only* the `FIXUP:` lines or `REQUEST_FILES:` line.
  261. """
  262. return prompt
  263. def parse_fixup_suggestions(ai_response_text, commits_in_range):
  264. """Parses AI response for FIXUP: lines and validates hashes."""
  265. fixup_pairs = []
  266. commit_hashes = {
  267. c.split()[0] for c in commits_in_range
  268. } # Set of valid short hashes
  269. for line in ai_response_text.splitlines():
  270. line = line.strip()
  271. if line.startswith("FIXUP:"):
  272. match = re.match(r"FIXUP:\s*(\w+)\s+INTO\s+(\w+)", line, re.IGNORECASE)
  273. if match:
  274. fixup_hash = match.group(1)
  275. target_hash = match.group(2)
  276. # Validate that both hashes were in the original commit list
  277. if fixup_hash in commit_hashes and target_hash in commit_hashes:
  278. fixup_pairs.append({"fixup": fixup_hash, "target": target_hash})
  279. logging.debug(
  280. f"Parsed fixup suggestion: {fixup_hash} into {target_hash}"
  281. )
  282. else:
  283. logging.warning(
  284. f"Ignoring invalid fixup suggestion (hash not in range): {line}"
  285. )
  286. else:
  287. logging.warning(f"Could not parse FIXUP line: {line}")
  288. return fixup_pairs
  289. # --- request_files_from_user function remains the same ---
  290. def request_files_from_user(requested_files_str, commits_in_range):
  291. """
  292. Parses AI request string "REQUEST_FILES: [hash:path, ...]", verifies hashes,
  293. asks user permission, fetches file contents, and returns formatted context.
  294. """
  295. file_requests = []
  296. try:
  297. content_match = re.search(
  298. r"REQUEST_FILES:\s*\[(.*)\]", requested_files_str, re.IGNORECASE | re.DOTALL
  299. )
  300. if not content_match:
  301. logging.warning("Could not parse file request format from AI response.")
  302. return None, None
  303. items_str = content_match.group(1).strip()
  304. if not items_str:
  305. logging.info("AI requested files but the list was empty.")
  306. return None, None
  307. items = [item.strip() for item in items_str.split(",") if item.strip()]
  308. commit_hash_map = {c.split()[0]: c.split()[0] for c in commits_in_range}
  309. for item in items:
  310. if ":" not in item:
  311. logging.warning(
  312. f"Invalid format in requested file item (missing ':'): {item}"
  313. )
  314. continue
  315. commit_hash, file_path = item.split(":", 1)
  316. commit_hash = commit_hash.strip()
  317. file_path = file_path.strip()
  318. if commit_hash not in commit_hash_map:
  319. logging.warning(
  320. f"AI requested file for unknown/out-of-range commit hash '{commit_hash}'. Skipping."
  321. )
  322. continue
  323. file_requests.append({"hash": commit_hash, "path": file_path})
  324. except Exception as e:
  325. logging.error(f"Error parsing requested files string: {e}")
  326. return None, None
  327. if not file_requests:
  328. logging.info("No valid file requests found after parsing AI response.")
  329. return None, None
  330. print("\n----------------------------------------")
  331. print("❓ AI Request for File Content ❓")
  332. print("----------------------------------------")
  333. print("The AI needs the content of the following files at specific commits")
  334. print("to provide more accurate fixup suggestions:")
  335. files_to_fetch = []
  336. for i, req in enumerate(file_requests):
  337. print(f" {i + 1}. File: '{req['path']}' at commit {req['hash']}")
  338. files_to_fetch.append(req)
  339. if not files_to_fetch:
  340. print("\nNo valid files to fetch based on the request.")
  341. return None, None
  342. print("----------------------------------------")
  343. while True:
  344. try:
  345. answer = (
  346. input("Allow fetching these file contents? (yes/no): ").lower().strip()
  347. )
  348. except EOFError:
  349. logging.warning("Input stream closed. Assuming 'no'.")
  350. answer = "no"
  351. if answer == "yes":
  352. logging.info("User approved fetching file content.")
  353. fetched_content_list = []
  354. for req in files_to_fetch:
  355. content = get_file_content_at_commit(req["hash"], req["path"])
  356. if content is not None:
  357. fetched_content_list.append(
  358. f"--- Content of '{req['path']}' at commit {req['hash']} ---\n"
  359. f"```\n{content}\n```\n"
  360. f"--- End Content for {req['path']} at {req['hash']} ---"
  361. )
  362. else:
  363. fetched_content_list.append(
  364. f"--- Could not fetch content of '{req['path']}' at commit {req['hash']} ---"
  365. )
  366. return "\n\n".join(fetched_content_list), requested_files_str
  367. elif answer == "no":
  368. logging.info("User denied fetching file content.")
  369. return None, requested_files_str
  370. else:
  371. print("Please answer 'yes' or 'no'.")
  372. # --- Automatic Rebase Logic ---
  373. def create_rebase_editor_script(script_path, fixup_plan):
  374. """Creates the python script to be used by GIT_SEQUENCE_EDITOR."""
  375. # Create a set of hashes that need to be fixed up
  376. fixups_to_apply = {pair["fixup"] for pair in fixup_plan}
  377. script_content = f"""#!/usr/bin/env python3
  378. import sys
  379. import logging
  380. import re
  381. import os
  382. # Define log file path relative to the script itself
  383. log_file = __file__ + ".log"
  384. # Setup logging within the editor script to write to the log file
  385. logging.basicConfig(filename=log_file, filemode='w', level=logging.WARN, format="%(asctime)s - %(levelname)s: %(message)s")
  386. todo_file_path = sys.argv[1]
  387. logging.info(f"GIT_SEQUENCE_EDITOR script started for: {{todo_file_path}}")
  388. # Hashes that should be changed to 'fixup'
  389. fixups_to_apply = {fixups_to_apply!r}
  390. logging.info(f"Applying fixups for hashes: {{fixups_to_apply}}")
  391. new_lines = []
  392. try:
  393. with open(todo_file_path, 'r', encoding='utf-8') as f:
  394. lines = f.readlines()
  395. for line in lines:
  396. stripped_line = line.strip()
  397. # Skip comments and blank lines
  398. if not stripped_line or stripped_line.startswith('#'):
  399. new_lines.append(line)
  400. continue
  401. # Use regex for more robust parsing of todo lines (action hash ...)
  402. match = re.match(r"^(\w+)\s+([0-9a-fA-F]+)(.*)", stripped_line)
  403. if match:
  404. action = match.group(1).lower()
  405. commit_hash = match.group(2)
  406. rest_of_line = match.group(3)
  407. # Check if this commit should be fixed up
  408. if commit_hash in fixups_to_apply and action == 'pick':
  409. logging.info(f"Changing 'pick {{commit_hash}}' to 'fixup {{commit_hash}}'")
  410. # Replace 'pick' with 'fixup', preserving the rest of the line
  411. new_line = f'f {{commit_hash}}{{rest_of_line}}\\n'
  412. new_lines.append(new_line)
  413. else:
  414. # Keep the original line
  415. new_lines.append(line)
  416. else:
  417. # Keep lines that don't look like standard todo lines
  418. logging.warning(f"Could not parse todo line: {{stripped_line}}")
  419. new_lines.append(line)
  420. logging.info(f"Writing {{len(new_lines)}} lines back to {{todo_file_path}}")
  421. with open(todo_file_path, 'w', encoding='utf-8') as f:
  422. f.writelines(new_lines)
  423. logging.info("GIT_SEQUENCE_EDITOR script finished successfully.")
  424. sys.exit(0) # Explicitly exit successfully
  425. except Exception as e:
  426. logging.error(f"Error in GIT_SEQUENCE_EDITOR script: {{e}}", exc_info=True)
  427. sys.exit(1) # Exit with error code
  428. """
  429. try:
  430. with open(script_path, "w", encoding="utf-8") as f:
  431. f.write(script_content)
  432. # Make the script executable (important on Linux/macOS)
  433. os.chmod(script_path, 0o755)
  434. logging.info(f"Created GIT_SEQUENCE_EDITOR script: {script_path}")
  435. return True
  436. except Exception as e:
  437. logging.error(f"Failed to create GIT_SEQUENCE_EDITOR script: {e}")
  438. return False
  439. def attempt_auto_fixup(merge_base, fixup_plan):
  440. """Attempts to perform the rebase automatically applying fixups."""
  441. if not fixup_plan:
  442. logging.info("No fixup suggestions provided by AI. Skipping auto-rebase.")
  443. return True # Nothing to do, considered success
  444. # Use a temporary directory to hold the script and its log
  445. temp_dir = tempfile.mkdtemp(prefix="git_rebase_")
  446. editor_script_path = os.path.join(temp_dir, "rebase_editor.py")
  447. logging.debug(f"Temporary directory: {temp_dir}")
  448. logging.debug(f"Temporary editor script path: {editor_script_path}")
  449. try:
  450. if not create_rebase_editor_script(editor_script_path, fixup_plan):
  451. return False # Failed to create script
  452. # Prepare environment for the git command
  453. rebase_env = os.environ.copy()
  454. rebase_env["GIT_SEQUENCE_EDITOR"] = editor_script_path
  455. # Prevent Git from opening a standard editor for messages etc.
  456. # 'true' simply exits successfully, accepting default messages
  457. rebase_env["GIT_EDITOR"] = "true"
  458. print("\nAttempting automatic rebase with suggested fixups...")
  459. logging.info(f"Running: git rebase -i {merge_base}")
  460. # Run rebase non-interactively, check=False to handle failures manually
  461. rebase_result = run_git_command(
  462. ["rebase", "-i", merge_base],
  463. check=False, # Don't raise exception on failure, check exit code
  464. capture_output=True, # Capture output to see potential errors
  465. env=rebase_env,
  466. )
  467. # Check the result (run_git_command returns None on CalledProcessError)
  468. if rebase_result is not None:
  469. # Command finished, exit code was likely 0 (success)
  470. print("✅ Automatic fixup rebase completed successfully.")
  471. logging.info("Automatic fixup rebase seems successful.")
  472. return True
  473. else:
  474. # Command failed (non-zero exit code, run_git_command returned None)
  475. print("\n❌ Automatic fixup rebase failed.")
  476. print(
  477. " This likely means merge conflicts occurred or another rebase error happened."
  478. )
  479. logging.warning("Automatic fixup rebase failed. Aborting...")
  480. # Attempt to abort the failed rebase
  481. print(" Attempting to abort the failed rebase (`git rebase --abort`)...")
  482. # Run abort without capturing output, just check success/failure
  483. abort_result = run_git_command(
  484. ["rebase", "--abort"], check=False, capture_output=False
  485. )
  486. # run_git_command returns None on failure (CalledProcessError)
  487. if abort_result is not None:
  488. print(
  489. " Rebase aborted successfully. Your branch is back to its original state."
  490. )
  491. logging.info("Failed rebase aborted successfully.")
  492. else:
  493. print(" ⚠️ Failed to automatically abort the rebase.")
  494. print(" Please run `git rebase --abort` manually to clean up.")
  495. logging.error("Failed to automatically abort the rebase.")
  496. return False
  497. except Exception as e:
  498. logging.error(
  499. f"An unexpected error occurred during auto-fixup attempt: {e}",
  500. exc_info=True,
  501. )
  502. # Might need manual cleanup here too
  503. print("\n❌ An unexpected error occurred during the automatic fixup attempt.")
  504. print(
  505. " You may need to manually check your Git status and potentially run `git rebase --abort`."
  506. )
  507. return False
  508. finally:
  509. # Determine if rebase failed *before* potential cleanup errors
  510. # Note: rebase_result is defined in the outer scope of the try block
  511. rebase_failed = "rebase_result" in locals() and rebase_result is None
  512. # Check if we need to display the editor script log
  513. editor_log_path = editor_script_path + ".log"
  514. verbose_logging = logging.getLogger().isEnabledFor(logging.DEBUG)
  515. if (rebase_failed or verbose_logging) and os.path.exists(editor_log_path):
  516. try:
  517. with open(editor_log_path, "r", encoding="utf-8") as log_f:
  518. log_content = log_f.read()
  519. if log_content:
  520. print("\n--- Rebase Editor Script Log ---")
  521. print(log_content.strip())
  522. print("--- End Log ---")
  523. else:
  524. # Only log if verbose, otherwise it's just noise
  525. if verbose_logging:
  526. logging.debug(
  527. f"Rebase editor script log file was empty: {editor_log_path}"
  528. )
  529. except Exception as log_e:
  530. logging.warning(
  531. f"Could not read rebase editor script log file {editor_log_path}: {log_e}"
  532. )
  533. # Clean up the temporary directory and its contents
  534. if temp_dir and os.path.exists(temp_dir):
  535. try:
  536. if os.path.exists(editor_log_path):
  537. os.remove(editor_log_path)
  538. if os.path.exists(editor_script_path):
  539. os.remove(editor_script_path)
  540. os.rmdir(temp_dir)
  541. logging.debug(f"Cleaned up temporary directory: {temp_dir}")
  542. except OSError as e:
  543. logging.warning(
  544. f"Could not completely remove temporary directory {temp_dir}: {e}"
  545. )
  546. # --- Main Execution ---
  547. def main():
  548. """Main function to orchestrate Git analysis and AI interaction."""
  549. parser = argparse.ArgumentParser(
  550. description="Uses Gemini AI to suggest and automatically attempt Git 'fixup' operations.",
  551. formatter_class=argparse.ArgumentDefaultsHelpFormatter,
  552. )
  553. parser.add_argument(
  554. "upstream_ref",
  555. nargs="?",
  556. default="upstream/main",
  557. help="The upstream reference point or commit hash to compare against "
  558. "(e.g., 'origin/main', 'upstream/develop', specific_commit_hash). "
  559. "Ensure this reference exists and is fetched.",
  560. )
  561. # --- Argument Change ---
  562. parser.add_argument(
  563. "--instruct",
  564. action="store_true",
  565. help="Only show AI suggestions and instructions; disable automatic fixup attempt.",
  566. )
  567. parser.add_argument(
  568. "-v", "--verbose", action="store_true", help="Enable verbose debug logging."
  569. )
  570. args = parser.parse_args()
  571. if args.verbose:
  572. logging.getLogger().setLevel(logging.DEBUG)
  573. logging.debug("Verbose logging enabled.")
  574. if not check_git_repository():
  575. logging.error("This script must be run from within a Git repository.")
  576. sys.exit(1)
  577. current_branch = get_current_branch()
  578. if not current_branch:
  579. logging.error("Could not determine the current Git branch.")
  580. sys.exit(1)
  581. logging.info(f"Current branch: {current_branch}")
  582. upstream_ref = args.upstream_ref
  583. logging.info(f"Comparing against reference: {upstream_ref}")
  584. # --- Safety: Create Backup Branch ---
  585. # Always create backup, especially if attempting auto-rebase
  586. backup_branch = create_backup_branch(current_branch)
  587. if not backup_branch:
  588. try:
  589. confirm = input(
  590. "⚠️ Failed to create backup branch. Continue without backup? (yes/no): "
  591. ).lower()
  592. except EOFError:
  593. logging.warning("Input stream closed. Aborting.")
  594. confirm = "no"
  595. if confirm != "yes":
  596. logging.info("Aborting.")
  597. sys.exit(1)
  598. else:
  599. logging.warning("Proceeding without a backup branch. Be careful!")
  600. else:
  601. print("-" * 40)
  602. print(f"✅ Backup branch created: {backup_branch}")
  603. print(" If anything goes wrong, you can restore using:")
  604. print(f" git checkout {current_branch}")
  605. print(f" git reset --hard {backup_branch}")
  606. print("-" * 40)
  607. # --- Gather Git Context ---
  608. print("\nGathering Git context...")
  609. commit_range, merge_base = get_commit_range(upstream_ref, current_branch)
  610. if not commit_range:
  611. sys.exit(1)
  612. logging.info(f"Analyzing commit range: {commit_range} (Merge Base: {merge_base})")
  613. commits = get_commits_in_range(commit_range)
  614. if not commits:
  615. logging.info(
  616. f"No commits found between '{merge_base}' and '{current_branch}'. Nothing to do."
  617. )
  618. sys.exit(0)
  619. file_structure, changed_files_list = get_changed_files_in_range(commit_range)
  620. diff = get_diff_in_range(commit_range)
  621. if not diff and not changed_files_list:
  622. logging.warning(
  623. f"No file changes or diff found between '{merge_base}' and '{current_branch}',"
  624. )
  625. logging.warning("even though commits exist. AI suggestions might be limited.")
  626. # Don't exit automatically, let AI try
  627. # --- Interact with AI ---
  628. print("\nGenerating prompt for AI fixup suggestions...")
  629. initial_prompt = generate_fixup_suggestion_prompt(
  630. commit_range, merge_base, commits, file_structure, diff
  631. )
  632. logging.debug("\n--- Initial AI Prompt Snippet ---")
  633. logging.debug(initial_prompt[:1000] + "...")
  634. logging.debug("--- End Prompt Snippet ---\n")
  635. print(f"Sending request to Gemini AI ({MODEL_NAME})...")
  636. ai_response_text = ""
  637. fixup_suggestions_text = "" # Store the raw suggestions for later display if needed
  638. try:
  639. convo = model.start_chat(history=[])
  640. response = convo.send_message(initial_prompt)
  641. ai_response_text = response.text
  642. # Loop for file requests
  643. while "REQUEST_FILES:" in ai_response_text.upper():
  644. logging.info("AI requested additional file content.")
  645. additional_context, original_request = request_files_from_user(
  646. ai_response_text, commits
  647. )
  648. if additional_context:
  649. logging.info("Sending fetched file content back to AI...")
  650. follow_up_prompt = f"""
  651. Okay, here is the content of the files you requested:
  652. {additional_context}
  653. Please use this new information to refine your **fixup suggestions** based on the original request and context. Provide the final list of `FIXUP: ...` lines now. Remember to *only* suggest fixup actions and output *only* `FIXUP:` lines. Do not ask for more files.
  654. """
  655. logging.debug("\n--- Follow-up AI Prompt Snippet ---")
  656. logging.debug(follow_up_prompt[:500] + "...")
  657. logging.debug("--- End Follow-up Snippet ---\n")
  658. response = convo.send_message(follow_up_prompt)
  659. ai_response_text = response.text
  660. else:
  661. logging.info(
  662. "Proceeding without providing files as requested by AI or user."
  663. )
  664. no_files_prompt = f"""
  665. I cannot provide the content for the files you requested ({original_request}).
  666. Please proceed with generating the **fixup suggestions** based *only* on the initial context (commit list, file structure, diff) I provided earlier. Make your best suggestions without the file content. Provide the final list of `FIXUP: ...` lines now. Remember to *only* suggest fixup actions.
  667. """
  668. logging.debug("\n--- No-Files AI Prompt ---")
  669. logging.debug(no_files_prompt)
  670. logging.debug("--- End No-Files Prompt ---\n")
  671. response = convo.send_message(no_files_prompt)
  672. ai_response_text = response.text
  673. break
  674. # Store the final AI response containing suggestions
  675. fixup_suggestions_text = ai_response_text.strip()
  676. # Parse the suggestions
  677. fixup_plan = parse_fixup_suggestions(fixup_suggestions_text, commits)
  678. if not fixup_plan:
  679. print("\n💡 AI did not suggest any specific fixup operations.")
  680. else:
  681. print("\n💡 --- AI Fixup Suggestions --- 💡")
  682. # Print the parsed plan for clarity
  683. for i, pair in enumerate(fixup_plan):
  684. print(
  685. f" {i + 1}. Fixup commit `{pair['fixup']}` into `{pair['target']}`"
  686. )
  687. print("💡 --- End AI Suggestions --- 💡")
  688. # --- Attempt Automatic Rebase or Show Instructions ---
  689. # --- Logic Change ---
  690. if not args.instruct: # Default behavior: attempt auto-fixup
  691. if fixup_plan:
  692. success = attempt_auto_fixup(merge_base, fixup_plan)
  693. if not success:
  694. # Failure message already printed by attempt_auto_fixup
  695. print("\n" + "=" * 60)
  696. print("🛠️ MANUAL REBASE REQUIRED 🛠️")
  697. print("=" * 60)
  698. print(
  699. "The automatic fixup rebase failed (likely due to conflicts)."
  700. )
  701. print("Please perform the rebase manually:")
  702. print(f" 1. Run: `git rebase -i {merge_base}`")
  703. print(
  704. " 2. In the editor, change 'pick' to 'f' (or 'fixup') for the commits"
  705. )
  706. print(
  707. " suggested by the AI above (and any other changes you want)."
  708. )
  709. print(" Original AI suggestions:")
  710. print(" ```text")
  711. # Print raw suggestions which might be easier to copy/paste
  712. print(
  713. fixup_suggestions_text
  714. if fixup_suggestions_text
  715. else " (No specific fixup lines found in AI response)"
  716. )
  717. print(" ```")
  718. print(" 3. Save the editor and resolve any conflicts Git reports.")
  719. print(
  720. " Use `git status`, edit files, `git add <files>`, `git rebase --continue`."
  721. )
  722. if backup_branch:
  723. print(f" 4. Remember backup branch: {backup_branch}")
  724. print("=" * 60)
  725. sys.exit(1) # Exit with error status after failure
  726. else:
  727. # Auto fixup succeeded
  728. print("\nBranch history has been modified by automatic fixups.")
  729. if backup_branch:
  730. print(
  731. f"Backup branch '{backup_branch}' still exists if needed."
  732. )
  733. else:
  734. print("\nNo automatic rebase attempted as AI suggested no fixups.")
  735. elif fixup_plan: # --instruct flag was used AND suggestions exist
  736. print("\n" + "=" * 60)
  737. print("📝 MANUAL REBASE INSTRUCTIONS (--instruct used) 📝")
  738. print("=" * 60)
  739. print("AI suggested the fixups listed above.")
  740. print("To apply them (or other changes):")
  741. print(f" 1. Run: `git rebase -i {merge_base}`")
  742. print(" 2. Edit the 'pick' lines in the editor based on the suggestions")
  743. print(" (changing 'pick' to 'f' or 'fixup').")
  744. print(" 3. Save the editor and follow Git's instructions.")
  745. if backup_branch:
  746. print(f" 4. Remember backup branch: {backup_branch}")
  747. print("=" * 60)
  748. # If --instruct and no fixup_plan, nothing specific needs to be printed here
  749. except Exception as e:
  750. logging.error(f"\nAn unexpected error occurred: {e}", exc_info=True)
  751. # Attempt to print feedback if available
  752. try:
  753. if response and hasattr(response, "prompt_feedback"):
  754. logging.error(f"AI Prompt Feedback: {response.prompt_feedback}")
  755. if response and hasattr(response, "candidates"):
  756. for candidate in response.candidates:
  757. logging.error(
  758. f"AI Candidate Finish Reason: {candidate.finish_reason}"
  759. )
  760. if hasattr(candidate, "safety_ratings"):
  761. logging.error(f"AI Safety Ratings: {candidate.safety_ratings}")
  762. except Exception as feedback_e:
  763. logging.error(
  764. f"Could not retrieve detailed feedback from AI response: {feedback_e}"
  765. )
  766. print("\n❌ An unexpected error occurred during the process.")
  767. print(" Please check the logs and your Git status.")
  768. print(" You may need to run `git rebase --abort` manually.")
  769. if __name__ == "__main__":
  770. main()