codemaster-cli 1.0.1__py3-none-any.whl

vibe/core/tools/builtins/todo.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+from enum import StrEnum, auto
+from typing import ClassVar
+
+from pydantic import BaseModel, Field
+
+from vibe.core.tools.base import (
+    BaseTool,
+    BaseToolConfig,
+    BaseToolState,
+    InvokeContext,
+    ToolError,
+    ToolPermission,
+)
+from vibe.core.tools.ui import ToolCallDisplay, ToolResultDisplay, ToolUIData
+from vibe.core.types import ToolCallEvent, ToolResultEvent, ToolStreamEvent
+
+
+class TodoStatus(StrEnum):
+    PENDING = auto()
+    IN_PROGRESS = auto()
+    COMPLETED = auto()
+    CANCELLED = auto()
+
+
+class TodoPriority(StrEnum):
+    LOW = auto()
+    MEDIUM = auto()
+    HIGH = auto()
+
+
+class TodoItem(BaseModel):
+    id: str
+    content: str
+    status: TodoStatus = TodoStatus.PENDING
+    priority: TodoPriority = TodoPriority.MEDIUM
+
+
+class TodoArgs(BaseModel):
+    action: str = Field(description="Either 'read' or 'write'")
+    todos: list[TodoItem] | None = Field(
+        default=None, description="Complete list of todos when writing."
+    )
+
+
+class TodoResult(BaseModel):
+    message: str
+    todos: list[TodoItem]
+    total_count: int
+
+
+class TodoConfig(BaseToolConfig):
+    permission: ToolPermission = ToolPermission.ALWAYS
+    max_todos: int = 100
+
+
+class TodoState(BaseToolState):
+    todos: list[TodoItem] = Field(default_factory=list)
+
+
+class Todo(
+    BaseTool[TodoArgs, TodoResult, TodoConfig, TodoState],
+    ToolUIData[TodoArgs, TodoResult],
+):
+    description: ClassVar[str] = (
+        "Manage todos. Use action='read' to view, action='write' with complete list to update."
+    )
+
+    @classmethod
+    def get_call_display(cls, event: ToolCallEvent) -> ToolCallDisplay:
+        if not isinstance(event.args, TodoArgs):
+            return ToolCallDisplay(summary="Invalid arguments")
+
+        args = event.args
+
+        match args.action:
+            case "read":
+                return ToolCallDisplay(summary="Reading todos")
+            case "write":
+                count = len(args.todos) if args.todos else 0
+                return ToolCallDisplay(summary=f"Writing {count} todos")
+            case _:
+                return ToolCallDisplay(summary=f"Unknown action: {args.action}")
+
+    @classmethod
+    def get_result_display(cls, event: ToolResultEvent) -> ToolResultDisplay:
+        if not isinstance(event.result, TodoResult):
+            return ToolResultDisplay(success=True, message="Success")
+
+        result = event.result
+
+        return ToolResultDisplay(success=True, message=result.message)
+
+    @classmethod
+    def get_status_text(cls) -> str:
+        return "Managing todos"
+
+    async def run(
+        self, args: TodoArgs, ctx: InvokeContext | None = None
+    ) -> AsyncGenerator[ToolStreamEvent | TodoResult, None]:
+        match args.action:
+            case "read":
+                yield self._read_todos()
+            case "write":
+                yield self._write_todos(args.todos or [])
+            case _:
+                raise ToolError(
+                    f"Invalid action '{args.action}'. Use 'read' or 'write'."
+                )
+
+    def _read_todos(self) -> TodoResult:
+        return TodoResult(
+            message=f"Retrieved {len(self.state.todos)} todos",
+            todos=self.state.todos,
+            total_count=len(self.state.todos),
+        )
+
+    def _write_todos(self, todos: list[TodoItem]) -> TodoResult:
+        if len(todos) > self.config.max_todos:
+            raise ToolError(f"Cannot store more than {self.config.max_todos} todos")
+
+        ids = [todo.id for todo in todos]
+        if len(ids) != len(set(ids)):
+            raise ToolError("Todo IDs must be unique")
+
+        self.state.todos = todos
+
+        return TodoResult(
+            message=f"Updated {len(todos)} todos",
+            todos=self.state.todos,
+            total_count=len(self.state.todos),
+        )

