codebuddy-agent-sdk 0.3.7__py3-none-manylinux_2_17_x86_64.whl

codebuddy_agent_sdk/mcp/types.py

@@ -0,0 +1,300 @@
+"""Type definitions for SDK MCP Server."""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Literal, TypedDict, cast
+
+if TYPE_CHECKING:
+    from .sdk_control_server_transport import SdkControlServerTransport
+
+
+# ============= JSON-RPC Types =============
+
+
+class JSONRPCRequest(TypedDict, total=False):
+    """JSON-RPC 2.0 request."""
+
+    jsonrpc: Literal["2.0"]
+    id: str | int
+    method: str
+    params: dict[str, Any] | list[Any] | None
+
+
+class JSONRPCError(TypedDict, total=False):
+    """JSON-RPC 2.0 error."""
+
+    code: int
+    message: str
+    data: Any
+
+
+class JSONRPCResponse(TypedDict, total=False):
+    """JSON-RPC 2.0 response."""
+
+    jsonrpc: Literal["2.0"]
+    id: str | int | None
+    result: Any
+    error: JSONRPCError
+
+
+class JSONRPCNotification(TypedDict, total=False):
+    """JSON-RPC 2.0 notification."""
+
+    jsonrpc: Literal["2.0"]
+    method: str
+    params: dict[str, Any] | list[Any] | None
+
+
+JSONRPCMessage = JSONRPCRequest | JSONRPCResponse | JSONRPCNotification
+
+
+# ============= MCP Tool Types =============
+
+
+class TextContent(TypedDict, total=False):
+    """Text content in tool result."""
+
+    type: Literal["text"]
+    text: str
+
+
+class ImageContent(TypedDict, total=False):
+    """Image content in tool result."""
+
+    type: Literal["image"]
+    data: str
+    mimeType: str
+
+
+class EmbeddedResource(TypedDict, total=False):
+    """Embedded resource content in tool result."""
+
+    type: Literal["resource"]
+    resource: dict[str, Any]
+
+
+ToolResultContent = TextContent | ImageContent | EmbeddedResource
+
+
+class CallToolResult(TypedDict, total=False):
+    """Result from calling a tool."""
+
+    content: list[ToolResultContent]
+    isError: bool
+
+
+# Tool handler type - takes arguments dict and returns CallToolResult
+ToolHandler = Callable[[dict[str, Any]], CallToolResult | Awaitable[CallToolResult]]
+
+
+@dataclass
+class ToolInputProperty:
+    """Property definition for tool input schema."""
+
+    type: str
+    description: str | None = None
+    enum: list[str] | None = None
+    default: Any = None
+    minimum: float | None = None
+    maximum: float | None = None
+
+
+@dataclass
+class ToolInputSchema:
+    """JSON Schema for tool input."""
+
+    type: Literal["object"] = "object"
+    properties: dict[str, dict[str, Any]] = field(default_factory=dict)
+    required: list[str] = field(default_factory=list)
+
+
+@dataclass
+class SdkMcpToolDefinition:
+    """
+    Tool definition for SDK MCP Server.
+
+    Example:
+        ```python
+        tool_def = SdkMcpToolDefinition(
+            name="get_weather",
+            description="Get the current weather for a location",
+            input_schema=ToolInputSchema(
+                properties={
+                    "location": {"type": "string", "description": "The city name"},
+                    "units": {"type": "string", "enum": ["celsius", "fahrenheit"]},
+                },
+                required=["location"],
+            ),
+            handler=get_weather_handler,
+        )
+        ```
+    """
+
+    name: str
+    description: str
+    input_schema: ToolInputSchema
+    handler: ToolHandler
+
+
+@dataclass
+class SdkMcpServerOptions:
+    """
+    Options for creating an SDK MCP Server.
+
+    Attributes:
+        name: Server name (must be unique within the session)
+        version: Server version (defaults to "1.0.0")
+        tools: List of tool definitions to register
+    """
+
+    name: str
+    version: str = "1.0.0"
+    tools: list[SdkMcpToolDefinition] = field(default_factory=list)
+
+
+@dataclass
+class SdkMcpServerResult:
+    """
+    Result type for create_sdk_mcp_server.
+
+    Attributes:
+        type: Type discriminator - always "sdk" for SDK MCP servers
+        name: Server name
+        server: The MCP server instance
+    """
+
+    type: Literal["sdk"]
+    name: str
+    server: SdkMcpServer
+
+
+class SdkMcpServer:
+    """
+    SDK MCP Server implementation.
+
+    This class implements an MCP server that runs within the SDK process
+    and communicates with the CLI via the control protocol.
+    """
+
+    def __init__(self, options: SdkMcpServerOptions):
+        self.name = options.name
+        self.version = options.version
+        self.tools: dict[str, SdkMcpToolDefinition] = {}
+        self._transport: SdkControlServerTransport | None = None
+
+        # Register tools
+        for tool_def in options.tools:
+            self.tools[tool_def.name] = tool_def
+
+    def connect(self, transport: SdkControlServerTransport) -> None:
+        """Connect the server to a transport."""
+        self._transport = transport
+
+    async def handle_message(self, message: JSONRPCMessage) -> JSONRPCMessage | None:
+        """Handle an incoming JSON-RPC message."""
+        # Check if it's a request (has method and id)
+        if "method" not in message:
+            return None
+
+        method = cast(str, message.get("method", ""))
+        msg_id = cast("str | int | None", message.get("id"))
+        params = cast("dict[str, Any] | None", message.get("params", {}))
+
+        if method == "initialize":
+            return await self._handle_initialize(msg_id, params)
+        elif method == "tools/list":
+            return await self._handle_tools_list(msg_id)
+        elif method == "tools/call":
+            return await self._handle_tools_call(msg_id, params)
+        elif method == "notifications/initialized":
+            # Notification, no response needed
+            return None
+        else:
+            # Unknown method
+            if msg_id is not None:
+                return self._create_error_response(msg_id, -32601, f"Method not found: {method}")
+            return None
+
+    async def _handle_initialize(
+        self, msg_id: str | int | None, params: dict[str, Any] | None
+    ) -> JSONRPCMessage:
+        """Handle initialize request."""
+        result = {
+            "protocolVersion": "2024-11-05",
+            "capabilities": {
+                "tools": {},
+            },
+            "serverInfo": {
+                "name": self.name,
+                "version": self.version,
+            },
+        }
+        return self._create_response(msg_id, result)
+
+    async def _handle_tools_list(self, msg_id: str | int | None) -> JSONRPCMessage:
+        """Handle tools/list request."""
+        tools_list = []
+        for tool_def in self.tools.values():
+            tools_list.append(
+                {
+                    "name": tool_def.name,
+                    "description": tool_def.description,
+                    "inputSchema": {
+                        "type": tool_def.input_schema.type,
+                        "properties": tool_def.input_schema.properties,
+                        "required": tool_def.input_schema.required,
+                    },
+                }
+            )
+        return self._create_response(msg_id, {"tools": tools_list})
+
+    async def _handle_tools_call(
+        self, msg_id: str | int | None, params: dict[str, Any] | None
+    ) -> JSONRPCMessage:
+        """Handle tools/call request."""
+        if not isinstance(params, dict):
+            return self._create_error_response(msg_id, -32602, "Invalid params")
+
+        tool_name = params.get("name", "")
+        arguments = params.get("arguments", {})
+
+        tool_def = self.tools.get(tool_name)
+        if not tool_def:
+            return self._create_error_response(msg_id, -32602, f"Tool not found: {tool_name}")
+
+        try:
+            # Call the handler
+            result = tool_def.handler(arguments)
+            # Handle async handlers
+            if inspect.isawaitable(result):
+                result = await result
+
+            return self._create_response(msg_id, result)
+        except Exception as e:
+            # Return error as tool result
+            error_result: CallToolResult = {
+                "content": [{"type": "text", "text": str(e)}],
+                "isError": True,
+            }
+            return self._create_response(msg_id, error_result)
+
+    def _create_response(self, msg_id: str | int | None, result: Any) -> JSONRPCResponse:
+        """Create a JSON-RPC response."""
+        return {
+            "jsonrpc": "2.0",
+            "id": msg_id,
+            "result": result,
+        }
+
+    def _create_error_response(
+        self, msg_id: str | int | None, code: int, message: str
+    ) -> JSONRPCResponse:
+        """Create a JSON-RPC error response."""
+        return {
+            "jsonrpc": "2.0",
+            "id": msg_id,
+            "error": {"code": code, "message": message},
+        }

