mita-code 0.1.0__py3-none-any.whl

mita/tools/safety.py

@@ -0,0 +1,91 @@
+"""Destructive action detection and banned command checking."""
+
+from __future__ import annotations
+
+import re
+
+from mita.config.schema import ToolSettings
+from mita.tools.schema import ToolCall, ToolDefinition
+
+# Patterns that indicate destructive shell commands
+DESTRUCTIVE_PATTERNS: list[re.Pattern[str]] = [
+    re.compile(r"\brm\s+(-[a-zA-Z]*f|-[a-zA-Z]*r)", re.IGNORECASE),
+    re.compile(r"\bgit\s+(push\s+--force|reset\s+--hard|clean\s+-[a-zA-Z]*f)", re.IGNORECASE),
+    re.compile(r"\bchmod\s+-R\s+0?7", re.IGNORECASE),
+    re.compile(r"\bchown\s+-R\b", re.IGNORECASE),
+    re.compile(r"\b(truncate|shred)\b", re.IGNORECASE),
+    re.compile(r">\s*/dev/\w+", re.IGNORECASE),
+]
+
+# Git subcommands that are destructive and need confirmation
+DESTRUCTIVE_GIT_PATTERNS: list[re.Pattern[str]] = [
+    re.compile(r"^push\s+--force", re.IGNORECASE),
+    re.compile(r"^push\s+-f\b", re.IGNORECASE),
+    re.compile(r"^reset\s+--hard", re.IGNORECASE),
+    re.compile(r"^clean\s+-[a-zA-Z]*f", re.IGNORECASE),
+    re.compile(r"^checkout\s+--\s", re.IGNORECASE),
+    re.compile(r"^branch\s+-[dD]\b", re.IGNORECASE),
+]
+
+
+def is_command_banned(command: str, banned_commands: list[str]) -> bool:
+    """Check if a shell command matches any banned command pattern."""
+    normalized = command.strip()
+    for banned in banned_commands:
+        if banned in normalized:
+            return True
+    return False
+
+
+def is_command_destructive(command: str) -> bool:
+    """Check if a shell command looks destructive."""
+    for pattern in DESTRUCTIVE_PATTERNS:
+        if pattern.search(command):
+            return True
+    return False
+
+
+def is_git_command_destructive(subcommand: str) -> bool:
+    """Check if a git subcommand is destructive."""
+    stripped = subcommand.strip()
+    for pattern in DESTRUCTIVE_GIT_PATTERNS:
+        if pattern.search(stripped):
+            return True
+    return False
+
+
+def needs_confirmation(
+    tool_call: ToolCall,
+    tool_def: ToolDefinition,
+    settings: ToolSettings,
+) -> bool:
+    """Determine if a tool call needs user confirmation before execution.
+
+    Returns True if:
+    - The tool is marked destructive AND confirm_destructive is on
+      AND the tool is not in auto_approve.
+    - Or the tool is 'shell' and the command looks destructive.
+    - Or the tool is 'git' and the subcommand is destructive.
+    """
+    if not settings.confirm_destructive:
+        return False
+
+    if tool_call.name in settings.auto_approve:
+        return False
+
+    if tool_def.destructive:
+        return True
+
+    # Extra check for shell commands that look destructive
+    if tool_call.name == "shell":
+        command = tool_call.arguments.get("command", "")
+        if is_command_destructive(command):
+            return True
+
+    # Check git subcommands for destructive operations
+    if tool_call.name == "git":
+        subcommand = tool_call.arguments.get("subcommand", "")
+        if is_git_command_destructive(subcommand):
+            return True
+
+    return False

mita/tools/schema.py

