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

codebuddy_agent_sdk/client.py

@@ -0,0 +1,394 @@
+"""CodeBuddy SDK Client for interactive conversations."""
+
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import AsyncIterable, AsyncIterator
+from types import TracebackType
+from typing import Any
+
+from ._errors import CLIConnectionError, ExecutionError
+from ._message_parser import parse_message
+from .transport import SubprocessTransport, Transport
+from .types import (
+    CanUseToolOptions,
+    CodeBuddyAgentOptions,
+    ErrorMessage,
+    Message,
+    ResultMessage,
+)
+
+
+class CodeBuddySDKClient:
+    """
+    Client for bidirectional, interactive conversations with CodeBuddy.
+
+    This client provides full control over the conversation flow with support
+    for streaming, interrupts, and dynamic message sending. For simple one-shot
+    queries, consider using the query() function instead.
+
+    Key features:
+    - Bidirectional: Send and receive messages at any time
+    - Stateful: Maintains conversation context across messages
+    - Interactive: Send follow-ups based on responses
+    - Control flow: Support for interrupts and session management
+
+    Example:
+        ```python
+        async with CodeBuddySDKClient() as client:
+            await client.query("Hello!")
+            async for msg in client.receive_response():
+                print(msg)
+        ```
+    """
+
+    def __init__(
+        self,
+        options: CodeBuddyAgentOptions | None = None,
+        transport: Transport | None = None,
+    ):
+        """Initialize CodeBuddy SDK client."""
+        self.options = options or CodeBuddyAgentOptions()
+        self._custom_transport = transport
+        self._transport: Transport | None = None
+        self._connected = False
+        # Hook callback registry: callback_id -> hook function
+        self._hook_callbacks: dict[str, Any] = {}
+
+        os.environ["CODEBUDDY_CODE_ENTRYPOINT"] = "sdk-py-client"
+
+    async def connect(self, prompt: str | AsyncIterable[dict[str, Any]] | None = None) -> None:
+        """Connect to CodeBuddy with an optional initial prompt."""
+        if self._custom_transport:
+            self._transport = self._custom_transport
+        else:
+            self._transport = SubprocessTransport(
+                options=self.options,
+                prompt=prompt,
+            )
+
+        await self._transport.connect()
+        self._connected = True
+        await self._send_initialize()
+
+    async def _send_initialize(self) -> None:
+        """Send initialization control request."""
+        if not self._transport:
+            return
+
+        hooks_config, self._hook_callbacks = self._build_hooks_config()
+
+        # Get SDK MCP server names from transport
+        sdk_mcp_server_names = self._transport.sdk_mcp_server_names
+
+        request = {
+            "type": "control_request",
+            "request_id": f"init_{id(self)}",
+            "request": {
+                "subtype": "initialize",
+                "hooks": hooks_config,
+                # Include SDK MCP server names from transport
+                "sdkMcpServers": sdk_mcp_server_names if sdk_mcp_server_names else None,
+            },
+        }
+        await self._transport.write(json.dumps(request))
+
+    def _build_hooks_config(self) -> tuple[dict[str, list[dict[str, Any]]] | None, dict[str, Any]]:
+        """Build hooks configuration for CLI and callback registry.
+
+        Returns:
+            Tuple of (config for CLI, callback_id -> hook function mapping)
+        """
+        callbacks: dict[str, Any] = {}
+
+        if not self.options.hooks:
+            return None, callbacks
+
+        config: dict[str, list[dict[str, Any]]] = {}
+
+        for event, matchers in self.options.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
+
+    async def query(
+        self,
+        prompt: str | AsyncIterable[dict[str, Any]],
+        session_id: str = "default",
+    ) -> None:
+        """
+        Send a new request.
+
+        Args:
+            prompt: Either a string message or an async iterable of message dicts
+            session_id: Session identifier for the conversation
+        """
+        if not self._connected or not self._transport:
+            raise CLIConnectionError("Not connected. Call connect() first.")
+
+        if isinstance(prompt, str):
+            message = {
+                "type": "user",
+                "message": {"role": "user", "content": prompt},
+                "parent_tool_use_id": None,
+                "session_id": session_id,
+            }
+            await self._transport.write(json.dumps(message))
+        else:
+            async for msg in prompt:
+                if "session_id" not in msg:
+                    msg["session_id"] = session_id
+                await self._transport.write(json.dumps(msg))
+
+    async def receive_messages(self) -> AsyncIterator[Message]:
+        """Receive all messages from CodeBuddy."""
+        if not self._transport:
+            raise CLIConnectionError("Not connected.")
+
+        async for line in self._transport.read():
+            if not line:
+                continue
+
+            try:
+                data = json.loads(line)
+
+                # Handle control requests (permissions, hooks)
+                if data.get("type") == "control_request":
+                    await self._handle_control_request(data)
+                    continue
+
+                message = parse_message(data)
+                if message:
+                    yield message
+            except json.JSONDecodeError:
+                continue
+
+    async def _handle_control_request(self, data: dict[str, Any]) -> None:
+        """Handle control request from CLI."""
+        if not self._transport:
+            return
+
+        request_id = data.get("request_id", "")
+        request = data.get("request", {})
+        subtype = request.get("subtype", "")
+
+        if subtype == "can_use_tool":
+            await self._handle_permission_request(request_id, request)
+        elif subtype == "hook_callback":
+            callback_id = request.get("callback_id", "")
+            hook_input = request.get("input", {})
+            tool_use_id = request.get("tool_use_id")
+
+            # Execute the hook
+            hook_response = await self._execute_hook(callback_id, hook_input, tool_use_id)
+
+            response = {
+                "type": "control_response",
+                "response": {
+                    "subtype": "success",
+                    "request_id": request_id,
+                    "response": hook_response,
+                },
+            }
+            await self._transport.write(json.dumps(response))
+        elif subtype == "mcp_message":
+            # MCP messages are handled at the transport level
+            from .transport import SubprocessTransport
+
+            if isinstance(self._transport, SubprocessTransport):
+                await self._transport.handle_mcp_message_request(request_id, request)
+
+    async def _handle_permission_request(self, request_id: str, request: dict[str, Any]) -> None:
+        """Handle permission request from CLI."""
+        if not self._transport:
+            return
+
+        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 = self.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 self._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 self._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 self._transport.write(json.dumps(response))
+
+    async def _execute_hook(
+        self,
+        callback_id: str,
+        hook_input: dict[str, Any],
+        tool_use_id: str | None,
+    ) -> dict[str, Any]:
+        """Execute a hook callback by looking up in the callback registry."""
+        hook = self._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)}
+
+    async def receive_response(self) -> AsyncIterator[Message]:
+        """
+        Receive messages until and including a ResultMessage or ErrorMessage.
+
+        Yields each message as it's received and terminates after
+        yielding a ResultMessage or ErrorMessage.
+        Raises ExecutionError if ResultMessage indicates an error.
+        """
+        async for message in self.receive_messages():
+            # 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
+                return
+
+            yield message
+
+            if isinstance(message, ErrorMessage):
+                return
+
+    async def interrupt(self) -> None:
+        """Send interrupt signal."""
+        if not self._transport:
+            raise CLIConnectionError("Not connected.")
+
+        request = {
+            "type": "control_request",
+            "request_id": f"interrupt_{id(self)}",
+            "request": {"subtype": "interrupt"},
+        }
+        await self._transport.write(json.dumps(request))
+
+    async def set_permission_mode(self, mode: str) -> None:
+        """Change permission mode during conversation."""
+        if not self._transport:
+            raise CLIConnectionError("Not connected.")
+
+        request = {
+            "type": "control_request",
+            "request_id": f"perm_{id(self)}",
+            "request": {"subtype": "set_permission_mode", "mode": mode},
+        }
+        await self._transport.write(json.dumps(request))
+
+    async def set_model(self, model: str | None = None) -> None:
+        """Change the AI model during conversation."""
+        if not self._transport:
+            raise CLIConnectionError("Not connected.")
+
+        request = {
+            "type": "control_request",
+            "request_id": f"model_{id(self)}",
+            "request": {"subtype": "set_model", "model": model},
+        }
+        await self._transport.write(json.dumps(request))
+
+    async def disconnect(self) -> None:
+        """Disconnect from CodeBuddy."""
+        if self._transport:
+            await self._transport.close()
+            self._transport = None
+        self._connected = False
+
+    async def __aenter__(self) -> CodeBuddySDKClient:
+        """Enter async context - automatically connects."""
+        await self.connect()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool:
+        """Exit async context - always disconnects."""
+        await self.disconnect()
+        return False