codebuddy_agent_sdk/query.py

@@ -0,0 +1,340 @@
+"""Query function for one-shot interactions with CodeBuddy."""
+
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import AsyncIterable, AsyncIterator
+from dataclasses import asdict
+from typing import Any
+
+from ._errors import ExecutionError
+from ._message_parser import parse_message
+from .transport import SubprocessTransport, Transport
+from .types import (
+    AppendSystemPrompt,
+    CanUseToolOptions,
+    CodeBuddyAgentOptions,
+    ErrorMessage,
+    HookCallback,
+    HookMatcher,
+    Message,
+    ResultMessage,
+)
+
+
+async def query(
+    *,
+    prompt: str | AsyncIterable[dict[str, Any]],
+    options: CodeBuddyAgentOptions | None = None,
+    transport: Transport | None = None,
+) -> AsyncIterator[Message]:
+    """
+    Query CodeBuddy for one-shot or unidirectional streaming interactions.
+
+    This function is ideal for simple, stateless queries where you don't need
+    bidirectional communication or conversation management. For interactive,
+    stateful conversations, use CodeBuddySDKClient instead.
+
+    Args:
+        prompt: The prompt to send to CodeBuddy. Can be a string for single-shot
+                queries or an AsyncIterable[dict] for streaming mode.
+        options: Optional configuration (defaults to CodeBuddyAgentOptions() if None).
+        transport: Optional transport implementation. If provided, this will be used
+                  instead of the default subprocess transport.
+
+    Yields:
+        Messages from the conversation.
+
+    Example:
+        ```python
+        async for message in query(prompt="What is 2+2?"):
+            print(message)
+        ```
+    """
+    if options is None:
+        options = CodeBuddyAgentOptions()
+
+    os.environ["CODEBUDDY_CODE_ENTRYPOINT"] = "sdk-py"
+
+    # Transport handles SDK MCP server extraction automatically
+    if transport is None:
+        transport = SubprocessTransport(options=options, prompt=prompt)
+
+    await transport.connect()
+
+    # Hook callback registry
+    hook_callbacks: dict[str, HookCallback] = {}
+
+    try:
+        hook_callbacks = await _send_initialize(transport, options)
+        await _send_prompt(transport, prompt)
+
+        async for line in transport.read():
+            if not line:
+                continue
+
+            try:
+                data = json.loads(line)
+
+                # Handle control requests (hooks, permissions, MCP messages)
+                if data.get("type") == "control_request":
+                    await _handle_control_request(transport, data, options, hook_callbacks)
+                    continue
+
+                message = parse_message(data)
+                if message:
+                    # Check for execution error BEFORE yielding
+                    if isinstance(message, ResultMessage):
+                        if message.is_error and message.errors and len(message.errors) > 0:
+                            raise ExecutionError(message.errors, message.subtype)
+                        yield message
+                        break
+
+                    yield message
+
+                    if isinstance(message, ErrorMessage):
+                        break
+
+            except json.JSONDecodeError:
+                continue  # Ignore non-JSON lines
+
+    finally:
+        await transport.close()
+
+
+async def _send_initialize(
+    transport: Transport,
+    options: CodeBuddyAgentOptions,
+) -> dict[str, HookCallback]:
+    """Send initialization control request.
+
+    Returns:
+        Hook callbacks registry (callback_id -> hook function)
+    """
+    hooks_config, hook_callbacks = _build_hooks_config(options.hooks)
+    agents_config = (
+        {name: asdict(agent) for name, agent in options.agents.items()} if options.agents else None
+    )
+
+    # Parse system_prompt config
+    system_prompt: str | None = None
+    append_system_prompt: str | None = None
+    if isinstance(options.system_prompt, str):
+        system_prompt = options.system_prompt
+    elif isinstance(options.system_prompt, AppendSystemPrompt):
+        append_system_prompt = options.system_prompt.append
+
+    # Get SDK MCP server names from transport
+    sdk_mcp_server_names = transport.sdk_mcp_server_names
+
+    request = {
+        "type": "control_request",
+        "request_id": f"init_{id(options)}",
+        "request": {
+            "subtype": "initialize",
+            "hooks": hooks_config,
+            "systemPrompt": system_prompt,
+            "appendSystemPrompt": append_system_prompt,
+            "agents": agents_config,
+            # Include SDK MCP server names from transport
+            "sdkMcpServers": sdk_mcp_server_names if sdk_mcp_server_names else None,
+        },
+    }
+    await transport.write(json.dumps(request))
+    return hook_callbacks
+
+
+async def _send_prompt(transport: Transport, prompt: str | AsyncIterable[dict[str, Any]]) -> None:
+    """Send user prompt."""
+    if isinstance(prompt, str):
+        message = {
+            "type": "user",
+            "session_id": "",
+            "message": {"role": "user", "content": prompt},
+            "parent_tool_use_id": None,
+        }
+        await transport.write(json.dumps(message))
+    else:
+        async for msg in prompt:
+            await transport.write(json.dumps(msg))
+
+
+async def _handle_control_request(
+    transport: Transport,
+    data: dict[str, Any],
+    options: CodeBuddyAgentOptions,
+    hook_callbacks: dict[str, HookCallback],
+) -> None:
+    """Handle control request from CLI."""
+    request_id = data.get("request_id", "")
+    request = data.get("request", {})
+    subtype = request.get("subtype", "")
+
+    if subtype == "hook_callback":
+        # Handle hook callback
+        callback_id = request.get("callback_id", "")
+        hook_input = request.get("input", {})
+        tool_use_id = request.get("tool_use_id")
+
+        # Find and execute the hook using callback registry
+        response = await _execute_hook(callback_id, hook_input, tool_use_id, hook_callbacks)
+
+        # Send response
+        control_response = {
+            "type": "control_response",
+            "response": {
+                "subtype": "success",
+                "request_id": request_id,
+                "response": response,
+            },
+        }
+        await transport.write(json.dumps(control_response))
+
+    elif subtype == "can_use_tool":
+        await _handle_permission_request(transport, request_id, request, options)
+
+    elif subtype == "mcp_message":
+        # MCP messages are handled at the transport level
+        if isinstance(transport, SubprocessTransport):
+            await transport.handle_mcp_message_request(request_id, request)
+
+
+async def _handle_permission_request(
+    transport: Transport,
+    request_id: str,
+    request: dict[str, Any],
+    options: CodeBuddyAgentOptions,
+) -> None:
+    """Handle permission request from CLI."""
+    tool_name = request.get("tool_name", "")
+    input_data = request.get("input", {})
+    tool_use_id = request.get("tool_use_id", "")
+    agent_id = request.get("agent_id")
+
+    can_use_tool = options.can_use_tool
+
+    # Default deny if no callback provided
+    if not can_use_tool:
+        response = {
+            "type": "control_response",
+            "response": {
+                "subtype": "success",
+                "request_id": request_id,
+                "response": {
+                    "allowed": False,
+                    "reason": "No permission handler provided",
+                    "tool_use_id": tool_use_id,
+                },
+            },
+        }
+        await transport.write(json.dumps(response))
+        return
+
+    try:
+        callback_options = CanUseToolOptions(
+            tool_use_id=tool_use_id,
+            signal=None,
+            agent_id=agent_id,
+            suggestions=request.get("permission_suggestions"),
+            blocked_path=request.get("blocked_path"),
+            decision_reason=request.get("decision_reason"),
+        )
+
+        result = await can_use_tool(tool_name, input_data, callback_options)
+
+        if result.behavior == "allow":
+            response_data = {
+                "allowed": True,
+                "updatedInput": result.updated_input,
+                "tool_use_id": tool_use_id,
+            }
+        else:
+            response_data = {
+                "allowed": False,
+                "reason": result.message,
+                "interrupt": result.interrupt,
+                "tool_use_id": tool_use_id,
+            }
+
+        response = {
+            "type": "control_response",
+            "response": {
+                "subtype": "success",
+                "request_id": request_id,
+                "response": response_data,
+            },
+        }
+        await transport.write(json.dumps(response))
+
+    except Exception as e:
+        response = {
+            "type": "control_response",
+            "response": {
+                "subtype": "success",
+                "request_id": request_id,
+                "response": {
+                    "allowed": False,
+                    "reason": str(e),
+                    "tool_use_id": tool_use_id,
+                },
+            },
+        }
+        await transport.write(json.dumps(response))
+
+
+async def _execute_hook(
+    callback_id: str,
+    hook_input: dict[str, Any],
+    tool_use_id: str | None,
+    hook_callbacks: dict[str, HookCallback],
+) -> dict[str, Any]:
+    """Execute a hook callback by looking up in the callback registry."""
+    hook = hook_callbacks.get(callback_id)
+    if not hook:
+        return {"continue": True}
+
+    try:
+        result = await hook(hook_input, tool_use_id, {"signal": None})
+        return dict(result)
+    except Exception as e:
+        return {"continue": False, "stopReason": str(e)}
+
+
+def _build_hooks_config(
+    hooks: dict[Any, list[HookMatcher]] | None,
+) -> tuple[dict[str, list[dict[str, Any]]] | None, dict[str, HookCallback]]:
+    """Build hooks configuration for CLI and callback registry.
+
+    Returns:
+        Tuple of (config for CLI, callback_id -> hook function mapping)
+    """
+    callbacks: dict[str, HookCallback] = {}
+
+    if not hooks:
+        return None, callbacks
+
+    config: dict[str, list[dict[str, Any]]] = {}
+
+    for event, matchers in hooks.items():
+        event_str = str(event)
+        matcher_configs = []
+
+        for i, m in enumerate(matchers):
+            callback_ids = []
+            for j, hook in enumerate(m.hooks):
+                callback_id = f"hook_{event_str}_{i}_{j}"
+                callback_ids.append(callback_id)
+                callbacks[callback_id] = hook
+
+            matcher_configs.append(
+                {
+                    "matcher": m.matcher,
+                    "hookCallbackIds": callback_ids,
+                    "timeout": m.timeout,
+                }
+            )
+
+        config[event_str] = matcher_configs
+
+    return (config if config else None), callbacks

codebuddy_agent_sdk/transport/__init__.py

@@ -0,0 +1,6 @@
+"""Transport layer for CLI communication."""
+
+from .base import Transport
+from .subprocess import SubprocessTransport
+
+__all__ = ["Transport", "SubprocessTransport"]

codebuddy_agent_sdk/transport/base.py

@@ -0,0 +1,31 @@
+"""Transport base class for CLI communication."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator
+
+
+class Transport(ABC):
+    """Abstract transport layer for CLI communication."""
+
+    @abstractmethod
+    async def connect(self) -> None:
+        """Establish connection to CLI."""
+
+    @abstractmethod
+    def read(self) -> AsyncIterator[str]:
+        """Read messages from CLI as an async iterator."""
+
+    @abstractmethod
+    async def write(self, data: str) -> None:
+        """Write data to CLI."""
+
+    @abstractmethod
+    async def close(self) -> None:
+        """Close the connection."""
+
+    @property
+    def sdk_mcp_server_names(self) -> list[str]:
+        """Get the list of SDK MCP server names. Override in subclasses."""
+        return []