monkeyscode 1.0.0__py3-none-any.whl

monkeyscode/mcp.py

@@ -0,0 +1,283 @@
+"""
+MonkeysCode SDK — MCP Server Manager.
+
+Launch and manage Model Context Protocol (MCP) server processes.
+Discover and call tools exposed by MCP servers.
+
+Usage::
+
+    from monkeyscode.mcp import McpManager
+
+    mcp = McpManager()
+    await mcp.connect("postgres", McpServerConfig(
+        command="npx",
+        args=["@modelcontextprotocol/server-postgres"],
+        env={"DATABASE_URL": "postgresql://localhost/mydb"},
+    ))
+
+    tools = await mcp.list_tools("postgres")
+    result = await mcp.call_tool("postgres", "query", {"sql": "SELECT 1"})
+
+    mcp.disconnect_all()
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+from dataclasses import dataclass, field
+from typing import Any
+
+# ── Types ────────────────────────────────────────────────────────────
+
+
+@dataclass
+class McpServerConfig:
+    """MCP server configuration."""
+
+    command: str
+    args: list[str] = field(default_factory=list)
+    env: dict[str, str] = field(default_factory=dict)
+    cwd: str | None = None
+    timeout: float = 30.0
+
+
+@dataclass
+class McpTool:
+    """An MCP tool descriptor."""
+
+    name: str
+    description: str
+    input_schema: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class McpToolResult:
+    """Result from calling an MCP tool."""
+
+    content: list[dict[str, Any]] = field(default_factory=list)
+    is_error: bool = False
+
+
+@dataclass
+class McpServerStatus:
+    """Status of an MCP server."""
+
+    name: str
+    pid: int | None = None
+    connected: bool = False
+    tool_count: int = 0
+
+
+# ── JSON-RPC Helpers ─────────────────────────────────────────────────
+
+
+class _JsonRpcTransport:
+    """Communicate with an MCP server over stdin/stdout JSON-RPC."""
+
+    def __init__(self, process: asyncio.subprocess.Process) -> None:
+        self._process = process
+        self._request_id = 0
+
+    async def request(self, method: str, params: dict[str, Any] | None = None) -> Any:
+        self._request_id += 1
+        msg = {
+            "jsonrpc": "2.0",
+            "id": self._request_id,
+            "method": method,
+            "params": params or {},
+        }
+
+        stdin = self._process.stdin
+        stdout = self._process.stdout
+        assert stdin is not None and stdout is not None
+
+        line = json.dumps(msg) + "\n"
+        stdin.write(line.encode())
+        await stdin.drain()
+
+        response_line = await asyncio.wait_for(stdout.readline(), timeout=30.0)
+        if not response_line:
+            raise RuntimeError("MCP server closed connection")
+
+        data = json.loads(response_line.decode())
+        if "error" in data:
+            raise RuntimeError(f"MCP error: {data['error']}")
+        return data.get("result")
+
+    async def notify(self, method: str, params: dict[str, Any] | None = None) -> None:
+        msg = {
+            "jsonrpc": "2.0",
+            "method": method,
+            "params": params or {},
+        }
+        stdin = self._process.stdin
+        assert stdin is not None
+        line = json.dumps(msg) + "\n"
+        stdin.write(line.encode())
+        await stdin.drain()
+
+
+# ── MCP Manager ──────────────────────────────────────────────────────
+
+
+class McpManager:
+    """
+    Manages MCP server lifecycle and tool invocations.
+
+    Launches servers as child processes, communicates via JSON-RPC over
+    stdin/stdout, and provides tool discovery and invocation.
+    """
+
+    def __init__(self) -> None:
+        self._servers: dict[str, _McpConnection] = {}
+
+    @property
+    def server_names(self) -> list[str]:
+        return list(self._servers.keys())
+
+    @property
+    def server_count(self) -> int:
+        return len(self._servers)
+
+    async def connect(self, name: str, config: McpServerConfig) -> McpServerStatus:
+        """Launch and connect to an MCP server."""
+        env = {**os.environ, **config.env}
+
+        process = await asyncio.create_subprocess_exec(
+            config.command,
+            *config.args,
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            cwd=config.cwd,
+            env=env,
+        )
+
+        transport = _JsonRpcTransport(process)
+
+        # Initialize the server
+        try:
+            await transport.request("initialize", {
+                "protocolVersion": "2024-11-05",
+                "capabilities": {},
+                "clientInfo": {"name": "monkeyscode-sdk-python", "version": "1.0.0"},
+            })
+            await transport.notify("notifications/initialized")
+        except Exception:
+            process.kill()
+            raise
+
+        # Discover tools
+        try:
+            result = await transport.request("tools/list")
+            tools = [
+                McpTool(
+                    name=t["name"],
+                    description=t.get("description", ""),
+                    input_schema=t.get("inputSchema", {}),
+                )
+                for t in (result.get("tools", []) if result else [])
+            ]
+        except Exception:
+            tools = []
+
+        conn = _McpConnection(
+            name=name,
+            config=config,
+            process=process,
+            transport=transport,
+            tools=tools,
+        )
+        self._servers[name] = conn
+
+        return McpServerStatus(
+            name=name,
+            pid=process.pid,
+            connected=True,
+            tool_count=len(tools),
+        )
+
+    async def disconnect(self, name: str) -> None:
+        """Disconnect and kill an MCP server."""
+        conn = self._servers.pop(name, None)
+        if conn is None:
+            return
+
+        try:
+            conn.process.terminate()
+            await asyncio.wait_for(conn.process.wait(), timeout=5.0)
+        except (asyncio.TimeoutError, ProcessLookupError):
+            conn.process.kill()
+
+    def disconnect_all(self) -> None:
+        """Disconnect all MCP servers synchronously."""
+        for name in list(self._servers.keys()):
+            conn = self._servers.pop(name, None)
+            if conn and conn.process.returncode is None:
+                try:
+                    conn.process.kill()
+                except ProcessLookupError:
+                    pass
+
+    async def list_tools(self, name: str) -> list[McpTool]:
+        """List tools available from a named server."""
+        conn = self._servers.get(name)
+        if conn is None:
+            raise KeyError(f"MCP server not connected: {name}")
+        return conn.tools
+
+    async def call_tool(
+        self,
+        server_name: str,
+        tool_name: str,
+        arguments: dict[str, Any] | None = None,
+    ) -> McpToolResult:
+        """Call a tool on a named MCP server."""
+        conn = self._servers.get(server_name)
+        if conn is None:
+            raise KeyError(f"MCP server not connected: {server_name}")
+
+        result = await conn.transport.request("tools/call", {
+            "name": tool_name,
+            "arguments": arguments or {},
+        })
+
+        if result is None:
+            return McpToolResult(is_error=True)
+
+        return McpToolResult(
+            content=result.get("content", []),
+            is_error=result.get("isError", False),
+        )
+
+    def status(self, name: str) -> McpServerStatus:
+        """Get the status of a named server."""
+        conn = self._servers.get(name)
+        if conn is None:
+            return McpServerStatus(name=name, connected=False)
+        return McpServerStatus(
+            name=name,
+            pid=conn.process.pid,
+            connected=conn.process.returncode is None,
+            tool_count=len(conn.tools),
+        )
+
+    def all_tools(self) -> list[McpTool]:
+        """List tools from all connected servers."""
+        tools: list[McpTool] = []
+        for conn in self._servers.values():
+            tools.extend(conn.tools)
+        return tools
+
+
+@dataclass
+class _McpConnection:
+    """Internal connection state for an MCP server."""
+
+    name: str
+    config: McpServerConfig
+    process: asyncio.subprocess.Process
+    transport: _JsonRpcTransport
+    tools: list[McpTool] = field(default_factory=list)

monkeyscode/orchestrator.py

@@ -0,0 +1,170 @@
+"""
+MonkeysCode SDK — Multi-Agent Orchestrator.
+
+Run multiple agents in parallel or fan out across targets.
+
+Usage::
+
+    from monkeyscode.orchestrator import run_parallel, run_fan_out
+
+    result = await run_parallel(agent, [
+        "Fix auth.py",
+        "Fix db.py",
+        "Fix api.py",
+    ], concurrency=3)
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+from monkeyscode.events import AgentResult, TokenUsage
+
+# ── Types ────────────────────────────────────────────────────────────
+
+
+@dataclass
+class AgentSlotResult:
+    """Result from one agent slot in a parallel run."""
+
+    index: int
+    prompt: str
+    result: AgentResult | None = None
+    error: str | None = None
+    duration_ms: float = 0
+
+
+@dataclass
+class OrchestratorResult:
+    """Combined result of a parallel or fan-out run."""
+
+    results: list[AgentSlotResult] = field(default_factory=list)
+    total_cost: float = 0.0
+    total_tokens: TokenUsage = field(default_factory=TokenUsage)
+    duration_ms: float = 0
+    success_count: int = 0
+    fail_count: int = 0
+
+
+# ── Parallel ─────────────────────────────────────────────────────────
+
+
+async def run_parallel(
+    parent: Any,
+    prompts: list[str],
+    *,
+    concurrency: int = 0,
+) -> OrchestratorResult:
+    """
+    Run multiple prompts concurrently using separate agent instances.
+
+    Args:
+        parent: The MonkeysCode agent to clone settings from.
+        prompts: List of prompts to run in parallel.
+        concurrency: Max simultaneous agents (default: len(prompts), max: 10).
+    """
+    if not prompts:
+        raise ValueError("At least one prompt is required")
+
+    if concurrency <= 0:
+        concurrency = len(prompts)
+    concurrency = min(concurrency, 10)
+
+    start = time.monotonic()
+    semaphore = asyncio.Semaphore(concurrency)
+    results: list[AgentSlotResult] = [
+        AgentSlotResult(index=i, prompt=p) for i, p in enumerate(prompts)
+    ]
+
+    async def _run_slot(idx: int, prompt: str) -> None:
+        async with semaphore:
+            slot_start = time.monotonic()
+            try:
+                from monkeyscode.agent import MonkeysCode
+                child = MonkeysCode(
+                    model=parent._model,
+                    api_key=parent._api_key,
+                    working_directory=parent._working_directory,
+                    proxy_url=parent._proxy_url,
+                    max_turns=parent._max_turns,
+                    timeout=parent._timeout,
+                    permissions_mode=parent._permissions_mode,
+                    debug=parent._debug,
+                )
+                result = await child.run(prompt)
+                elapsed = (time.monotonic() - slot_start) * 1000
+                results[idx] = AgentSlotResult(
+                    index=idx,
+                    prompt=prompt,
+                    result=result,
+                    duration_ms=elapsed,
+                )
+                child.close()
+            except Exception as e:
+                elapsed = (time.monotonic() - slot_start) * 1000
+                results[idx] = AgentSlotResult(
+                    index=idx,
+                    prompt=prompt,
+                    error=str(e),
+                    duration_ms=elapsed,
+                )
+
+    tasks = [_run_slot(i, p) for i, p in enumerate(prompts)]
+    await asyncio.gather(*tasks)
+
+    total_cost = 0.0
+    total_input = 0
+    total_output = 0
+    success_count = 0
+    fail_count = 0
+
+    for slot in results:
+        if slot.result is not None:
+            total_cost += slot.result.cost
+            total_input += slot.result.tokens.input
+            total_output += slot.result.tokens.output
+            if slot.result.success:
+                success_count += 1
+            else:
+                fail_count += 1
+        else:
+            fail_count += 1
+
+    return OrchestratorResult(
+        results=results,
+        total_cost=total_cost,
+        total_tokens=TokenUsage(
+            input=total_input,
+            output=total_output,
+            total=total_input + total_output,
+        ),
+        duration_ms=(time.monotonic() - start) * 1000,
+        success_count=success_count,
+        fail_count=fail_count,
+    )
+
+
+async def run_fan_out(
+    parent: Any,
+    prompt: str,
+    targets: list[str],
+    *,
+    concurrency: int = 0,
+) -> OrchestratorResult:
+    """
+    Run the same prompt against multiple targets concurrently.
+
+    Args:
+        parent: The MonkeysCode agent to clone settings from.
+        prompt: The base instruction.
+        targets: Items to fan out over (file paths, modules, etc.).
+        concurrency: Max simultaneous agents.
+    """
+    if not targets:
+        raise ValueError("At least one target is required")
+
+    prompts = [f"{prompt}\n\nTarget: {target}" for target in targets]
+    return await run_parallel(parent, prompts, concurrency=concurrency)