vibe/core/tools/builtins/write_file.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+from pathlib import Path
+from typing import ClassVar, final
+
+import anyio
+from pydantic import BaseModel, Field
+
+from vibe.core.tools.base import (
+    BaseTool,
+    BaseToolConfig,
+    BaseToolState,
+    InvokeContext,
+    ToolError,
+    ToolPermission,
+)
+from vibe.core.tools.ui import ToolCallDisplay, ToolResultDisplay, ToolUIData
+from vibe.core.types import ToolCallEvent, ToolResultEvent, ToolStreamEvent
+
+
+class WriteFileArgs(BaseModel):
+    path: str
+    content: str
+    overwrite: bool = Field(
+        default=False, description="Must be set to true to overwrite an existing file."
+    )
+
+
+class WriteFileResult(BaseModel):
+    path: str
+    bytes_written: int
+    file_existed: bool
+    content: str
+
+
+class WriteFileConfig(BaseToolConfig):
+    permission: ToolPermission = ToolPermission.ASK
+    max_write_bytes: int = 64_000
+    create_parent_dirs: bool = True
+
+
+class WriteFileState(BaseToolState):
+    recently_written_files: list[str] = Field(default_factory=list)
+
+
+class WriteFile(
+    BaseTool[WriteFileArgs, WriteFileResult, WriteFileConfig, WriteFileState],
+    ToolUIData[WriteFileArgs, WriteFileResult],
+):
+    description: ClassVar[str] = (
+        "Create or overwrite a UTF-8 file. Fails if file exists unless 'overwrite=True'."
+    )
+
+    @classmethod
+    def get_call_display(cls, event: ToolCallEvent) -> ToolCallDisplay:
+        if not isinstance(event.args, WriteFileArgs):
+            return ToolCallDisplay(summary="Invalid arguments")
+
+        args = event.args
+
+        return ToolCallDisplay(
+            summary=f"Writing {args.path}{' (overwrite)' if args.overwrite else ''}",
+            content=args.content,
+        )
+
+    @classmethod
+    def get_result_display(cls, event: ToolResultEvent) -> ToolResultDisplay:
+        if isinstance(event.result, WriteFileResult):
+            action = "Overwritten" if event.result.file_existed else "Created"
+            return ToolResultDisplay(
+                success=True, message=f"{action} {Path(event.result.path).name}"
+            )
+
+        return ToolResultDisplay(success=True, message="File written")
+
+    @classmethod
+    def get_status_text(cls) -> str:
+        return "Writing file"
+
+    def check_allowlist_denylist(self, args: WriteFileArgs) -> ToolPermission | None:
+        import fnmatch
+
+        file_path = Path(args.path).expanduser()
+        if not file_path.is_absolute():
+            file_path = Path.cwd() / file_path
+        file_str = str(file_path)
+
+        for pattern in self.config.denylist:
+            if fnmatch.fnmatch(file_str, pattern):
+                return ToolPermission.NEVER
+
+        for pattern in self.config.allowlist:
+            if fnmatch.fnmatch(file_str, pattern):
+                return ToolPermission.ALWAYS
+
+        return None
+
+    @final
+    async def run(
+        self, args: WriteFileArgs, ctx: InvokeContext | None = None
+    ) -> AsyncGenerator[ToolStreamEvent | WriteFileResult, None]:
+        file_path, file_existed, content_bytes = self._prepare_and_validate_path(args)
+
+        await self._write_file(args, file_path)
+
+        BUFFER_SIZE = 10
+        self.state.recently_written_files.append(str(file_path))
+        if len(self.state.recently_written_files) > BUFFER_SIZE:
+            self.state.recently_written_files.pop(0)
+
+        yield WriteFileResult(
+            path=str(file_path),
+            bytes_written=content_bytes,
+            file_existed=file_existed,
+            content=args.content,
+        )
+
+    def _prepare_and_validate_path(self, args: WriteFileArgs) -> tuple[Path, bool, int]:
+        if not args.path.strip():
+            raise ToolError("Path cannot be empty")
+
+        content_bytes = len(args.content.encode("utf-8"))
+        if content_bytes > self.config.max_write_bytes:
+            raise ToolError(
+                f"Content exceeds {self.config.max_write_bytes} bytes limit"
+            )
+
+        file_path = Path(args.path).expanduser()
+        if not file_path.is_absolute():
+            file_path = Path.cwd() / file_path
+        file_path = file_path.resolve()
+
+        try:
+            file_path.relative_to(Path.cwd().resolve())
+        except ValueError:
+            raise ToolError(f"Cannot write outside project directory: {file_path}")
+
+        file_existed = file_path.exists()
+
+        if file_existed and not args.overwrite:
+            raise ToolError(
+                f"File '{file_path}' exists. Set overwrite=True to replace."
+            )
+
+        if self.config.create_parent_dirs:
+            file_path.parent.mkdir(parents=True, exist_ok=True)
+        elif not file_path.parent.exists():
+            raise ToolError(f"Parent directory does not exist: {file_path.parent}")
+
+        return file_path, file_existed, content_bytes
+
+    async def _write_file(self, args: WriteFileArgs, file_path: Path) -> None:
+        try:
+            async with await anyio.Path(file_path).open(
+                mode="w", encoding="utf-8"
+            ) as f:
+                await f.write(args.content)
+        except Exception as e:
+            raise ToolError(f"Error writing {file_path}: {e}") from e

