Update src/specify_cli/__init__.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: mbachorik

src/specify_cli/__init__.py

@@ -2872,30 +2872,83 @@ def extension_update(
                     # Remove any NEW command files created by failed install
                     # (files that weren't in the original backup)
                     try:
+                        # Use the backup directory's mtime as an approximate
+                        # "start of install" timestamp. Files modified after this
+                        # and not present in the backup are candidates for cleanup
+                        backup_timestamp = None
+                        try:
+                            backup_timestamp = Path(backup_base).stat().st_mtime
+                        except Exception:
+                            # If we can't stat the backup directory, we still fall
+                            # back to registry-based cleanup where possible.
+                            pass
+
                         new_registry_entry = manager.registry.get(extension_id)
-                        if new_registry_entry is None:
-                            new_registered_commands = {}
-                        else:
+                        if new_registry_entry is not None:
                             new_registered_commands = new_registry_entry.get("registered_commands", {})
-                        for agent_name, cmd_names in new_registered_commands.items():
-                            if agent_name not in registrar.AGENT_CONFIGS:
-                                continue
-                            agent_config = registrar.AGENT_CONFIGS[agent_name]
-                            commands_dir = project_root / agent_config["dir"]
-
-                            for cmd_name in cmd_names:
-                                cmd_file = commands_dir / f"{cmd_name}{agent_config['extension']}"
-                                # Delete if it exists and wasn't in our backup
-                                if cmd_file.exists() and str(cmd_file) not in backed_up_command_files:
-                                    cmd_file.unlink()
-
-                                # Also handle copilot prompt files
+
+                            for agent_name, cmd_names in new_registered_commands.items():
+                                if agent_name not in registrar.AGENT_CONFIGS:
+                                    continue
+                                agent_config = registrar.AGENT_CONFIGS[agent_name]
+                                commands_dir = project_root / agent_config["dir"]
+
+                                for cmd_name in cmd_names:
+                                    cmd_file = commands_dir / f"{cmd_name}{agent_config['extension']}"
+                                    # Delete if it exists and wasn't in our backup
+                                    if cmd_file.exists() and str(cmd_file) not in backed_up_command_files:
+                                        # Optionally use timestamp heuristic if available
+                                        if backup_timestamp is None or cmd_file.stat().st_mtime >= backup_timestamp:
+                                            cmd_file.unlink()
+
+                                    # Also handle copilot prompt files
+                                    if agent_name == "copilot":
+                                        prompt_file = project_root / ".github" / "prompts" / f"{cmd_name}.prompt.md"
+                                        if prompt_file.exists() and str(prompt_file) not in backed_up_command_files:
+                                            if backup_timestamp is None or prompt_file.stat().st_mtime >= backup_timestamp:
+                                                prompt_file.unlink()
+                        else:
+                            # Fallback: we have no new registry entry (e.g., failure
+                            # occurred before it was written). Best-effort cleanup by
+                            # scanning agent command directories for files that weren't
+                            # in the backup and appear to have been created/modified
+                            # after the backup was taken.
+                            for agent_name, agent_config in registrar.AGENT_CONFIGS.items():
+                                commands_dir = project_root / agent_config["dir"]
+                                if commands_dir.is_dir():
+                                    for cmd_file in commands_dir.glob(f"*{agent_config['extension']}"):
+                                        cmd_path_str = str(cmd_file)
+                                        if cmd_path_str in backed_up_command_files:
+                                            continue
+                                        if backup_timestamp is not None:
+                                            try:
+                                                if cmd_file.stat().st_mtime < backup_timestamp:
+                                                    continue
+                                            except OSError:
+                                                # If we can't stat the file, skip it
+                                                continue
+                                        cmd_file.unlink()
+
+                                # Also handle copilot prompt files in fallback mode
                                 if agent_name == "copilot":
-                                    prompt_file = project_root / ".github" / "prompts" / f"{cmd_name}.prompt.md"
-                                    if prompt_file.exists() and str(prompt_file) not in backed_up_command_files:
-                                        prompt_file.unlink()
+                                    prompts_dir = project_root / ".github" / "prompts"
+                                    if prompts_dir.is_dir():
+                                        for prompt_file in prompts_dir.glob("*.prompt.md"):
+                                            prompt_path_str = str(prompt_file)
+                                            if prompt_path_str in backed_up_command_files:
+                                                continue
+                                            if backup_timestamp is not None:
+                                                try:
+                                                    if prompt_file.stat().st_mtime < backup_timestamp:
+                                                        continue
+                                                except OSError:
+                                                    continue
+                                            prompt_file.unlink()
                     except KeyError:
-                        pass  # No new registry entry exists, nothing to clean up
+                        # If the registry lookup itself fails, we conservatively do
+                        # nothing extra here; the backup-based restoration below will
+                        # still put back any original files we saved.
+                        pass
 
                     # Restore backed up command files
                     for original_path, backup_path in backed_up_command_files.items():