codebuddy_agent_sdk/mcp/__init__.py

@@ -0,0 +1,35 @@
+"""
+MCP (Model Context Protocol) Integration
+
+This module provides utilities for creating and managing SDK MCP servers
+that can be integrated with the CLI via the control protocol.
+"""
+
+from .create_sdk_mcp_server import (
+    create_sdk_mcp_server,
+    tool,
+)
+from .sdk_control_server_transport import SdkControlServerTransport
+from .types import (
+    CallToolResult,
+    SdkMcpServerOptions,
+    SdkMcpServerResult,
+    SdkMcpToolDefinition,
+    TextContent,
+    ToolHandler,
+)
+
+__all__ = [
+    # Factory functions
+    "create_sdk_mcp_server",
+    "tool",
+    # Transport
+    "SdkControlServerTransport",
+    # Types
+    "SdkMcpServerOptions",
+    "SdkMcpServerResult",
+    "SdkMcpToolDefinition",
+    "ToolHandler",
+    "CallToolResult",
+    "TextContent",
+]

codebuddy_agent_sdk/mcp/create_sdk_mcp_server.py

@@ -0,0 +1,154 @@
+"""
+create_sdk_mcp_server - Create an SDK MCP Server
+
+This function creates an MCP server that can be integrated into the SDK
+and used with the CLI via the control protocol.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import Any, TypeVar
+
+from .types import (
+    CallToolResult,
+    SdkMcpServer,
+    SdkMcpServerOptions,
+    SdkMcpServerResult,
+    SdkMcpToolDefinition,
+    ToolInputSchema,
+)
+
+# Tool handler type
+ToolHandler = Callable[[dict[str, Any]], CallToolResult | Awaitable[CallToolResult]]
+
+# Type variable for the decorated function
+F = TypeVar("F", bound=ToolHandler)
+
+
+def _python_type_to_json_schema(py_type: type) -> dict[str, Any]:
+    """Convert Python type to JSON Schema type."""
+    type_mapping = {
+        str: {"type": "string"},
+        int: {"type": "integer"},
+        float: {"type": "number"},
+        bool: {"type": "boolean"},
+        list: {"type": "array"},
+        dict: {"type": "object"},
+    }
+    return type_mapping.get(py_type, {"type": "string"})
+
+
+def _convert_schema(input_schema: dict[str, Any]) -> tuple[dict[str, dict[str, Any]], list[str]]:
+    """
+    Convert input schema to JSON Schema format.
+
+    Supports multiple formats:
+    1. Simple type mapping: {"latitude": float, "longitude": float}
+    2. JSON Schema format: {"properties": {...}, "required": [...]}
+
+    Returns:
+        Tuple of (properties dict, required list)
+    """
+    # Check if it's already in JSON Schema format
+    if "properties" in input_schema:
+        return input_schema.get("properties", {}), input_schema.get("required", [])
+
+    if "type" in input_schema and input_schema.get("type") == "object":
+        return input_schema.get("properties", {}), input_schema.get("required", [])
+
+    # Simple type mapping format: {"param_name": type}
+    properties = {}
+    required = []
+
+    for param_name, param_type in input_schema.items():
+        if isinstance(param_type, type):
+            properties[param_name] = _python_type_to_json_schema(param_type)
+            required.append(param_name)
+        elif isinstance(param_type, dict):
+            properties[param_name] = param_type
+            if "default" not in param_type:
+                required.append(param_name)
+        else:
+            properties[param_name] = {"type": "string"}
+            required.append(param_name)
+
+    return properties, required
+
+
+def tool(name: str, description: str, input_schema: dict[str, Any]) -> Callable[[F], F]:
+    """
+    Decorator to define an MCP tool.
+
+    Example:
+        ```python
+        @tool("get_weather", "Get current weather", {"latitude": float, "longitude": float})
+        async def get_weather(args: dict[str, Any]) -> dict[str, Any]:
+            return {"content": [{"type": "text", "text": f"Weather: sunny"}]}
+        ```
+
+    Args:
+        name: Tool name (unique within the server)
+        description: Tool description
+        input_schema: Input parameters schema. Supports:
+            - Simple types: {"latitude": float, "longitude": float}
+            - JSON Schema: {"properties": {...}, "required": [...]}
+
+    Returns:
+        Decorated function with tool metadata attached
+    """
+    properties, required = _convert_schema(input_schema)
+    tool_schema = ToolInputSchema(type="object", properties=properties, required=required)
+
+    def decorator(func: F) -> F:
+        func._tool_definition = SdkMcpToolDefinition(  # type: ignore[attr-defined]
+            name=name,
+            description=description,
+            input_schema=tool_schema,
+            handler=func,
+        )
+        return func
+
+    return decorator
+
+
+def create_sdk_mcp_server(
+    name: str,
+    version: str = "1.0.0",
+    tools: list[Callable[..., Any]] | None = None,
+) -> SdkMcpServerResult:
+    """
+    Create an SDK MCP Server.
+
+    Args:
+        name: Server name (unique within the session)
+        version: Server version (defaults to "1.0.0")
+        tools: List of functions decorated with @tool
+
+    Returns:
+        SDK MCP server result for use with query()
+
+    Example:
+        ```python
+        from codebuddy_agent_sdk import create_sdk_mcp_server, tool, query
+
+        @tool("get_weather", "Get weather", {"location": str})
+        async def get_weather(args: dict) -> dict:
+            return {"content": [{"type": "text", "text": f"Weather: sunny"}]}
+
+        server = create_sdk_mcp_server(
+            name="my-server",
+            tools=[get_weather]
+        )
+        ```
+    """
+    tool_definitions: list[SdkMcpToolDefinition] = []
+    if tools:
+        for func in tools:
+            if not hasattr(func, "_tool_definition"):
+                raise ValueError(f"Function {func.__name__} is not decorated with @tool")
+            tool_definitions.append(func._tool_definition)
+
+    server = SdkMcpServer(SdkMcpServerOptions(name=name, version=version, tools=tool_definitions))
+
+    return SdkMcpServerResult(type="sdk", name=name, server=server)

codebuddy_agent_sdk/mcp/sdk_control_server_transport.py

@@ -0,0 +1,95 @@
+"""
+SDK Control Server Transport
+
+Custom transport implementation that bridges SDK MCP servers to CLI process.
+This transport forwards MCP messages through the control protocol.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from .types import JSONRPCMessage
+
+# Callback function type for sending MCP messages to CLI
+SendMcpMessageCallback = Callable[[JSONRPCMessage], None]
+
+
+class SdkControlServerTransport:
+    """
+    SdkControlServerTransport - bridges MCP servers to CLI via control messages.
+
+    This transport implements a simple interface for forwarding MCP messages
+    between the SDK MCP server and the CLI via the control protocol.
+    """
+
+    def __init__(self, send_mcp_message: SendMcpMessageCallback):
+        """
+        Create a new SDK Control Server Transport.
+
+        Args:
+            send_mcp_message: Callback function to forward MCP messages to CLI
+        """
+        self._send_mcp_message = send_mcp_message
+        self._is_closed = False
+        self._on_message: Callable[[JSONRPCMessage], None] | None = None
+        self._on_close: Callable[[], None] | None = None
+        self._on_error: Callable[[Exception], None] | None = None
+
+    @property
+    def closed(self) -> bool:
+        """Check if the transport is closed."""
+        return self._is_closed
+
+    def set_on_message(self, callback: Callable[[JSONRPCMessage], None] | None) -> None:
+        """Set the message callback."""
+        self._on_message = callback
+
+    def set_on_close(self, callback: Callable[[], None] | None) -> None:
+        """Set the close callback."""
+        self._on_close = callback
+
+    def set_on_error(self, callback: Callable[[Exception], None] | None) -> None:
+        """Set the error callback."""
+        self._on_error = callback
+
+    async def start(self) -> None:
+        """
+        Start the transport.
+        No-op since connection is already established via stdio.
+        """
+        pass
+
+    async def send(self, message: JSONRPCMessage) -> None:
+        """
+        Send a message to the CLI via control_request.
+
+        Args:
+            message: The JSON-RPC message to send
+        """
+        if self._is_closed:
+            raise RuntimeError("Transport is closed")
+        # Forward message to CLI via control_request
+        self._send_mcp_message(message)
+
+    async def close(self) -> None:
+        """Close the transport."""
+        if self._is_closed:
+            return
+
+        self._is_closed = True
+        if self._on_close:
+            self._on_close()
+
+    def handle_incoming_message(self, message: JSONRPCMessage) -> None:
+        """
+        Handle incoming message from CLI.
+        This method should be called when the CLI sends a message to this server.
+
+        Args:
+            message: The JSON-RPC message from CLI
+        """
+        if self._is_closed:
+            return
+        if self._on_message:
+            self._on_message(message)