kivi-cli 2.0.6__py3-none-any.whl

kiviai/services/tools/copilot_tools/agent_orchestration/tools.py

@@ -0,0 +1,296 @@
+"""
+Agent orchestration: launch_task, read_agent, list_agents, execute_skill.
+
+Implements a hierarchical manager-worker pattern where a main agent delegates
+tasks to isolated sub-agents.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Callable, Optional
+
+
+# ---------------------------------------------------------------------------
+# Enums & constants
+# ---------------------------------------------------------------------------
+
+class AgentType(str, Enum):
+    EXPLORE = "explore"
+    TASK = "task"
+    GENERAL_PURPOSE = "general-purpose"
+    CODE_REVIEW = "code-review"
+
+
+class AgentStatus(str, Enum):
+    RUNNING = "running"
+    IDLE = "idle"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+DEFAULT_MODELS: dict[AgentType, str] = {
+    AgentType.EXPLORE: "claude-haiku-4.5",
+    AgentType.TASK: "claude-haiku-4.5",
+    AgentType.GENERAL_PURPOSE: "claude-sonnet-4.6",
+    AgentType.CODE_REVIEW: "claude-sonnet-4.6",
+}
+
+
+# ---------------------------------------------------------------------------
+# Agent data structures
+# ---------------------------------------------------------------------------
+
+@dataclass
+class Turn:
+    turn_index: int
+    role: str
+    content: str
+
+
+@dataclass
+class Agent:
+    agent_id: str
+    name: str
+    agent_type: AgentType
+    description: str
+    model: str
+    prompt: str
+    status: AgentStatus = AgentStatus.RUNNING
+    turns: list[Turn] = field(default_factory=list)
+    _thread: Optional[threading.Thread] = field(default=None, repr=False)
+    _handler: Optional[Callable[..., str]] = field(default=None, repr=False)
+
+
+# ---------------------------------------------------------------------------
+# AgentManager — singleton registry
+# ---------------------------------------------------------------------------
+
+class AgentManager:
+    def __init__(self) -> None:
+        self._agents: dict[str, Agent] = {}
+        self._lock = threading.Lock()
+
+    def register(self, agent: Agent) -> None:
+        with self._lock:
+            self._agents[agent.agent_id] = agent
+
+    def get(self, agent_id: str) -> Optional[Agent]:
+        return self._agents.get(agent_id)
+
+    def all_agents(self, *, include_completed: bool = True) -> list[Agent]:
+        agents = list(self._agents.values())
+        if not include_completed:
+            agents = [
+                a for a in agents
+                if a.status in (AgentStatus.RUNNING, AgentStatus.IDLE)
+            ]
+        return agents
+
+
+_manager = AgentManager()
+
+
+# ---------------------------------------------------------------------------
+# Default handler (stub — replace with real LLM call in production)
+# ---------------------------------------------------------------------------
+
+def _default_handler(agent: Agent) -> None:
+    """Placeholder that simulates agent work.  Override in production."""
+    agent.turns.append(
+        Turn(
+            turn_index=0,
+            role="assistant",
+            content=(
+                f"[stub] Agent '{agent.name}' ({agent.agent_type.value}) "
+                f"received prompt: {agent.prompt[:200]}..."
+            ),
+        )
+    )
+    agent.status = AgentStatus.COMPLETED
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+def launch_task(
+    name: str,
+    prompt: str,
+    agent_type: str,
+    description: str,
+    *,
+    model: Optional[str] = None,
+    mode: str = "background",
+    handler: Optional[Callable[..., None]] = None,
+) -> dict:
+    """Launch a specialised sub-agent.
+
+    Parameters
+    ----------
+    name:
+        Short human-readable name (e.g. ``"auth-flow"``).
+    prompt:
+        Complete task description with all context.
+    agent_type:
+        ``"explore"`` | ``"task"`` | ``"general-purpose"`` | ``"code-review"``.
+    description:
+        3-5 word UI summary.
+    model:
+        Override the default model for the agent type.
+    mode:
+        ``"background"`` (default) or ``"sync"``.
+    handler:
+        Custom callable ``(Agent) -> None``.  Defaults to a stub.
+
+    Returns
+    -------
+    dict with ``agent_id`` and ``status``.
+    """
+    atype = AgentType(agent_type)
+    agent_id = f"{name}-{uuid.uuid4().hex[:4]}"
+
+    agent = Agent(
+        agent_id=agent_id,
+        name=name,
+        agent_type=atype,
+        description=description,
+        model=model or DEFAULT_MODELS[atype],
+        prompt=prompt,
+    )
+    _manager.register(agent)
+
+    work_fn = handler or _default_handler
+
+    if mode == "sync":
+        try:
+            work_fn(agent)
+        except Exception as exc:
+            agent.status = AgentStatus.FAILED
+            agent.turns.append(
+                Turn(turn_index=len(agent.turns), role="error", content=str(exc))
+            )
+        return {"agent_id": agent_id, "status": agent.status.value}
+
+    # background
+    def _run() -> None:
+        try:
+            work_fn(agent)
+        except Exception as exc:
+            agent.status = AgentStatus.FAILED
+            agent.turns.append(
+                Turn(turn_index=len(agent.turns), role="error", content=str(exc))
+            )
+
+    t = threading.Thread(target=_run, daemon=True)
+    agent._thread = t
+    t.start()
+
+    return {"agent_id": agent_id, "status": "running"}
+
+
+def read_agent(
+    agent_id: str,
+    *,
+    wait: bool = False,
+    timeout: float = 30,
+    since_turn: Optional[int] = None,
+) -> dict:
+    """Read status and results from a background agent.
+
+    Parameters
+    ----------
+    agent_id:
+        ID returned by :func:`launch_task`.
+    wait:
+        Block until the agent completes.
+    timeout:
+        Max seconds to wait (only with ``wait=True``).  Max 180.
+    since_turn:
+        Only return turns after this index (exclusive).
+    """
+    agent = _manager.get(agent_id)
+    if agent is None:
+        raise KeyError(f"No agent with id={agent_id!r}")
+
+    timeout = min(timeout, 180)
+
+    if wait and agent._thread is not None:
+        agent._thread.join(timeout=timeout)
+
+    turns = agent.turns
+    if since_turn is not None:
+        turns = [t for t in turns if t.turn_index > since_turn]
+
+    return {
+        "agent_id": agent.agent_id,
+        "status": agent.status.value,
+        "turns": [
+            {"turn_index": t.turn_index, "role": t.role, "content": t.content}
+            for t in turns
+        ],
+    }
+
+
+def list_agents(*, include_completed: bool = True) -> list[dict]:
+    """List all agents.
+
+    Parameters
+    ----------
+    include_completed:
+        When ``False``, only return ``running`` and ``idle`` agents.
+    """
+    return [
+        {
+            "agent_id": a.agent_id,
+            "name": a.name,
+            "agent_type": a.agent_type.value,
+            "status": a.status.value,
+            "description": a.description,
+        }
+        for a in _manager.all_agents(include_completed=include_completed)
+    ]
+
+
+# ---------------------------------------------------------------------------
+# Skills
+# ---------------------------------------------------------------------------
+
+# Skill registry — register custom skills with ``register_skill(name, fn)``
+_skill_registry: dict[str, Callable[..., Any]] = {}
+
+
+def register_skill(name: str, fn: Callable[..., Any]) -> None:
+    """Register a skill handler for use with :func:`execute_skill`."""
+    _skill_registry[name] = fn
+
+
+def execute_skill(skill: str) -> Any:
+    """Execute a named skill in the main conversation context.
+
+    Parameters
+    ----------
+    skill:
+        Skill name (e.g. ``"pdf"``, ``"xlsx"``, ``"customize-cloud-agent"``).
+
+    Returns
+    -------
+    Whatever the skill handler returns.
+
+    Raises
+    ------
+    KeyError
+        If the skill is not registered.
+    """
+    fn = _skill_registry.get(skill)
+    if fn is None:
+        raise KeyError(
+            f"Skill {skill!r} is not registered. "
+            f"Available: {list(_skill_registry.keys())}"
+        )
+    return fn()

kiviai/services/tools/copilot_tools/bash_execution/__init__.py

@@ -0,0 +1,20 @@
+"""
+Bash Execution Tools — interactive shell session management.
+
+Provides five tools that give a full interactive shell:
+    bash        — Run a command (sync or async), create/reuse sessions
+    write_bash  — Send input (text + key sequences) to a running session
+    read_bash   — Read accumulated output from a session
+    stop_bash   — Terminate a session and clean up resources
+    list_bash   — List all active sessions
+"""
+
+from copilot_tools.bash_execution.sessions import (
+    bash,
+    write_bash,
+    read_bash,
+    stop_bash,
+    list_bash,
+)
+
+__all__ = ["bash", "write_bash", "read_bash", "stop_bash", "list_bash"]

kiviai/services/tools/copilot_tools/bash_execution/sessions.py

@@ -0,0 +1,335 @@
+"""
+Session-oriented PTY shell management.
+
+Each session owns a pseudo-terminal (PTY), a PID, and persistent env vars.
+Sessions are identified by a string ``shellId`` and managed through a global
+``SessionManager``.
+"""
+
+from __future__ import annotations
+
+import os
+import signal
+import subprocess
+import threading
+import time
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Optional
+
+
+# ---------------------------------------------------------------------------
+# Constants & key-map
+# ---------------------------------------------------------------------------
+
+class Mode(str, Enum):
+    SYNC = "sync"
+    ASYNC = "async"
+
+
+KEY_MAP: dict[str, str] = {
+    "{up}": "\x1b[A",
+    "{down}": "\x1b[B",
+    "{right}": "\x1b[C",
+    "{left}": "\x1b[D",
+    "{enter}": "\n",
+    "{backspace}": "\x7f",
+}
+
+DEFAULT_INITIAL_WAIT = 30
+MIN_INITIAL_WAIT = 30
+MAX_INITIAL_WAIT = 600
+
+
+# ---------------------------------------------------------------------------
+# Session dataclass
+# ---------------------------------------------------------------------------
+
+@dataclass
+class Session:
+    shell_id: str
+    command: str
+    mode: Mode
+    pid: Optional[int] = None
+    process: Optional[subprocess.Popen] = None
+    detach: bool = False
+    _output_buffer: str = ""
+    _lock: threading.Lock = field(default_factory=threading.Lock)
+
+    # -- output helpers -----------------------------------------------------
+
+    def append_output(self, text: str) -> None:
+        with self._lock:
+            self._output_buffer += text
+
+    def read_output(self) -> str:
+        with self._lock:
+            buf = self._output_buffer
+            self._output_buffer = ""
+            return buf
+
+    def full_output(self) -> str:
+        with self._lock:
+            return self._output_buffer
+
+    # -- status helpers -----------------------------------------------------
+
+    @property
+    def is_running(self) -> bool:
+        if self.process is None:
+            return False
+        return self.process.poll() is None
+
+    @property
+    def exit_code(self) -> Optional[int]:
+        if self.process is None:
+            return None
+        return self.process.poll()
+
+    @property
+    def has_unread_output(self) -> bool:
+        with self._lock:
+            return len(self._output_buffer) > 0
+
+
+# ---------------------------------------------------------------------------
+# SessionManager — singleton pool of sessions
+# ---------------------------------------------------------------------------
+
+class SessionManager:
+    def __init__(self) -> None:
+        self._sessions: dict[str, Session] = {}
+        self._lock = threading.Lock()
+
+    def get_or_create(self, shell_id: Optional[str], mode: Mode) -> Session:
+        if shell_id and shell_id in self._sessions:
+            return self._sessions[shell_id]
+        sid = shell_id or f"session-{uuid.uuid4().hex[:8]}"
+        session = Session(shell_id=sid, command="", mode=mode)
+        with self._lock:
+            self._sessions[sid] = session
+        return session
+
+    def get(self, shell_id: str) -> Optional[Session]:
+        return self._sessions.get(shell_id)
+
+    def remove(self, shell_id: str) -> None:
+        with self._lock:
+            self._sessions.pop(shell_id, None)
+
+    def all_sessions(self) -> list[Session]:
+        return list(self._sessions.values())
+
+
+_manager = SessionManager()
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+def _translate_keys(raw_input: str) -> str:
+    """Replace ``{enter}``, ``{up}`` etc. with ANSI escape sequences."""
+    result = raw_input
+    for token, escape in KEY_MAP.items():
+        result = result.replace(token, escape)
+    return result
+
+
+def _collect_output(session: Session, timeout: float) -> str:
+    """Block until the process finishes *or* ``timeout`` expires, collecting stdout/stderr."""
+    if session.process is None:
+        return ""
+    try:
+        stdout, _ = session.process.communicate(timeout=timeout)
+        session.append_output(stdout)
+    except subprocess.TimeoutExpired:
+        pass
+    return session.full_output()
+
+
+def _start_process(command: str, detach: bool) -> subprocess.Popen:
+    kwargs: dict = dict(
+        shell=True,
+        executable="/bin/bash",
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+    )
+    if detach:
+        kwargs["start_new_session"] = True
+    return subprocess.Popen(command, **kwargs)
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+def bash(
+    command: str,
+    description: str,
+    *,
+    mode: str = "sync",
+    initial_wait: int = DEFAULT_INITIAL_WAIT,
+    shell_id: Optional[str] = None,
+    detach: bool = False,
+) -> dict:
+    """Run a Bash command in an interactive session.
+
+    Parameters
+    ----------
+    command:
+        The shell command to run.
+    description:
+        Human-readable description (max 100 chars), displayed in the UI.
+    mode:
+        ``"sync"`` (default) — wait for output; ``"async"`` — run in background.
+    initial_wait:
+        Seconds (30-600) to wait for initial output in sync mode.
+    shell_id:
+        Reuse an existing session.  Auto-generated when ``None``.
+    detach:
+        Async-only.  If ``True`` the process survives agent shutdown.
+
+    Returns
+    -------
+    dict with keys: ``output``, ``exitCode``, ``shellId``, ``status``.
+    """
+    initial_wait = max(MIN_INITIAL_WAIT, min(initial_wait, MAX_INITIAL_WAIT))
+    run_mode = Mode(mode)
+
+    session = _manager.get_or_create(shell_id, run_mode)
+    session.command = command
+    session.mode = run_mode
+    session.detach = detach
+
+    proc = _start_process(command, detach=(run_mode == Mode.ASYNC and detach))
+    session.process = proc
+    session.pid = proc.pid
+
+    if run_mode == Mode.SYNC:
+        output = _collect_output(session, timeout=initial_wait)
+
+        if session.is_running:
+            return {
+                "output": output,
+                "exitCode": None,
+                "shellId": session.shell_id,
+                "status": "running",
+            }
+
+        return {
+            "output": session.full_output(),
+            "exitCode": session.exit_code,
+            "shellId": session.shell_id,
+            "status": "completed",
+        }
+
+    # async — return immediately
+    return {
+        "output": "",
+        "shellId": session.shell_id,
+        "status": "running",
+    }
+
+
+def write_bash(
+    shell_id: str,
+    *,
+    input_text: Optional[str] = None,
+    delay: float = 10,
+) -> dict:
+    """Send input to a running session, then read output after *delay* seconds.
+
+    Parameters
+    ----------
+    shell_id:
+        Session to write to.
+    input_text:
+        Text / key tokens (``{enter}``, ``{up}``, …) to send.
+    delay:
+        Seconds to wait before reading the output buffer.
+    """
+    session = _manager.get(shell_id)
+    if session is None:
+        raise KeyError(f"No session with shellId={shell_id!r}")
+
+    if input_text and session.process and session.process.stdin:
+        translated = _translate_keys(input_text)
+        session.process.stdin.write(translated)
+        session.process.stdin.flush()
+
+    time.sleep(delay)
+    return {"output": session.read_output()}
+
+
+def read_bash(shell_id: str, *, delay: float = 5) -> dict:
+    """Read accumulated output from a running session.
+
+    Parameters
+    ----------
+    shell_id:
+        Session to read from.
+    delay:
+        Seconds to wait before reading.
+    """
+    session = _manager.get(shell_id)
+    if session is None:
+        raise KeyError(f"No session with shellId={shell_id!r}")
+
+    time.sleep(delay)
+
+    if session.process and session.is_running:
+        try:
+            stdout, _ = session.process.communicate(timeout=0.5)
+            session.append_output(stdout)
+        except subprocess.TimeoutExpired:
+            pass
+
+    return {"output": session.read_output()}
+
+
+def stop_bash(shell_id: str) -> dict:
+    """Terminate a Bash session and clean up resources.
+
+    Sends SIGTERM → grace period → SIGKILL if needed.
+    """
+    session = _manager.get(shell_id)
+    if session is None:
+        raise KeyError(f"No session with shellId={shell_id!r}")
+
+    if session.process and session.is_running:
+        pid = session.process.pid
+        try:
+            pgid = os.getpgid(pid)
+            os.killpg(pgid, signal.SIGTERM)
+        except (ProcessLookupError, PermissionError):
+            pass
+
+        try:
+            session.process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            try:
+                pgid = os.getpgid(pid)
+                os.killpg(pgid, signal.SIGKILL)
+            except (ProcessLookupError, PermissionError):
+                pass
+
+    _manager.remove(shell_id)
+    return {"status": "terminated", "shellId": shell_id}
+
+
+def list_bash() -> list[dict]:
+    """Return info about all active sessions."""
+    return [
+        {
+            "shellId": s.shell_id,
+            "command": s.command,
+            "mode": s.mode.value,
+            "pid": s.pid,
+            "status": "running" if s.is_running else "completed",
+            "unreadOutput": s.has_unread_output,
+        }
+        for s in _manager.all_sessions()
+    ]

kiviai/services/tools/copilot_tools/code_search/__init__.py

@@ -0,0 +1,10 @@
+"""
+Code Search Tools — fast content and filename searching.
+
+    grep        — Search file *contents* via regex (wraps ripgrep)
+    glob_search — Search file *names* via glob patterns
+"""
+
+from copilot_tools.code_search.tools import grep, glob_search
+
+__all__ = ["grep", "glob_search"]