shotgun-sh 0.1.12.dev3__tar.gz → 0.1.13.dev1__tar.gz

{shotgun_sh-0.1.12.dev3 → shotgun_sh-0.1.13.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: shotgun-sh
-Version: 0.1.12.dev3
+Version: 0.1.13.dev1
 Summary: AI-powered research, planning, and task management CLI tool
 Project-URL: Homepage, https://shotgun.sh/
 Project-URL: Repository, https://github.com/shotgun-sh/shotgun

{shotgun_sh-0.1.12.dev3 → shotgun_sh-0.1.13.dev1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "shotgun-sh"
-version = "0.1.12.dev3"
+version = "0.1.13.dev1"
 description = "AI-powered research, planning, and task management CLI tool"
 readme = "README.md"
 license = { text = "MIT" }

{shotgun_sh-0.1.12.dev3 → shotgun_sh-0.1.13.dev1}/src/shotgun/agents/agent_manager.py

@@ -31,6 +31,7 @@ from pydantic_ai.messages import (
     SystemPromptPart,
     ToolCallPart,
     ToolCallPartDelta,
+    ToolReturnPart,
 )
 from textual.message import Message
 from textual.widget import Widget
@@ -44,7 +45,7 @@ from shotgun.utils.source_detection import detect_source
 from .export import create_export_agent
 from .history.compaction import apply_persistent_compaction
 from .messages import AgentSystemPrompt
-from .models import AgentDeps, AgentRuntimeOptions
+from .models import AgentDeps, AgentRuntimeOptions, UserAnswer
 from .plan import create_plan_agent
 from .research import create_research_agent
 from .specify import create_specify_agent
@@ -272,6 +273,8 @@ class AgentManager(Widget):
         # Use merged deps (shared state + agent-specific system prompt) if not provided
         if deps is None:
             deps = self._create_merged_deps(self._current_agent_type)
+        if not deferred_tool_results:
+            self.ensure_agent_canecelled_safely()
 
         # Ensure deps is not None
         if deps is None:
@@ -674,6 +677,32 @@ class AgentManager(Widget):
         self.ui_message_history.append(message)
         self._post_messages_updated()
 
+    def ensure_agent_canecelled_safely(self) -> None:
+        if not self.message_history:
+            return
+        self.last_response = self.message_history[-1]
+        ## we're searching for unanswered ask_user parts
+        found_tool = None
+        for part in self.message_history[-1].parts:
+            if isinstance(part, ToolCallPart) and part.tool_name == "ask_user":
+                found_tool = part
+                break
+        if not found_tool:
+            return
+        tool_result = ModelRequest(
+            parts=[
+                ToolReturnPart(
+                    tool_call_id=found_tool.tool_call_id,
+                    tool_name=found_tool.tool_name,
+                    content=UserAnswer(
+                        answer="⚠️ Operation cancelled by user",
+                        tool_call_id=found_tool.tool_call_id,
+                    ),
+                )
+            ]
+        )
+        self.message_history.append(tool_result)
+
 
 # Re-export AgentType for backward compatibility
 __all__ = [

{shotgun_sh-0.1.12.dev3 → shotgun_sh-0.1.13.dev1}/src/shotgun/cli/codebase/commands.py

@@ -6,8 +6,21 @@ from pathlib import Path
 from typing import Annotated
 
 import typer
+from rich.console import Console
+from rich.progress import (
+    BarColumn,
+    Progress,
+    SpinnerColumn,
+    TaskProgressColumn,
+    TextColumn,
+    TimeElapsedColumn,
+)
 
-from shotgun.codebase.models import CodebaseGraph, QueryType
+from shotgun.codebase.models import (
+    CodebaseGraph,
+    IndexProgress,
+    QueryType,
+)
 from shotgun.logging_config import get_logger
 from shotgun.sdk.codebase import CodebaseSDK
 from shotgun.sdk.exceptions import CodebaseNotFoundError, InvalidPathError
@@ -59,10 +72,66 @@ def index(
 ) -> None:
     """Index a new codebase."""
     sdk = CodebaseSDK()
+    console = Console()
+
+    # Create progress display
+    progress = Progress(
+        SpinnerColumn(),
+        TextColumn("[bold blue]{task.description}"),
+        BarColumn(),
+        TaskProgressColumn(),
+        TimeElapsedColumn(),
+        console=console,
+    )
+
+    # Track tasks by phase
+    tasks = {}
+
+    def progress_callback(progress_info: IndexProgress) -> None:
+        """Update progress display based on indexing phase."""
+        phase = progress_info.phase
+
+        # Create task if it doesn't exist
+        if phase not in tasks:
+            if progress_info.total is not None:
+                tasks[phase] = progress.add_task(
+                    progress_info.phase_name, total=progress_info.total
+                )
+            else:
+                # Indeterminate progress (spinner only)
+                tasks[phase] = progress.add_task(progress_info.phase_name, total=None)
+
+        task_id = tasks[phase]
+
+        # Update task
+        if progress_info.total is not None:
+            progress.update(
+                task_id,
+                completed=progress_info.current,
+                total=progress_info.total,
+                description=f"[bold blue]{progress_info.phase_name}",
+            )
+        else:
+            # Just update description for indeterminate tasks
+            progress.update(
+                task_id,
+                description=f"[bold blue]{progress_info.phase_name} ({progress_info.current} items)",
+            )
+
+        # Mark as complete if phase is done
+        if progress_info.phase_complete:
+            if progress_info.total is not None:
+                progress.update(task_id, completed=progress_info.total)
 
     try:
         repo_path = Path(path)
-        result = asyncio.run(sdk.index_codebase(repo_path, name))
+
+        # Run indexing with progress display
+        with progress:
+            result = asyncio.run(
+                sdk.index_codebase(repo_path, name, progress_callback=progress_callback)
+            )
+
         output_result(result, format_type)
     except InvalidPathError as e:
         error_result = ErrorResult(error_message=str(e))

{shotgun_sh-0.1.12.dev3 → shotgun_sh-0.1.13.dev1}/src/shotgun/codebase/core/ingestor.py

@@ -535,6 +535,7 @@ class SimpleGraphBuilder:
         parsers: dict[str, Parser],
         queries: dict[str, Any],
         exclude_patterns: list[str] | None = None,
+        progress_callback: Any | None = None,
     ):
         self.ingestor = ingestor
         self.repo_path = repo_path
@@ -544,6 +545,7 @@ class SimpleGraphBuilder:
         self.ignore_dirs = IGNORE_PATTERNS
         if exclude_patterns:
             self.ignore_dirs = self.ignore_dirs.union(set(exclude_patterns))
+        self.progress_callback = progress_callback
 
         # Caches
         self.structural_elements: dict[Path, str | None] = {}
@@ -552,6 +554,34 @@ class SimpleGraphBuilder:
         self.simple_name_lookup: dict[str, set[str]] = defaultdict(set)
         self.class_inheritance: dict[str, list[str]] = {}  # class_qn -> [parent_qns]
 
+    def _report_progress(
+        self,
+        phase: str,
+        phase_name: str,
+        current: int,
+        total: int | None = None,
+        phase_complete: bool = False,
+    ) -> None:
+        """Report progress via callback if available."""
+        if not self.progress_callback:
+            return
+
+        try:
+            # Import here to avoid circular dependency
+            from shotgun.codebase.models import IndexProgress, ProgressPhase
+
+            progress = IndexProgress(
+                phase=ProgressPhase(phase),
+                phase_name=phase_name,
+                current=current,
+                total=total,
+                phase_complete=phase_complete,
+            )
+            self.progress_callback(progress)
+        except Exception as e:
+            # Don't let progress callback errors crash the build
+            logger.debug(f"Progress callback error: {e}")
+
     def run(self) -> None:
         """Run the three-pass graph building process."""
         logger.info(f"Building graph for project: {self.project_name}")
@@ -575,6 +605,7 @@ class SimpleGraphBuilder:
 
     def _identify_structure(self) -> None:
         """First pass: Walk directory to find packages and folders."""
+        dir_count = 0
         for root_str, dirs, _ in os.walk(self.repo_path, topdown=True):
             dirs[:] = [d for d in dirs if d not in self.ignore_dirs]
             root = Path(root_str)
@@ -584,6 +615,13 @@ class SimpleGraphBuilder:
             if root == self.repo_path:
                 continue
 
+            dir_count += 1
+            # Report progress every 10 directories
+            if dir_count % 10 == 0:
+                self._report_progress(
+                    "structure", "Identifying packages and folders", dir_count
+                )
+
             parent_rel_path = relative_root.parent
             parent_container_qn = self.structural_elements.get(parent_rel_path)
 
@@ -686,8 +724,34 @@ class SimpleGraphBuilder:
 
                 self.structural_elements[relative_root] = None
 
+        # Report phase completion
+        self._report_progress(
+            "structure",
+            "Identifying packages and folders",
+            dir_count,
+            phase_complete=True,
+        )
+
     def _process_files(self) -> None:
         """Second pass: Process files and extract definitions."""
+        # First pass: Count total files
+        total_files = 0
+        for root_str, _, files in os.walk(self.repo_path):
+            root = Path(root_str)
+
+            # Skip ignored directories
+            if any(part in self.ignore_dirs for part in root.parts):
+                continue
+
+            for filename in files:
+                filepath = root / filename
+                ext = filepath.suffix
+                lang_config = get_language_config(ext)
+
+                if lang_config and lang_config.name in self.parsers:
+                    total_files += 1
+
+        # Second pass: Process files with progress reporting
         file_count = 0
         for root_str, _, files in os.walk(self.repo_path):
             root = Path(root_str)
@@ -707,10 +771,27 @@ class SimpleGraphBuilder:
                     self._process_single_file(filepath, lang_config.name)
                     file_count += 1
 
+                    # Report progress after each file
+                    self._report_progress(
+                        "definitions",
+                        "Processing files and extracting definitions",
+                        file_count,
+                        total_files,
+                    )
+
                     if file_count % 100 == 0:
-                        logger.info(f"  Processed {file_count} files...")
+                        logger.info(f"  Processed {file_count}/{total_files} files...")
+
+        logger.info(f"  Total files processed: {file_count}/{total_files}")
 
-        logger.info(f"  Total files processed: {file_count}")
+        # Report phase completion
+        self._report_progress(
+            "definitions",
+            "Processing files and extracting definitions",
+            file_count,
+            total_files,
+            phase_complete=True,
+        )
 
     def _process_single_file(self, filepath: Path, language: str) -> None:
         """Process a single file."""
@@ -1143,7 +1224,8 @@ class SimpleGraphBuilder:
         self._process_inheritance()
 
         # Then process function calls
-        logger.info(f"Processing function calls for {len(self.ast_cache)} files...")
+        total_files = len(self.ast_cache)
+        logger.info(f"Processing function calls for {total_files} files...")
         logger.info(f"Function registry has {len(self.function_registry)} entries")
         logger.info(
             f"Simple name lookup has {len(self.simple_name_lookup)} unique names"
@@ -1157,10 +1239,29 @@ class SimpleGraphBuilder:
                     f"  Example: '{name}' -> {list(self.simple_name_lookup[name])[:3]}"
                 )
 
+        file_count = 0
         for filepath, (root_node, language) in self.ast_cache.items():
             self._process_calls(filepath, root_node, language)
             # NOTE: Add import processing. wtf does this mean?
 
+            file_count += 1
+            # Report progress after each file
+            self._report_progress(
+                "relationships",
+                "Processing relationships (calls, imports)",
+                file_count,
+                total_files,
+            )
+
+        # Report phase completion
+        self._report_progress(
+            "relationships",
+            "Processing relationships (calls, imports)",
+            file_count,
+            total_files,
+            phase_complete=True,
+        )
+
     def _process_inheritance(self) -> None:
         """Process inheritance relationships between classes."""
         logger.info("Processing inheritance relationships...")
@@ -1444,6 +1545,7 @@ class CodebaseIngestor:
         db_path: str,
         project_name: str | None = None,
         exclude_patterns: list[str] | None = None,
+        progress_callback: Any | None = None,
     ):
         """Initialize the ingestor.
 
@@ -1451,10 +1553,12 @@ class CodebaseIngestor:
             db_path: Path to Kuzu database
             project_name: Optional project name
             exclude_patterns: Patterns to exclude from processing
+            progress_callback: Optional callback for progress reporting
         """
         self.db_path = Path(db_path)
         self.project_name = project_name
         self.exclude_patterns = exclude_patterns or []
+        self.progress_callback = progress_callback
 
     def build_graph_from_directory(self, repo_path: str) -> None:
         """Build a code knowledge graph from a directory.
@@ -1484,7 +1588,12 @@ class CodebaseIngestor:
 
             # Build graph
             builder = SimpleGraphBuilder(
-                ingestor, repo_path_obj, parsers, queries, self.exclude_patterns
+                ingestor,
+                repo_path_obj,
+                parsers,
+                queries,
+                self.exclude_patterns,
+                self.progress_callback,
             )
             if self.project_name:
                 builder.project_name = self.project_name

{shotgun_sh-0.1.12.dev3 → shotgun_sh-0.1.13.dev1}/src/shotgun/codebase/core/manager.py

@@ -329,6 +329,7 @@ class CodebaseGraphManager:
         languages: list[str] | None = None,
         exclude_patterns: list[str] | None = None,
         indexed_from_cwd: str | None = None,
+        progress_callback: Any | None = None,
     ) -> CodebaseGraph:
         """Build a new code knowledge graph.
 
@@ -337,6 +338,7 @@ class CodebaseGraphManager:
             name: Optional human-readable name
             languages: Languages to parse (default: all supported)
             exclude_patterns: Patterns to exclude
+            progress_callback: Optional callback for progress reporting
 
         Returns:
             Created graph metadata
@@ -353,7 +355,19 @@ class CodebaseGraphManager:
 
         # Check if graph already exists
         if graph_path.exists():
-            raise CodebaseAlreadyIndexedError(repo_path)
+            # Verify it's not corrupted by checking if we can load the Project node
+            existing_graph = await self.get_graph(graph_id)
+            if existing_graph:
+                # Valid existing graph
+                raise CodebaseAlreadyIndexedError(repo_path)
+            else:
+                # Corrupted database - remove and re-index
+                logger.warning(
+                    f"Found corrupted database at {graph_path}, removing for re-indexing..."
+                )
+                import shutil
+
+                shutil.rmtree(graph_path)
 
         # Import the builder from local core module
         from shotgun.codebase.core import CodebaseIngestor
@@ -379,6 +393,7 @@ class CodebaseGraphManager:
             db_path=str(graph_path),
             project_name=name,
             exclude_patterns=exclude_patterns or [],
+            progress_callback=progress_callback,
         )
 
         # Run build in thread pool