vibe/core/tools/manager.py

@@ -0,0 +1,341 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator
+import hashlib
+import importlib.util
+import inspect
+from logging import getLogger
+from pathlib import Path
+import re
+import sys
+from typing import TYPE_CHECKING, Any
+
+from vibe.core.paths.config_paths import resolve_local_tools_dir
+from vibe.core.paths.global_paths import DEFAULT_TOOL_DIR, GLOBAL_TOOLS_DIR
+from vibe.core.tools.base import BaseTool, BaseToolConfig
+from vibe.core.tools.mcp import (
+    RemoteTool,
+    create_mcp_http_proxy_tool_class,
+    create_mcp_stdio_proxy_tool_class,
+    list_tools_http,
+    list_tools_stdio,
+)
+from vibe.core.utils import name_matches, run_sync
+
+logger = getLogger("vibe")
+
+if TYPE_CHECKING:
+    from vibe.core.config import MCPHttp, MCPStdio, MCPStreamableHttp, VibeConfig
+
+
+def _try_canonical_module_name(path: Path) -> str | None:
+    """Extract canonical module name for vibe package files.
+
+    Prevents Pydantic class identity mismatches when the same module
+    is imported via dynamic discovery and regular imports.
+    """
+    try:
+        parts = path.resolve().parts
+    except (OSError, ValueError):
+        return None
+
+    try:
+        vibe_idx = parts.index("vibe")
+    except ValueError:
+        return None
+
+    if vibe_idx + 1 >= len(parts):
+        return None
+
+    module_parts = [p.removesuffix(".py") for p in parts[vibe_idx:]]
+    return ".".join(module_parts)
+
+
+def _compute_module_name(path: Path) -> str:
+    """Return canonical module name for vibe files, hash-based synthetic name otherwise."""
+    if canonical := _try_canonical_module_name(path):
+        return canonical
+
+    resolved = path.resolve()
+    path_hash = hashlib.md5(str(resolved).encode()).hexdigest()[:8]
+    stem = re.sub(r"[^0-9A-Za-z_]", "_", path.stem) or "mod"
+    return f"vibe_tools_discovered_{stem}_{path_hash}"
+
+
+class NoSuchToolError(Exception):
+    """Exception raised when a tool is not found."""
+
+
+class ToolManager:
+    """Manages tool discovery and instantiation for an Agent.
+
+    Discovers available tools from the provided search paths. Each Agent
+    should have its own ToolManager instance.
+    """
+
+    def __init__(self, config_getter: Callable[[], VibeConfig]) -> None:
+        self._config_getter = config_getter
+        self._instances: dict[str, BaseTool] = {}
+        self._search_paths: list[Path] = self._compute_search_paths(self._config)
+
+        self._available: dict[str, type[BaseTool]] = {
+            cls.get_name(): cls for cls in self._iter_tool_classes(self._search_paths)
+        }
+        self._integrate_mcp()
+
+    @property
+    def _config(self) -> VibeConfig:
+        return self._config_getter()
+
+    @staticmethod
+    def _compute_search_paths(config: VibeConfig) -> list[Path]:
+        paths: list[Path] = [DEFAULT_TOOL_DIR.path]
+
+        paths.extend(config.tool_paths)
+
+        if (tools_dir := resolve_local_tools_dir(Path.cwd())) is not None:
+            paths.append(tools_dir)
+
+        paths.append(GLOBAL_TOOLS_DIR.path)
+
+        unique: list[Path] = []
+        seen: set[Path] = set()
+        for p in paths:
+            rp = p.resolve()
+            if rp not in seen:
+                seen.add(rp)
+                unique.append(rp)
+        return unique
+
+    @staticmethod
+    def _iter_tool_classes(search_paths: list[Path]) -> Iterator[type[BaseTool]]:
+        """Iterate over all search_paths to find tool classes.
+
+        Note: if a search path is not a directory, it is treated as a single tool file.
+        """
+        for base in search_paths:
+            if not base.is_dir() and base.name.endswith(".py"):
+                if tools := ToolManager._load_tools_from_file(base):
+                    for tool in tools:
+                        yield tool
+
+            for path in base.rglob("*.py"):
+                if tools := ToolManager._load_tools_from_file(path):
+                    for tool in tools:
+                        yield tool
+
+    @staticmethod
+    def _load_tools_from_file(file_path: Path) -> list[type[BaseTool]] | None:
+        if not file_path.is_file():
+            return
+        name = file_path.name
+        if name.startswith("_"):
+            return
+
+        module_name = _compute_module_name(file_path)
+
+        if module_name in sys.modules:
+            module = sys.modules[module_name]
+        else:
+            spec = importlib.util.spec_from_file_location(module_name, file_path)
+            if spec is None or spec.loader is None:
+                return
+            module = importlib.util.module_from_spec(spec)
+            sys.modules[module_name] = module
+            try:
+                spec.loader.exec_module(module)
+            except Exception:
+                return
+
+        tools = []
+        for tool_obj in vars(module).values():
+            if not inspect.isclass(tool_obj):
+                continue
+            if not issubclass(tool_obj, BaseTool) or tool_obj is BaseTool:
+                continue
+            if inspect.isabstract(tool_obj):
+                continue
+            tools.append(tool_obj)
+        return tools
+
+    @staticmethod
+    def discover_tool_defaults(
+        search_paths: list[Path] | None = None,
+    ) -> dict[str, dict[str, Any]]:
+        if search_paths is None:
+            search_paths = [DEFAULT_TOOL_DIR.path]
+
+        defaults: dict[str, dict[str, Any]] = {}
+        for cls in ToolManager._iter_tool_classes(search_paths):
+            try:
+                tool_name = cls.get_name()
+                config_class = cls._get_tool_config_class()
+                defaults[tool_name] = config_class().model_dump(exclude_none=True)
+            except Exception as e:
+                logger.warning(
+                    "Failed to get defaults for tool %s: %s", cls.__name__, e
+                )
+                continue
+        return defaults
+
+    @property
+    def available_tools(self) -> dict[str, type[BaseTool]]:
+        if self._config.enabled_tools:
+            return {
+                name: cls
+                for name, cls in self._available.items()
+                if name_matches(name, self._config.enabled_tools)
+            }
+        if self._config.disabled_tools:
+            return {
+                name: cls
+                for name, cls in self._available.items()
+                if not name_matches(name, self._config.disabled_tools)
+            }
+        return dict(self._available)
+
+    def _integrate_mcp(self) -> None:
+        if not self._config.mcp_servers:
+            return
+        run_sync(self._integrate_mcp_async())
+
+    async def _integrate_mcp_async(self) -> None:
+        try:
+            http_count = 0
+            stdio_count = 0
+
+            for srv in self._config.mcp_servers:
+                match srv.transport:
+                    case "http" | "streamable-http":
+                        http_count += await self._register_http_server(srv)
+                    case "stdio":
+                        stdio_count += await self._register_stdio_server(srv)
+                    case _:
+                        logger.warning("Unsupported MCP transport: %r", srv.transport)
+
+            logger.info(
+                "MCP integration registered %d tools (http=%d, stdio=%d)",
+                http_count + stdio_count,
+                http_count,
+                stdio_count,
+            )
+        except Exception as exc:
+            logger.warning("Failed to integrate MCP tools: %s", exc)
+
+    async def _register_http_server(self, srv: MCPHttp | MCPStreamableHttp) -> int:
+        url = (srv.url or "").strip()
+        if not url:
+            logger.warning("MCP server '%s' missing url for http transport", srv.name)
+            return 0
+
+        headers = srv.http_headers()
+        try:
+            tools: list[RemoteTool] = await list_tools_http(
+                url, headers=headers, startup_timeout_sec=srv.startup_timeout_sec
+            )
+        except Exception as exc:
+            logger.warning("MCP HTTP discovery failed for %s: %s", url, exc)
+            return 0
+
+        added = 0
+        for remote in tools:
+            try:
+                proxy_cls = create_mcp_http_proxy_tool_class(
+                    url=url,
+                    remote=remote,
+                    alias=srv.name,
+                    server_hint=srv.prompt,
+                    headers=headers,
+                    startup_timeout_sec=srv.startup_timeout_sec,
+                    tool_timeout_sec=srv.tool_timeout_sec,
+                )
+                self._available[proxy_cls.get_name()] = proxy_cls
+                added += 1
+            except Exception as exc:
+                logger.warning(
+                    "Failed to register MCP HTTP tool '%s' from %s: %r",
+                    getattr(remote, "name", "<unknown>"),
+                    url,
+                    exc,
+                )
+        return added
+
+    async def _register_stdio_server(self, srv: MCPStdio) -> int:
+        cmd = srv.argv()
+        if not cmd:
+            logger.warning("MCP stdio server '%s' has invalid/empty command", srv.name)
+            return 0
+
+        try:
+            tools: list[RemoteTool] = await list_tools_stdio(
+                cmd, env=srv.env or None, startup_timeout_sec=srv.startup_timeout_sec
+            )
+        except Exception as exc:
+            logger.warning("MCP stdio discovery failed for %r: %s", cmd, exc)
+            return 0
+
+        added = 0
+        for remote in tools:
+            try:
+                proxy_cls = create_mcp_stdio_proxy_tool_class(
+                    command=cmd,
+                    remote=remote,
+                    alias=srv.name,
+                    server_hint=srv.prompt,
+                    env=srv.env or None,
+                    startup_timeout_sec=srv.startup_timeout_sec,
+                    tool_timeout_sec=srv.tool_timeout_sec,
+                )
+                self._available[proxy_cls.get_name()] = proxy_cls
+                added += 1
+            except Exception as exc:
+                logger.warning(
+                    "Failed to register MCP stdio tool '%s' from %r: %r",
+                    getattr(remote, "name", "<unknown>"),
+                    cmd,
+                    exc,
+                )
+        return added
+
+    def get_tool_config(self, tool_name: str) -> BaseToolConfig:
+        tool_class = self._available.get(tool_name)
+
+        if tool_class:
+            config_class = tool_class._get_tool_config_class()
+            default_config = config_class()
+        else:
+            config_class = BaseToolConfig
+            default_config = BaseToolConfig()
+
+        user_overrides = self._config.tools.get(tool_name)
+        if user_overrides is None:
+            merged_dict = default_config.model_dump()
+        else:
+            merged_dict = {**default_config.model_dump(), **user_overrides.model_dump()}
+
+        return config_class.model_validate(merged_dict)
+
+    def get(self, tool_name: str) -> BaseTool:
+        """Get a tool instance, creating it lazily on first call.
+
+        Raises:
+            NoSuchToolError: If the requested tool is not available.
+        """
+        if tool_name in self._instances:
+            return self._instances[tool_name]
+
+        if tool_name not in self._available:
+            raise NoSuchToolError(
+                f"Unknown tool: {tool_name}. Available: {list(self._available.keys())}"
+            )
+
+        tool_class = self._available[tool_name]
+        tool_config = self.get_tool_config(tool_name)
+        self._instances[tool_name] = tool_class.from_config(tool_config)
+        return self._instances[tool_name]
+
+    def reset_all(self) -> None:
+        self._instances.clear()
+
+    def invalidate_tool(self, tool_name: str) -> None:
+        self._instances.pop(tool_name, None)