monkeyscode/otel.py

@@ -0,0 +1,181 @@
+"""
+MonkeysCode SDK — OpenTelemetry Exporter.
+
+Send OTLP HTTP traces to Grafana, Datadog, Honeycomb, etc.
+
+Usage::
+
+    from monkeyscode.otel import OtelExporter
+
+    exporter = OtelExporter(endpoint="http://localhost:4318")
+    # Events are recorded automatically when wired into the agent
+    await exporter.flush()
+"""
+
+from __future__ import annotations
+
+import random
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+import httpx
+
+# ── Types ────────────────────────────────────────────────────────────
+
+
+@dataclass
+class OtelOptions:
+    """OpenTelemetry exporter configuration."""
+
+    endpoint: str = ""
+    headers: dict[str, str] = field(default_factory=dict)
+    service_name: str = "monkeyscode-sdk-python"
+    debug: bool = False
+
+
+@dataclass
+class _OtelSpan:
+    trace_id: str
+    span_id: str
+    parent_span_id: str = ""
+    name: str = ""
+    kind: int = 0
+    start_time_ns: int = 0
+    end_time_ns: int = 0
+    status_code: int = 0
+    attributes: list[dict[str, Any]] = field(default_factory=list)
+
+
+# ── Exporter ─────────────────────────────────────────────────────────
+
+
+class OtelExporter:
+    """
+    OTLP HTTP exporter for agent traces.
+
+    Records tool calls, run durations, and cost as spans.
+    Flush to send to your observability backend.
+    """
+
+    def __init__(self, options: OtelOptions | None = None) -> None:
+        opts = options or OtelOptions()
+        endpoint = opts.endpoint.rstrip("/")
+        if endpoint and not endpoint.endswith("/v1/traces"):
+            endpoint += "/v1/traces"
+
+        self._endpoint = endpoint
+        self._headers = opts.headers
+        self._service_name = opts.service_name
+        self._debug = opts.debug
+        self._trace_id = _generate_hex(32)
+        self._spans: list[_OtelSpan] = []
+        self._pending: dict[str, _OtelSpan] = {}
+
+    @property
+    def enabled(self) -> bool:
+        return bool(self._endpoint)
+
+    @property
+    def trace_id(self) -> str:
+        return self._trace_id
+
+    @property
+    def span_count(self) -> int:
+        return len(self._spans)
+
+    def record_tool_start(self, tool: str) -> None:
+        """Record the start of a tool call."""
+        if not self.enabled:
+            return
+        span = _OtelSpan(
+            trace_id=self._trace_id,
+            span_id=_generate_hex(16),
+            name=f"tool.{tool}",
+            kind=3,  # CLIENT
+            start_time_ns=time.time_ns(),
+            attributes=[
+                {"key": "tool.name", "value": {"stringValue": tool}},
+            ],
+        )
+        self._pending[tool] = span
+
+    def record_tool_end(self, tool: str, *, is_error: bool = False) -> None:
+        """Record the end of a tool call."""
+        span = self._pending.pop(tool, None)
+        if span is None:
+            return
+        span.end_time_ns = time.time_ns()
+        if is_error:
+            span.status_code = 2  # ERROR
+        self._spans.append(span)
+
+    def record_run_complete(
+        self,
+        *,
+        cost: float = 0.0,
+        tokens_total: int = 0,
+        duration_ms: float = 0,
+    ) -> None:
+        """Record the overall run as a root span."""
+        if not self.enabled:
+            return
+        now = time.time_ns()
+        self._spans.append(_OtelSpan(
+            trace_id=self._trace_id,
+            span_id=_generate_hex(16),
+            name="agent.run",
+            kind=2,  # SERVER
+            start_time_ns=now - int(duration_ms * 1e6),
+            end_time_ns=now,
+            attributes=[
+                {"key": "agent.cost", "value": {"stringValue": f"{cost:.4f}"}},
+                {"key": "agent.tokens.total", "value": {"intValue": str(tokens_total)}},
+            ],
+        ))
+
+    async def flush(self) -> None:
+        """Send all accumulated spans to the OTLP endpoint."""
+        if not self.enabled or not self._spans:
+            return
+
+        spans = self._spans[:]
+        self._spans.clear()
+
+        payload = {
+            "resourceSpans": [{
+                "resource": {
+                    "attributes": [
+                        {"key": "service.name", "value": {"stringValue": self._service_name}},
+                    ],
+                },
+                "scopeSpans": [{
+                    "scope": {"name": "monkeyscode-sdk-python", "version": "1.0.0"},
+                    "spans": [
+                        {
+                            "traceId": s.trace_id,
+                            "spanId": s.span_id,
+                            "parentSpanId": s.parent_span_id,
+                            "name": s.name,
+                            "kind": s.kind,
+                            "startTimeUnixNano": str(s.start_time_ns),
+                            "endTimeUnixNano": str(s.end_time_ns),
+                            "status": {"code": s.status_code},
+                            "attributes": s.attributes,
+                        }
+                        for s in spans
+                    ],
+                }],
+            }],
+        }
+
+        try:
+            async with httpx.AsyncClient() as client:
+                headers = {"Content-Type": "application/json", **self._headers}
+                await client.post(self._endpoint, json=payload, headers=headers, timeout=10.0)
+        except Exception:
+            pass  # OTEL should never break the app
+
+
+def _generate_hex(length: int) -> str:
+    return "".join(f"{random.randint(0, 255):02x}" for _ in range(length // 2))

monkeyscode/py.typed

@@ -0,0 +1 @@
+