@@ -1193,6 +1208,136 @@ class CodebaseGraphManager:
             )
             return None
 
+    async def cleanup_corrupted_databases(self) -> list[str]:
+        """Detect and remove corrupted Kuzu databases.
+
+        This method iterates through all .kuzu files in the storage directory,
+        attempts to open them, and removes any that are corrupted or unreadable.
+
+        Returns:
+            List of graph_ids that were removed due to corruption
+        """
+        import shutil
+
+        removed_graphs = []
+
+        # Find all .kuzu files (can be files or directories)
+        for path in self.storage_dir.glob("*.kuzu"):
+            graph_id = path.stem
+
+            # If it's a plain file (not a directory), it's corrupted
+            # Valid Kuzu databases are always directories
+            if path.is_file():
+                logger.warning(
+                    f"Detected corrupted database file (should be directory) at {path}, removing it"
+                )
+                try:
+                    await anyio.to_thread.run_sync(path.unlink)
+                    removed_graphs.append(graph_id)
+                    logger.info(f"Removed corrupted database file: {graph_id}")
+                except Exception as e:
+                    logger.error(
+                        f"Failed to remove corrupted database file {graph_id}: {e}"
+                    )
+                continue
+
+            try:
+                # Try to open the database with a timeout to prevent hanging
+                async def try_open_database(
+                    gid: str = graph_id, db_path: Path = path
+                ) -> bool:
+                    lock = await self._get_lock()
+                    async with lock:
+                        # Close existing connections if any
+                        if gid in self._connections:
+                            try:
+                                self._connections[gid].close()
+                            except Exception as e:
+                                logger.debug(
+                                    f"Failed to close connection for {gid}: {e}"
+                                )
+                            del self._connections[gid]
+                        if gid in self._databases:
+                            try:
+                                self._databases[gid].close()
+                            except Exception as e:
+                                logger.debug(f"Failed to close database for {gid}: {e}")
+                            del self._databases[gid]
+
+                        # Try to open the database
+                        def _open_and_query(g: str = gid, p: Path = db_path) -> bool:
+                            db = kuzu.Database(str(p))
+                            conn = kuzu.Connection(db)
+                            try:
+                                result = conn.execute(
+                                    "MATCH (p:Project {graph_id: $graph_id}) RETURN p",
+                                    {"graph_id": g},
+                                )
+                                has_results = (
+                                    result.has_next()
+                                    if hasattr(result, "has_next")
+                                    else False
+                                )
+                                return has_results
+                            finally:
+                                conn.close()
+                                db.close()
+
+                        return await anyio.to_thread.run_sync(_open_and_query)
+
+                # Try to open with 5 second timeout
+                has_project = await asyncio.wait_for(try_open_database(), timeout=5.0)
+
+                if not has_project:
+                    # Database exists but has no Project node - consider it corrupted
+                    raise ValueError("No Project node found in database")
+
+            except (Exception, asyncio.TimeoutError) as e:
+                # Database is corrupted or timed out - remove it
+                error_type = (
+                    "timed out" if isinstance(e, asyncio.TimeoutError) else "corrupted"
+                )
+                logger.warning(
+                    f"Detected {error_type} database at {path}, removing it. "
+                    f"Error: {str(e) if not isinstance(e, asyncio.TimeoutError) else 'Operation timed out after 5 seconds'}"
+                )
+
+                try:
+                    # Clean up any open connections
+                    lock = await self._get_lock()
+                    async with lock:
+                        if graph_id in self._connections:
+                            try:
+                                self._connections[graph_id].close()
+                            except Exception as e:
+                                logger.debug(
+                                    f"Failed to close connection during cleanup for {graph_id}: {e}"
+                                )
+                            del self._connections[graph_id]
+                        if graph_id in self._databases:
+                            try:
+                                self._databases[graph_id].close()
+                            except Exception as e:
+                                logger.debug(
+                                    f"Failed to close database during cleanup for {graph_id}: {e}"
+                                )
+                            del self._databases[graph_id]
+
+                    # Remove the database (could be file or directory)
+                    if path.is_dir():
+                        await anyio.to_thread.run_sync(shutil.rmtree, path)
+                    else:
+                        await anyio.to_thread.run_sync(path.unlink)
+                    removed_graphs.append(graph_id)
+                    logger.info(f"Removed {error_type} database: {graph_id}")
+
+                except Exception as cleanup_error:
+                    logger.error(
+                        f"Failed to remove corrupted database {graph_id}: {cleanup_error}"
+                    )
+
+        return removed_graphs
+
     async def list_graphs(self) -> list[CodebaseGraph]:
         """List all available graphs.
 
@@ -1464,6 +1609,7 @@ class CodebaseGraphManager:
         languages: list[str] | None,
         exclude_patterns: list[str] | None,
         indexed_from_cwd: str | None = None,
+        progress_callback: Any | None = None,
     ) -> CodebaseGraph:
         """Internal implementation of graph building (runs in background)."""
         operation_id = str(uuid.uuid4())
@@ -1487,7 +1633,13 @@ class CodebaseGraphManager:
 
             # Do the actual build work
             graph = await self._do_build_graph(
-                graph_id, repo_path, name, languages, exclude_patterns, indexed_from_cwd
+                graph_id,
+                repo_path,
+                name,
+                languages,
+                exclude_patterns,
+                indexed_from_cwd,
+                progress_callback,
             )
 
             # Update operation stats
@@ -1536,6 +1688,7 @@ class CodebaseGraphManager:
         languages: list[str] | None,
         exclude_patterns: list[str] | None,
         indexed_from_cwd: str | None = None,
+        progress_callback: Any | None = None,
     ) -> CodebaseGraph:
         """Execute the actual graph building logic (extracted from original build_graph)."""
         # The database and Project node already exist from _initialize_graph_metadata
@@ -1591,6 +1744,7 @@ class CodebaseGraphManager:
                 parsers=parsers,
                 queries=queries,
                 exclude_patterns=exclude_patterns,
+                progress_callback=progress_callback,
             )
 
             # Build the graph
@@ -1616,6 +1770,7 @@ class CodebaseGraphManager:
         languages: list[str] | None = None,
         exclude_patterns: list[str] | None = None,
         indexed_from_cwd: str | None = None,
+        progress_callback: Any | None = None,
     ) -> str:
         """Start building a new code knowledge graph asynchronously.
 