@@ -0,0 +1,87 @@
+"""Pydantic models for tool definitions, calls, and results."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class ToolParameter(BaseModel):
+    """A parameter in a tool definition."""
+
+    name: str
+    type: str  # "string", "integer", "boolean", "array", "object"
+    description: str
+    required: bool = True
+    default: Any | None = None
+
+
+class ToolDefinition(BaseModel):
+    """Schema sent to the LLM so it knows what tools are available."""
+
+    name: str
+    description: str
+    parameters: list[ToolParameter]
+    destructive: bool = False  # triggers confirmation flow
+    source: str = "builtin"  # "builtin" | "mcp:<plugin_name>"
+
+    def to_openai_schema(self) -> dict[str, Any]:
+        """Convert to OpenAI-compatible function schema for LiteLLM."""
+        properties: dict[str, Any] = {}
+        required: list[str] = []
+
+        for param in self.parameters:
+            prop: dict[str, Any] = {
+                "type": param.type,
+                "description": param.description,
+            }
+            if param.default is not None:
+                prop["default"] = param.default
+            properties[param.name] = prop
+            if param.required:
+                required.append(param.name)
+
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": {
+                    "type": "object",
+                    "properties": properties,
+                    "required": required,
+                },
+            },
+        }
+
+
+class ToolCall(BaseModel):
+    """Parsed from LLM output."""
+
+    id: str
+    name: str
+    arguments: dict[str, Any]
+
+
+class ToolResult(BaseModel):
+    """Returned after executing a tool."""
+
+    tool_call_id: str
+    success: bool
+    output: str = ""
+    error: str | None = None
+    truncated: bool = False
+
+    # Max output size before truncation (10K chars)
+    MAX_OUTPUT_SIZE: int = Field(default=10_000, exclude=True)
+
+    def truncate_output(self) -> ToolResult:
+        """Return a copy with output truncated if too long."""
+        if len(self.output) <= self.MAX_OUTPUT_SIZE:
+            return self
+        truncated_output = (
+            self.output[: self.MAX_OUTPUT_SIZE]
+            + f"\n\n... [truncated, {len(self.output):,} chars total]"
+        )
+        return self.model_copy(update={"output": truncated_output, "truncated": True})

mita/ui/__init__.py

@@ -0,0 +1 @@
+"""Terminal UI: REPL, Rich rendering, spinners, themes."""

mita/ui/display.py

@@ -0,0 +1,139 @@
+"""Rich terminal rendering: markdown, code blocks, diffs, tool calls."""
+
+from __future__ import annotations
+
+from rich.console import Console
+from rich.markdown import Markdown
+from rich.panel import Panel
+
+from mita.tools.schema import ToolCall, ToolResult
+from mita.ui.theme import MITA_THEME
+
+# Module-level console with mita theme
+_console: Console | None = None
+
+
+def get_console() -> Console:
+    """Get or create the themed console."""
+    global _console  # noqa: PLW0603
+    if _console is None:
+        _console = Console(theme=MITA_THEME)
+    return _console
+
+
+def display_welcome(console: Console) -> None:
+    """Display the welcome banner."""
+    console.print(
+        Panel(
+            "[mita.prompt]mita[/mita.prompt] — local-first coding assistant\n"
+            "[mita.dim]Type your prompt, or /quit to exit.[/mita.dim]",
+            border_style="mita.info",
+        )
+    )
+
+
+def display_goodbye(console: Console) -> None:
+    """Display the exit message."""
+    console.print("[mita.dim]Goodbye![/mita.dim]")
+
+
+def display_markdown(console: Console, text: str) -> None:
+    """Render markdown text in the terminal."""
+    if not text.strip():
+        return
+    md = Markdown(text)
+    console.print(md)
+
+
+def display_streaming_token(console: Console, token: str) -> None:
+    """Display a single streaming token (no newline)."""
+    console.print(token, end="", highlight=False)
+
+
+def display_streaming_end(console: Console) -> None:
+    """End a streaming output block."""
+    console.print()  # final newline
+
+
+def display_tool_call(console: Console, tool_call: ToolCall) -> None:
+    """Display a tool call that is about to be executed."""
+    args_str = ", ".join(f"{k}={v!r}" for k, v in tool_call.arguments.items())
+    # Truncate long args for display
+    if len(args_str) > 200:
+        args_str = args_str[:200] + "..."
+    console.print(
+        f"  [mita.tool_name]{tool_call.name}[/mita.tool_name]({args_str})",
+    )
+
+
+def display_tool_result(console: Console, result: ToolResult) -> None:
+    """Display the result of a tool execution."""
+    if result.success:
+        if result.output:
+            # Show truncated output
+            output = result.output
+            if len(output) > 500:
+                output = output[:500] + "\n..."
+            console.print(f"  [mita.dim]{output}[/mita.dim]")
+    else:
+        console.print(f"  [mita.tool_error]Error: {result.error}[/mita.tool_error]")
+
+
+def display_error(console: Console, message: str) -> None:
+    """Display an error message."""
+    console.print(f"[mita.error]{message}[/mita.error]")
+
+
+def display_warning(console: Console, message: str) -> None:
+    """Display a warning message."""
+    console.print(f"[mita.warning]{message}[/mita.warning]")
+
+
+def display_token_usage(console: Console, total_tokens: int) -> None:
+    """Display token usage after a response (estimated)."""
+    console.print(f"[mita.token_count]({total_tokens:,} tokens)[/mita.token_count]")
+
+
+def display_response_stats(
+    console: Console,
+    prompt_tokens: int,
+    completion_tokens: int,
+    total_time: float,
+    ttft: float | None = None,
+) -> None:
+    """Display detailed response statistics after an LLM call."""
+    parts: list[str] = []
+
+    parts.append(f"prompt: {prompt_tokens:,}")
+    parts.append(f"completion: {completion_tokens:,}")
+
+    if ttft is not None:
+        parts.append(f"ttft: {ttft:.1f}s")
+
+    parts.append(f"total: {total_time:.1f}s")
+
+    if completion_tokens > 0 and total_time > 0:
+        tps = completion_tokens / total_time
+        parts.append(f"{tps:.1f} tok/s")
+
+    console.print(f"[mita.token_count]({' · '.join(parts)})[/mita.token_count]")
+
+
+def display_error_with_suggestion(console: Console, error: str, suggestion: str) -> None:
+    """Display an error with an actionable suggestion."""
+    console.print(f"[red]{error}[/red]")
+    console.print(f"[dim]  → {suggestion}[/dim]")
+
+
+async def prompt_user_confirm(console: Console, question: str) -> bool:
+    """Ask the user for confirmation before a destructive action.
+
+    Returns True if the user approves, False otherwise.
+    """
+    console.print(f"[mita.warning]{question}[/mita.warning] ", end="")
+    try:
+        response = console.input("[y/N] ")
+        return response.strip().lower() in ("y", "yes")
+    except (EOFError, KeyboardInterrupt):
+        console.print()
+        return False

mita/ui/repl.py

@@ -0,0 +1,88 @@
+"""Interactive REPL input loop with prompt_toolkit."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Coroutine
+from typing import Any
+
+from prompt_toolkit import PromptSession
+from prompt_toolkit.history import InMemoryHistory
+from rich.console import Console
+
+from mita.ui.display import display_goodbye, display_welcome
+
+# Built-in REPL commands that should NOT be treated as skill invocations.
+_BUILTIN_COMMANDS = frozenset({"/quit", "/exit", "/q", "/clear"})
+
+
+async def repl_loop(
+    console: Console,
+    on_input: Callable[[str], Coroutine[Any, Any, None]],
+    on_clear: Callable[[], None] | None = None,
+    skills_paths: list[str] | None = None,
+) -> None:
+    """Run the interactive REPL.
+
+    Args:
+        console: Rich console for output.
+        on_input: Async callback called with each user input line.
+        on_clear: Callback invoked when the user types /clear.
+        skills_paths: Skill search paths for ``/`` prefix invocation.
+    """
+    display_welcome(console)
+
+    session: PromptSession[str] = PromptSession(history=InMemoryHistory())
+
+    while True:
+        try:
+            user_input = await session.prompt_async("mita> ")
+        except (EOFError, KeyboardInterrupt):
+            display_goodbye(console)
+            break
+
+        user_input = user_input.strip()
+        if not user_input:
+            continue
+
+        if user_input.lower() in ("/quit", "/exit", "/q"):
+            display_goodbye(console)
+            break
+
+        if user_input.lower() == "/clear":
+            if on_clear is not None:
+                on_clear()
+            console.print("[mita.dim]Conversation cleared.[/mita.dim]")
+            continue
+
+        # Skill invocation: /name [args]
+        if user_input.startswith("/") and user_input.split()[0].lower() not in _BUILTIN_COMMANDS:
+            rendered = _try_render_skill(user_input, skills_paths or [], console)
+            if rendered is not None:
+                try:
+                    await on_input(rendered)
+                except KeyboardInterrupt:
+                    console.print("\n[Interrupted]")
+                continue
+
+        try:
+            await on_input(user_input)
+        except KeyboardInterrupt:
+            console.print("\n[Interrupted]")
+
+
+def _try_render_skill(user_input: str, skills_paths: list[str], console: Console) -> str | None:
+    """Attempt to find and render a skill from user input.
+
+    Returns the rendered prompt string, or ``None`` if no matching skill was found.
+    """
+    from mita.skills.executor import render_skill
+    from mita.skills.loader import discover_skills, find_skill
+
+    skills = discover_skills(skills_paths)
+    skill = find_skill(skills, user_input)
+    if skill is None:
+        return None
+
+    rendered = render_skill(skill, user_input)
+    console.print(f"[dim]Running skill: {skill.frontmatter.name}[/dim]")
+    return rendered

