monkeyscode 1.0.0__py3-none-any.whl

monkeyscode/agent.py

@@ -0,0 +1,772 @@
+"""
+MonkeysCode SDK — Agent Class.
+
+The primary SDK interface for Python. Handles:
+- Connecting to the model-proxy API (SSE streaming)
+- Executing custom tools locally
+- Goal mode (autonomous multi-iteration)
+- Session management
+
+Usage::
+
+    from monkeyscode import MonkeysCode
+
+    agent = MonkeysCode(api_key="mc_...")
+
+    # Full result
+    result = await agent.run("Fix all type errors")
+    print(result.summary)
+
+    # Streaming
+    async for event in agent.stream("Refactor auth"):
+        if event.type == "text":
+            print(event.content, end="", flush=True)
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from collections import defaultdict
+from collections.abc import AsyncGenerator, Callable
+from dataclasses import dataclass
+from typing import Any
+
+from monkeyscode._http import stream_sse
+from monkeyscode.events import (
+    AgentCompleteEvent,
+    AgentEvent,
+    AgentResult,
+    AgentSpawnedEvent,
+    CompleteEvent,
+    CostUpdateEvent,
+    ErrorEvent,
+    FanOutProgressEvent,
+    FileChange,
+    FileEditEvent,
+    FileReadEvent,
+    GoalResult,
+    PlanReadyEvent,
+    PlanStep,
+    SearchResultEvent,
+    StartEvent,
+    SuggestionEvent,
+    TerminalCommandEvent,
+    TextEvent,
+    ThinkingEvent,
+    TokenUsage,
+    ToolCallEvent,
+    ToolResultEvent,
+    VerificationEvent,
+)
+from monkeyscode.hooks import (
+    HookManager,
+    PostToolContext,
+    PostToolHook,
+    PreToolContext,
+    PreToolHook,
+)
+from monkeyscode.subagent import (
+    MAX_SUBAGENTS,
+    SubagentOptions,
+    SubagentResult,
+    spawn_subagent,
+)
+from monkeyscode.tools import ToolDefinition
+
+
+@dataclass
+class MonkeysCodeOptions:
+    """Configuration options for the MonkeysCode agent."""
+
+    model: str = "auto"
+    api_key: str | None = None
+    working_directory: str | None = None
+    proxy_url: str | None = None
+    max_turns: int = 30
+    timeout: int = 300_000
+    permissions_mode: str = "ask"
+    max_cost: float | None = None
+    max_retries: int = 3
+    debug: bool = False
+    runtime: str = "local"  # 'local' | 'container'
+    sandbox: dict[str, Any] | None = None
+    telemetry: dict[str, Any] | None = None
+
+
+class MonkeysCode:
+    """
+    MonkeysCode agent — the primary SDK interface.
+
+    Connects to the model-proxy API and streams events back.
+    Supports custom tools, goal mode, and session management.
+    """
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        model: str = "auto",
+        working_directory: str | None = None,
+        proxy_url: str | None = None,
+        max_turns: int = 30,
+        timeout: int = 300_000,
+        permissions_mode: str = "ask",
+        max_cost: float | None = None,
+        max_retries: int = 3,
+        debug: bool = False,
+        runtime: str = "local",
+        sandbox: dict[str, Any] | None = None,
+        telemetry: dict[str, Any] | None = None,
+    ) -> None:
+        self._api_key = api_key or os.environ.get("MONKEYSCODE_API_KEY")
+        self._model = model
+        self._working_directory = working_directory or os.getcwd()
+        self._proxy_url = (
+            proxy_url
+            or os.environ.get("MONKEYSCODE_PROXY_URL")
+            or "https://models.monkeyscode.com"
+        )
+        self._max_turns = max_turns
+        self._timeout = timeout
+        self._permissions_mode = permissions_mode
+        self._max_cost = max_cost
+        self._max_retries = max_retries
+        self._debug = debug
+
+        self._session_id: str | None = None
+        self._total_cost: float = 0.0
+        self._total_tokens = TokenUsage()
+        self._custom_tools: dict[str, ToolDefinition] = {}
+        self._closed = False
+        self._listeners: dict[str, list[Callable]] = defaultdict(list)
+        self._hook_manager = HookManager()
+        self._subagent_count = 0
+        self._subagent_depth = 0
+        self._runtime = runtime
+        self._sandbox = sandbox
+        self._telemetry = telemetry
+
+    # ── Properties ───────────────────────────────────────────────────
+
+    @property
+    def session_id(self) -> str | None:
+        """Current session ID."""
+        return self._session_id
+
+    @property
+    def total_cost(self) -> float:
+        """Total accumulated cost in USD."""
+        return self._total_cost
+
+    @property
+    def total_tokens(self) -> TokenUsage:
+        """Total accumulated token usage."""
+        return TokenUsage(
+            input=self._total_tokens.input,
+            output=self._total_tokens.output,
+            thinking=self._total_tokens.thinking,
+            total=self._total_tokens.total,
+        )
+
+    @property
+    def model(self) -> str:
+        """Current model."""
+        return self._model
+
+    # ── Tools ────────────────────────────────────────────────────────
+
+    def add_tool(self, tool: ToolDefinition) -> None:
+        """Register a custom tool."""
+        self._custom_tools[tool.name] = tool
+
+    def remove_tool(self, name: str) -> None:
+        """Unregister a custom tool."""
+        self._custom_tools.pop(name, None)
+
+    def list_tools(self) -> list[str]:
+        """List registered custom tools."""
+        return list(self._custom_tools.keys())
+
+    # ── Event Listeners ──────────────────────────────────────────────
+
+    def on(self, event_type: str, callback: Callable) -> MonkeysCode:
+        """Register an event listener.
+
+        Example::
+
+            agent.on("text", lambda e: print(e.content, end=""))
+            agent.on("error", lambda e: print(f"Error: {e.message}"))
+        """
+        self._listeners[event_type].append(callback)
+        return self
+
+    def off(self, event_type: str, callback: Callable) -> MonkeysCode:
+        """Remove an event listener."""
+        try:
+            self._listeners[event_type].remove(callback)
+        except ValueError:
+            pass
+        return self
+
+    def _emit(self, event: AgentEvent) -> None:
+        """Notify all listeners for an event type."""
+        event_type = getattr(event, "type", None)
+        if event_type and event_type in self._listeners:
+            for cb in self._listeners[event_type]:
+                try:
+                    cb(event)
+                except Exception:
+                    pass  # Swallow listener errors
+
+    # ── Hooks (SDK Gap #1) ────────────────────────────────────────────
+
+    def on_pre_tool_use(self, hook: PreToolHook) -> MonkeysCode:
+        """Register a pre-tool hook to intercept tool calls.
+
+        Example::
+
+            async def block_rm(ctx):
+                if ctx.tool == "shell_command" and "rm -rf" in (ctx.args.get("command") or ""):
+                    return PreToolResult(action="block", reason="Blocked")
+                return PreToolResult(action="allow")
+
+            agent.on_pre_tool_use(block_rm)
+        """
+        self._hook_manager.add_pre_hook(hook)
+        return self
+
+    def on_post_tool_use(self, hook: PostToolHook) -> MonkeysCode:
+        """Register a post-tool hook to observe tool results."""
+        self._hook_manager.add_post_hook(hook)
+        return self
+
+    def off_pre_tool_use(self, hook: PreToolHook) -> bool:
+        """Remove a pre-tool hook. Returns True if found."""
+        return self._hook_manager.remove_pre_hook(hook)
+
+    def off_post_tool_use(self, hook: PostToolHook) -> bool:
+        """Remove a post-tool hook. Returns True if found."""
+        return self._hook_manager.remove_post_hook(hook)
+
+    def clear_hooks(self) -> None:
+        """Remove all hooks."""
+        self._hook_manager.clear_all()
+
+    # ── Subagent Spawning (SDK Gap #2) ────────────────────────────────
+
+    async def spawn(self, options: SubagentOptions | None = None, **kwargs: Any) -> SubagentResult:
+        """Spawn a focused child agent for a specific task.
+
+        Example::
+
+            sub = await agent.spawn(SubagentOptions(
+                task="Find SQL injection risks",
+                preset="reviewer",
+            ))
+            print(sub.summary)
+        """
+        if options is None:
+            options = SubagentOptions(**kwargs)
+
+        if self._subagent_count >= MAX_SUBAGENTS:
+            raise RuntimeError(
+                f"Max subagents reached ({MAX_SUBAGENTS}). Wait for one to finish.",
+            )
+
+        self._subagent_count += 1
+        try:
+            parent_opts = {
+                "model": self._model,
+                "api_key": self._api_key,
+                "working_directory": self._working_directory,
+                "proxy_url": self._proxy_url,
+                "debug": self._debug,
+            }
+            result = await spawn_subagent(parent_opts, options, self._subagent_depth)
+            self._total_cost += result.result.cost
+            return result
+        finally:
+            self._subagent_count -= 1
+
+    # ── Run (Full Result) ────────────────────────────────────────────
+
+    async def run(self, prompt: str) -> AgentResult:
+        """
+        Send a prompt and get the full result after the agent finishes.
+
+        Args:
+            prompt: The instruction for the agent.
+
+        Returns:
+            AgentResult with success status, summary, files changed, etc.
+        """
+        self._assert_not_closed()
+
+        events: list[AgentEvent] = []
+        async for event in self.stream(prompt):
+            events.append(event)
+
+        # Find the complete event
+        for event in events:
+            if isinstance(event, CompleteEvent):
+                return AgentResult(
+                    success=True,
+                    summary=event.summary,
+                    files_changed=event.files_changed,
+                    tokens=event.tokens,
+                    cost=event.cost,
+                    duration_ms=event.duration_ms,
+                    session_id=self._session_id or "",
+                    exit_code=0,
+                    model=self._model,
+                )
+
+        # Find error
+        for event in events:
+            if isinstance(event, ErrorEvent):
+                return AgentResult(
+                    success=False,
+                    summary=event.message,
+                    exit_code=1,
+                    model=self._model,
+                )
+
+        return AgentResult(
+            success=False,
+            summary="Agent failed with no completion event",
+            exit_code=1,
+            model=self._model,
+        )
+
+    # ── Stream (Async Iterator) ──────────────────────────────────────
+
+    async def stream(self, prompt: str) -> AsyncGenerator[AgentEvent, None]:
+        """
+        Send a prompt and stream events as they arrive.
+
+        Args:
+            prompt: The instruction for the agent.
+
+        Yields:
+            AgentEvent instances (TextEvent, ToolCallEvent, etc.)
+        """
+        self._assert_not_closed()
+
+        if not self._api_key:
+            yield ErrorEvent(
+                timestamp=int(time.time() * 1000),
+                code="AUTH_MISSING",
+                message="No API key provided. Set api_key or MONKEYSCODE_API_KEY env var.",
+            )
+            return
+
+        start_time = time.monotonic()
+
+        # Build request
+        body: dict[str, Any] = {
+            "prompt": prompt,
+            "model": self._model,
+            "workspace": self._working_directory,
+            "sessionId": self._session_id,
+            "maxTurns": self._max_turns,
+            "tools": [t.to_json_schema() for t in self._custom_tools.values()],
+            "permissions": {"mode": self._permissions_mode},
+        }
+
+        if self._max_cost is not None:
+            body["permissions"]["maxCost"] = self._max_cost
+
+        yield StartEvent(
+            timestamp=int(time.time() * 1000),
+            session_id=self._session_id or "pending",
+            model=self._model,
+        )
+
+        try:
+            files_changed: list[FileChange] = []
+            tokens = TokenUsage()
+            cost = 0.0
+            summary = ""
+
+            async for data in stream_sse(
+                f"{self._proxy_url}/v1/agent/run",
+                headers={"Authorization": f"Bearer {self._api_key}"},
+                json_body=body,
+                timeout=self._timeout / 1000,
+            ):
+                event = self._map_sse_event(data, files_changed)
+                if event is None:
+                    continue
+
+                # Track session ID
+                if "sessionId" in data:
+                    self._session_id = data["sessionId"]
+
+                # Track tokens and cost
+                if isinstance(event, CompleteEvent):
+                    tokens = event.tokens
+                    cost = event.cost
+                    summary = event.summary
+                    self._total_tokens.input += tokens.input
+                    self._total_tokens.output += tokens.output
+                    self._total_tokens.thinking += tokens.thinking
+                    self._total_tokens.total += tokens.total
+                    self._total_cost += cost
+
+                # Track mid-run cost
+                if isinstance(event, CostUpdateEvent):
+                    self._total_tokens.input = event.tokens_in
+                    self._total_tokens.output = event.tokens_out
+                    self._total_tokens.thinking = event.tokens_thinking
+
+                # ── Pre-tool hooks ──
+                if isinstance(event, ToolCallEvent) and self._hook_manager.has_hooks:
+                    hook_ctx = PreToolContext(
+                        tool=event.tool,
+                        args=event.args,
+                        run_id=f"run_{int(start_time)}",
+                        turn_index=0,
+                        session_id=self._session_id,
+                        timestamp=event.timestamp,
+                    )
+                    hook_result = await self._hook_manager.run_pre_hooks(hook_ctx)
+
+                    if hook_result.action == "block":
+                        yield event
+                        yield ToolResultEvent(
+                            timestamp=int(time.time() * 1000),
+                            tool=event.tool,
+                            result=f"Blocked by hook: {hook_result.reason or 'no reason'}",
+                            success=False,
+                            duration_ms=0,
+                        )
+                        continue
+
+                    if hook_result.action == "modify" and hook_result.args:
+                        event = ToolCallEvent(
+                            timestamp=event.timestamp,
+                            tool=event.tool,
+                            args=hook_result.args,
+                            tool_call_id=getattr(event, 'tool_call_id', ''),
+                        )
+
+                # Handle custom tool calls
+                if isinstance(event, ToolCallEvent) and event.tool in self._custom_tools:
+                    yield event
+                    tool = self._custom_tools[event.tool]
+                    tool_start = time.monotonic()
+                    try:
+                        result = await tool.execute(event.args)
+                        tool_result_event = ToolResultEvent(
+                            timestamp=int(time.time() * 1000),
+                            tool=event.tool,
+                            result=str(result),
+                            success=True,
+                            duration_ms=int((time.monotonic() - tool_start) * 1000),
+                        )
+                        yield tool_result_event
+
+                        # ── Post-tool hooks ──
+                        if self._hook_manager.has_hooks:
+                            await self._hook_manager.run_post_hooks(PostToolContext(
+                                tool=event.tool,
+                                args=event.args,
+                                result=result,
+                                is_error=False,
+                                duration_ms=int((time.monotonic() - tool_start) * 1000),
+                                run_id=f"run_{int(start_time)}",
+                                turn_index=0,
+                                session_id=self._session_id,
+                                timestamp=tool_result_event.timestamp,
+                            ))
+                    except Exception as e:
+                        tool_result_event = ToolResultEvent(
+                            timestamp=int(time.time() * 1000),
+                            tool=event.tool,
+                            result=str(e),
+                            success=False,
+                            duration_ms=int((time.monotonic() - tool_start) * 1000),
+                        )
+                        yield tool_result_event
+
+                        if self._hook_manager.has_hooks:
+                            await self._hook_manager.run_post_hooks(PostToolContext(
+                                tool=event.tool,
+                                args=event.args,
+                                result=str(e),
+                                is_error=True,
+                                duration_ms=int((time.monotonic() - tool_start) * 1000),
+                                run_id=f"run_{int(start_time)}",
+                                turn_index=0,
+                                session_id=self._session_id,
+                                timestamp=tool_result_event.timestamp,
+                            ))
+                    continue
+
+                yield event
+
+            # Final completion
+            elapsed_ms = int((time.monotonic() - start_time) * 1000)
+            yield CompleteEvent(
+                timestamp=int(time.time() * 1000),
+                summary=summary,
+                files_changed=files_changed,
+                tokens=tokens,
+                cost=cost,
+                duration_ms=elapsed_ms,
+            )
+
+        except Exception as e:
+            yield ErrorEvent(
+                timestamp=int(time.time() * 1000),
+                code="CONNECTION_ERROR",
+                message=str(e),
+            )
+
+    # ── Goal Mode ────────────────────────────────────────────────────
+
+    async def goal(
+        self,
+        prompt: str,
+        *,
+        verify_command: str | None = None,
+        max_iterations: int = 10,
+        timeout: int = 300_000,
+        max_cost: float | None = None,
+    ) -> GoalResult:
+        """
+        Run the agent in goal mode — iterates until the goal is met.
+
+        Args:
+            prompt: The goal to achieve.
+            verify_command: Command to verify success (e.g., "pytest").
+            max_iterations: Maximum iterations.
+            timeout: Timeout in ms.
+            max_cost: Maximum cost in USD.
+
+        Returns:
+            GoalResult with success status and iteration details.
+        """
+        start_time = time.monotonic()
+        iteration_summaries: list[str] = []
+        last_result: AgentResult | None = None
+        goal_met = False
+
+        goal_prompt = prompt
+        if verify_command:
+            goal_prompt = (
+                f"{prompt}\n\nVerification: Run `{verify_command}` after each attempt. "
+                "Keep iterating until it passes."
+            )
+
+        for iteration in range(1, max_iterations + 1):
+            elapsed = (time.monotonic() - start_time) * 1000
+            if elapsed > timeout:
+                break
+            if max_cost and self._total_cost >= max_cost:
+                break
+
+            if iteration == 1:
+                iter_prompt = goal_prompt
+            else:
+                prev = last_result.summary if last_result else 'unknown'
+                iter_prompt = f"Continue working on the goal. Previous: {prev}"
+
+            last_result = await self.run(iter_prompt)
+            iteration_summaries.append(last_result.summary)
+
+            if last_result.success and last_result.exit_code == 0:
+                goal_met = True
+                break
+
+        return GoalResult(
+            success=last_result.success if last_result else False,
+            summary=last_result.summary if last_result else "No iterations completed",
+            files_changed=last_result.files_changed if last_result else [],
+            tokens=self.total_tokens,
+            cost=self._total_cost,
+            duration_ms=int((time.monotonic() - start_time) * 1000),
+            session_id=self._session_id or "",
+            exit_code=0 if goal_met else 1,
+            model=self._model,
+            goal_met=goal_met,
+            iterations=len(iteration_summaries),
+            iteration_summaries=iteration_summaries,
+        )
+
+    # ── Lifecycle ────────────────────────────────────────────────────
+
+    def close(self) -> None:
+        """Close the agent. No further runs can be made."""
+        self._closed = True
+
+    async def __aenter__(self) -> MonkeysCode:
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        self.close()
+
+    # ── Private ──────────────────────────────────────────────────────
+
+    def _assert_not_closed(self) -> None:
+        if self._closed:
+            raise RuntimeError("Agent is closed")
+
+    def _map_sse_event(
+        self,
+        data: dict[str, Any],
+        files_changed: list[FileChange],
+    ) -> AgentEvent | None:
+        """Map a raw SSE JSON object to a typed event."""
+        event_type = data.get("type") or data.get("kind") or ""
+        ts = data.get("timestamp") or int(time.time() * 1000)
+
+        if event_type in ("text", "run.token", "chunk", "streaming_chunk"):
+            return TextEvent(
+                timestamp=ts,
+                content=data.get("content") or data.get("data") or "",
+            )
+
+        if event_type == "thinking":
+            return ThinkingEvent(
+                timestamp=ts,
+                content=data.get("content") or "",
+            )
+
+        if event_type in ("tool_call", "run.tool_call", "tool_call_start"):
+            return ToolCallEvent(
+                timestamp=ts,
+                tool=data.get("tool") or data.get("name") or "",
+                args=data.get("args") or data.get("input") or {},
+                summary=data.get("summary") or "",
+            )
+
+        if event_type in ("tool_result", "run.tool_result", "tool_call_complete"):
+            return ToolResultEvent(
+                timestamp=ts,
+                tool=data.get("tool") or data.get("name") or "",
+                result=str(data.get("result") or data.get("output") or ""),
+                success=data.get("success", True),
+                duration_ms=data.get("durationMs") or data.get("duration") or 0,
+            )
+
+        if event_type in ("file_edit", "run.diff"):
+            fc = FileChange(
+                path=data.get("path") or "",
+                action="modified",
+                diff=data.get("diff") or "",
+                lines_added=data.get("linesAdded") or 0,
+                lines_removed=data.get("linesRemoved") or 0,
+            )
+            files_changed.append(fc)
+            return FileEditEvent(
+                timestamp=ts,
+                path=fc.path,
+                diff=fc.diff or "",
+                lines_added=fc.lines_added,
+                lines_removed=fc.lines_removed,
+            )
+
+        if event_type == "file_read":
+            return FileReadEvent(
+                timestamp=ts,
+                path=data.get("path") or "",
+                lines=data.get("lines") or 0,
+            )
+
+        if event_type == "terminal_command":
+            return TerminalCommandEvent(
+                timestamp=ts,
+                command=data.get("command") or "",
+                exit_code=data.get("exitCode") or 0,
+                output=data.get("output") or "",
+            )
+
+        if event_type == "search_result":
+            return SearchResultEvent(
+                timestamp=ts,
+                query=data.get("query") or "",
+                matches=data.get("matches") or [],
+            )
+
+        # Daemon events
+        if event_type in ("plan_ready", "run.plan_ready"):
+            raw_steps = data.get("steps") or []
+            return PlanReadyEvent(
+                timestamp=ts,
+                steps=[PlanStep(**s) if isinstance(s, dict) else s for s in raw_steps],
+                summary=data.get("summary") or "",
+            )
+
+        if event_type in ("cost_update", "run.cost"):
+            return CostUpdateEvent(
+                timestamp=ts,
+                tokens_in=data.get("tokens_in") or data.get("tokensIn") or 0,
+                tokens_out=data.get("tokens_out") or data.get("tokensOut") or 0,
+                tokens_thinking=data.get("tokens_thinking") or data.get("tokensThinking") or 0,
+                cost=data.get("usd") or data.get("cost") or 0.0,
+            )
+
+        if event_type in ("agent_spawned", "run.agent_spawned"):
+            return AgentSpawnedEvent(
+                timestamp=ts,
+                agent_id=data.get("agentId") or data.get("agent_id") or "",
+                task=data.get("task") or "",
+            )
+
+        if event_type in ("agent_complete", "run.agent_complete"):
+            return AgentCompleteEvent(
+                timestamp=ts,
+                agent_id=data.get("agentId") or data.get("agent_id") or "",
+                success=data.get("success", True),
+                summary=data.get("summary") or "",
+            )
+
+        if event_type in ("fan_out_progress", "run.fan_out_progress"):
+            return FanOutProgressEvent(
+                timestamp=ts,
+                total=data.get("total") or 0,
+                completed=data.get("completed") or 0,
+                running=data.get("running") or 0,
+            )
+
+        if event_type in ("verification", "run.verification"):
+            return VerificationEvent(
+                timestamp=ts,
+                command=data.get("command") or "",
+                passed=data.get("passed", False),
+                output=data.get("output") or "",
+            )
+
+        if event_type in ("suggestion", "run.suggestion"):
+            return SuggestionEvent(
+                timestamp=ts,
+                suggestions=data.get("suggestions") or [],
+            )
+
+        if event_type == "error":
+            return ErrorEvent(
+                timestamp=ts,
+                code=data.get("code") or "UNKNOWN",
+                message=data.get("message") or "",
+            )
+
+        if event_type in ("complete", "run.complete", "done"):
+            tokens_data = data.get("tokens") or {}
+            return CompleteEvent(
+                timestamp=ts,
+                summary=data.get("summary") or "",
+                files_changed=files_changed.copy(),
+                tokens=TokenUsage(
+                    input=tokens_data.get("input") or 0,
+                    output=tokens_data.get("output") or 0,
+                    thinking=tokens_data.get("thinking") or 0,
+                    total=tokens_data.get("total") or 0,
+                ),
+                cost=data.get("cost") or 0.0,
+                duration_ms=data.get("durationMs") or data.get("duration") or 0,
+            )
+
+        return None
+