codemaster-cli 2.2.0__py3-none-any.whl

vibe/core/tools/builtins/task.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+from typing import ClassVar
+
+from pydantic import BaseModel, Field
+
+from vibe.core.agent_loop import AgentLoop
+from vibe.core.agents.models import AgentType
+from vibe.core.config import SessionLoggingConfig, VibeConfig
+from vibe.core.tools.base import (
+    BaseTool,
+    BaseToolConfig,
+    BaseToolState,
+    InvokeContext,
+    ToolError,
+    ToolPermission,
+)
+from vibe.core.tools.ui import (
+    ToolCallDisplay,
+    ToolResultDisplay,
+    ToolUIData,
+    ToolUIDataAdapter,
+)
+from vibe.core.types import (
+    AssistantEvent,
+    Role,
+    ToolCallEvent,
+    ToolResultEvent,
+    ToolStreamEvent,
+)
+
+
+class TaskArgs(BaseModel):
+    task: str = Field(description="The task to delegate to the subagent")
+    agent: str = Field(
+        default="explore",
+        description="Name of the agent profile to use (must be a subagent)",
+    )
+
+
+class TaskResult(BaseModel):
+    response: str = Field(description="The accumulated response from the subagent")
+    turns_used: int = Field(description="Number of turns the subagent used")
+    completed: bool = Field(description="Whether the task completed normally")
+
+
+class TaskToolConfig(BaseToolConfig):
+    permission: ToolPermission = ToolPermission.ASK
+
+
+class Task(
+    BaseTool[TaskArgs, TaskResult, TaskToolConfig, BaseToolState],
+    ToolUIData[TaskArgs, TaskResult],
+):
+    description: ClassVar[str] = (
+        "Delegate a task to a subagent for independent execution. "
+        "Useful for exploration, research, or parallel work that doesn't "
+        "require user interaction. The subagent runs in-memory without "
+        "saving interaction logs."
+    )
+
+    @classmethod
+    def get_call_display(cls, event: ToolCallEvent) -> ToolCallDisplay:
+        args = event.args
+        if isinstance(args, TaskArgs):
+            return ToolCallDisplay(summary=f"Running {args.agent} agent: {args.task}")
+        return ToolCallDisplay(summary="Running subagent")
+
+    @classmethod
+    def get_result_display(cls, event: ToolResultEvent) -> ToolResultDisplay:
+        result = event.result
+        if isinstance(result, TaskResult):
+            turn_word = "turn" if result.turns_used == 1 else "turns"
+            if not result.completed:
+                return ToolResultDisplay(
+                    success=False,
+                    message=f"Agent interrupted after {result.turns_used} {turn_word}",
+                )
+            return ToolResultDisplay(
+                success=True,
+                message=f"Agent completed in {result.turns_used} {turn_word}",
+            )
+        return ToolResultDisplay(success=True, message="Agent completed")
+
+    @classmethod
+    def get_status_text(cls) -> str:
+        return "Running subagent"
+
+    async def run(
+        self, args: TaskArgs, ctx: InvokeContext | None = None
+    ) -> AsyncGenerator[ToolStreamEvent | TaskResult, None]:
+        if not ctx or not ctx.agent_manager:
+            raise ToolError("Task tool requires agent_manager in context")
+
+        agent_manager = ctx.agent_manager
+
+        try:
+            agent_profile = agent_manager.get_agent(args.agent)
+        except ValueError as e:
+            raise ToolError(f"Unknown agent: {args.agent}") from e
+
+        if agent_profile.agent_type != AgentType.SUBAGENT:
+            raise ToolError(
+                f"Agent '{args.agent}' is a {agent_profile.agent_type.value} agent. "
+                f"Only subagents can be used with the task tool. "
+                f"This is a security constraint to prevent recursive spawning."
+            )
+
+        base_config = VibeConfig.load(
+            session_logging=SessionLoggingConfig(enabled=False)
+        )
+        subagent_loop = AgentLoop(config=base_config, agent_name=args.agent)
+
+        if ctx and ctx.approval_callback:
+            subagent_loop.set_approval_callback(ctx.approval_callback)
+
+        accumulated_response: list[str] = []
+        completed = True
+        try:
+            async for event in subagent_loop.act(args.task):
+                if isinstance(event, AssistantEvent) and event.content:
+                    accumulated_response.append(event.content)
+                    if event.stopped_by_middleware:
+                        completed = False
+                elif isinstance(event, ToolResultEvent):
+                    if event.skipped:
+                        completed = False
+                    elif event.result and event.tool_class:
+                        adapter = ToolUIDataAdapter(event.tool_class)
+                        display = adapter.get_result_display(event)
+                        message = f"{event.tool_name}: {display.message}"
+                        yield ToolStreamEvent(
+                            tool_name=self.get_name(),
+                            message=message,
+                            tool_call_id=ctx.tool_call_id,
+                        )
+
+            turns_used = sum(
+                msg.role == Role.assistant for msg in subagent_loop.messages
+            )
+
+        except Exception as e:
+            completed = False
+            accumulated_response.append(f"\n[Subagent error: {e}]")
+            turns_used = sum(
+                msg.role == Role.assistant for msg in subagent_loop.messages
+            )
+
+        yield TaskResult(
+            response="".join(accumulated_response),
+            turns_used=turns_used,
+            completed=completed,
+        )

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