mita/ui/spinner.py

@@ -0,0 +1,48 @@
+"""Thinking/loading indicators for the terminal."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Generator
+from contextlib import contextmanager
+
+from rich.console import Console, RenderableType
+from rich.live import Live
+from rich.spinner import Spinner as RichSpinner
+from rich.table import Table
+from rich.text import Text
+
+
+class _TimedSpinner:
+    """A spinner that shows elapsed time alongside the message."""
+
+    def __init__(self, message: str = "Thinking...") -> None:
+        self._message = message
+        self._start = time.monotonic()
+        self._spinner = RichSpinner("dots")
+
+    def __rich__(self) -> RenderableType:
+        elapsed = time.monotonic() - self._start
+        table = Table.grid(padding=(0, 1))
+        table.add_row(
+            self._spinner,
+            Text(self._message, style="mita.dim"),
+            Text(f"({elapsed:.0f}s)", style="dim"),
+        )
+        return table
+
+
+@contextmanager
+def thinking_spinner(
+    console: Console,
+    message: str = "Thinking...",
+) -> Generator[Live, None, None]:
+    """Display a spinner with elapsed time while the agent is thinking.
+
+    Usage:
+        with thinking_spinner(console):
+            await long_operation()
+    """
+    spinner = _TimedSpinner(message)
+    with Live(spinner, console=console, transient=True, refresh_per_second=4) as live:
+        yield live

mita/ui/theme.py

@@ -0,0 +1,23 @@
+"""Color scheme and style constants for terminal UI."""
+
+from __future__ import annotations
+
+from rich.style import Style
+from rich.theme import Theme
+
+# Style constants
+STYLES = {
+    "mita.prompt": Style(color="cyan", bold=True),
+    "mita.assistant": Style(color="white"),
+    "mita.tool_name": Style(color="yellow", bold=True),
+    "mita.tool_result": Style(color="green"),
+    "mita.tool_error": Style(color="red"),
+    "mita.info": Style(color="blue"),
+    "mita.warning": Style(color="yellow"),
+    "mita.error": Style(color="red", bold=True),
+    "mita.dim": Style(dim=True),
+    "mita.success": Style(color="green", bold=True),
+    "mita.token_count": Style(color="cyan", dim=True),
+}
+
+MITA_THEME = Theme(STYLES)