@@ -1654,7 +1809,13 @@ class CodebaseGraphManager:
         # Start the build operation in background
         task = asyncio.create_task(
             self._build_graph_impl(
-                graph_id, repo_path, name, languages, exclude_patterns, indexed_from_cwd
+                graph_id,
+                repo_path,
+                name,
+                languages,
+                exclude_patterns,
+                indexed_from_cwd,
+                progress_callback,
             )
         )
         self._operations[graph_id] = task

{shotgun_sh-0.1.12.dev3 → shotgun_sh-0.1.13.dev1}/src/shotgun/codebase/models.py

@@ -1,5 +1,6 @@
 """Data models for codebase service."""
 
+from collections.abc import Callable
 from enum import Enum
 from typing import Any
 
@@ -22,6 +23,30 @@ class QueryType(str, Enum):
     CYPHER = "cypher"
 
 
+class ProgressPhase(str, Enum):
+    """Phase of codebase indexing progress."""
+
+    STRUCTURE = "structure"  # Identifying packages and folders
+    DEFINITIONS = "definitions"  # Processing files and extracting definitions
+    RELATIONSHIPS = "relationships"  # Processing relationships (calls, imports)
+
+
+class IndexProgress(BaseModel):
+    """Progress information for codebase indexing."""
+
+    phase: ProgressPhase = Field(..., description="Current indexing phase")
+    phase_name: str = Field(..., description="Human-readable phase name")
+    current: int = Field(..., description="Current item count")
+    total: int | None = Field(None, description="Total items (None if unknown)")
+    phase_complete: bool = Field(
+        default=False, description="Whether this phase is complete"
+    )
+
+
+# Type alias for progress callback function
+ProgressCallback = Callable[[IndexProgress], None]
+
+
 class OperationStats(BaseModel):
     """Statistics for a graph operation (build/update)."""