codebuddy-agent-sdk 0.1.23__py3-none-musllinux_1_1_x86_64.whl → 0.2.2__py3-none-musllinux_1_1_x86_64.whl

codebuddy_agent_sdk/__init__.py

@@ -5,10 +5,22 @@ from ._errors import (
     CLIJSONDecodeError,
     CLINotFoundError,
     CodeBuddySDKError,
+    ExecutionError,
     ProcessError,
 )
 from ._version import __version__
 from .client import CodeBuddySDKClient
+from .mcp import (
+    CallToolResult,
+    SdkControlServerTransport,
+    SdkMcpServerOptions,
+    SdkMcpServerResult,
+    SdkMcpToolDefinition,
+    TextContent,
+    ToolHandler,
+    create_sdk_mcp_server,
+    tool,
+)
 from .query import query
 from .transport import Transport
 from .types import (
@@ -20,14 +32,20 @@ from .types import (
     AssistantMessage,
     CanUseTool,
     CanUseToolOptions,
+    Checkpoint,
+    CheckpointFileChangeStats,
     CodeBuddyAgentOptions,
     ContentBlock,
+    ErrorMessage,
+    FileVersion,
     HookCallback,
     HookContext,
     HookEvent,
     HookJSONOutput,
     HookMatcher,
+    McpSdkServerConfig,
     McpServerConfig,
+    McpStdioServerConfig,
     Message,
     PermissionMode,
     PermissionResult,
@@ -50,6 +68,16 @@ __all__ = [
     "CodeBuddySDKClient",
     "Transport",
     "__version__",
+    # MCP Server API
+    "create_sdk_mcp_server",
+    "tool",
+    "SdkControlServerTransport",
+    "SdkMcpServerOptions",
+    "SdkMcpServerResult",
+    "SdkMcpToolDefinition",
+    "ToolHandler",
+    "CallToolResult",
+    "TextContent",
     # Types - Permission
     "PermissionMode",
     # Types - Messages
@@ -59,6 +87,7 @@ __all__ = [
     "SystemMessage",
     "ResultMessage",
     "StreamEvent",
+    "ErrorMessage",
     # Types - Content blocks
     "ContentBlock",
     "TextBlock",
@@ -86,12 +115,19 @@ __all__ = [
     "HookMatcher",
     "HookJSONOutput",
     "HookContext",
+    # Types - Checkpoint
+    "Checkpoint",
+    "CheckpointFileChangeStats",
+    "FileVersion",
     # Types - MCP
     "McpServerConfig",
+    "McpStdioServerConfig",
+    "McpSdkServerConfig",
     # Errors
     "CodeBuddySDKError",
     "CLIConnectionError",
     "CLINotFoundError",
     "CLIJSONDecodeError",
     "ProcessError",
+    "ExecutionError",
 ]

codebuddy_agent_sdk/_binary.py

@@ -5,6 +5,8 @@ This module provides functions to locate the CodeBuddy CLI binary.
 The binary is bundled in platform-specific wheels.
 """
 
+from __future__ import annotations
+
 import os
 import platform
 import sys

codebuddy_agent_sdk/_errors.py

@@ -1,5 +1,7 @@
 """Error definitions for CodeBuddy Agent SDK."""
 
+from __future__ import annotations
+
 
 class CodeBuddySDKError(Exception):
     """Base exception for CodeBuddy SDK errors."""
@@ -37,3 +39,16 @@ class ProcessError(CodeBuddySDKError):
     """Raised when CLI process encounters an error."""
 
     pass
+
+
+class ExecutionError(CodeBuddySDKError):
+    """Raised when execution fails (e.g., authentication error, API error).
+
+    Contains the errors array from the ResultMessage.
+    """
+
+    def __init__(self, errors: list[str], subtype: str):
+        message = errors[0] if errors else "Execution failed"
+        super().__init__(message)
+        self.errors = errors
+        self.subtype = subtype

codebuddy_agent_sdk/_message_parser.py

@@ -1,10 +1,13 @@
 """Message parser for CLI output."""
 
+from __future__ import annotations
+
 from typing import Any
 
 from .types import (
     AssistantMessage,
     ContentBlock,
+    ErrorMessage,
     Message,
     ResultMessage,
     StreamEvent,
@@ -99,6 +102,7 @@ def parse_message(data: dict[str, Any]) -> Message | None:
             total_cost_usd=data.get("total_cost_usd"),
             usage=data.get("usage"),
             result=data.get("result"),
+            errors=data.get("errors"),
         )
 
     if msg_type == "stream_event":
@@ -109,4 +113,10 @@ def parse_message(data: dict[str, Any]) -> Message | None:
             parent_tool_use_id=data.get("parent_tool_use_id"),
         )
 
+    if msg_type == "error":
+        return ErrorMessage(
+            error=data.get("error", ""),
+            session_id=data.get("session_id"),
+        )
+
     return None

codebuddy_agent_sdk/_version.py

@@ -1,3 +1,3 @@
 """Version information for CodeBuddy Agent SDK."""
 
-__version__ = "0.1.23"
+__version__ = "0.2.2"

codebuddy_agent_sdk/client.py

@@ -1,14 +1,17 @@
 """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
+from ._errors import CLIConnectionError, ExecutionError
 from ._message_parser import parse_message
 from .transport import SubprocessTransport, Transport
-from .types import CanUseToolOptions, CodeBuddyAgentOptions, Message, ResultMessage
+from .types import CanUseToolOptions, CodeBuddyAgentOptions, ErrorMessage, HookMatcher, Message, ResultMessage
 
 
 class CodeBuddySDKClient:
@@ -44,12 +47,12 @@ class CodeBuddySDKClient:
         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:
+    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
@@ -68,13 +71,52 @@ class CodeBuddySDKClient:
         if not self._transport:
             return
 
+        hooks_config, self._hook_callbacks = self._build_hooks_config()
+
         request = {
             "type": "control_request",
             "request_id": f"init_{id(self)}",
-            "request": {"subtype": "initialize"},
+            "request": {
+                "subtype": "initialize",
+                "hooks": hooks_config,
+            },
         }
         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]],
@@ -139,20 +181,24 @@ class CodeBuddySDKClient:
         if subtype == "can_use_tool":
             await self._handle_permission_request(request_id, request)
         elif subtype == "hook_callback":
-            # Default: continue
+            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": {"continue": True},
+                    "response": hook_response,
                 },
             }
             await self._transport.write(json.dumps(response))
 
-    async def _handle_permission_request(
-        self, request_id: str, request: dict[str, Any]
-    ) -> None:
+    async def _handle_permission_request(self, request_id: str, request: dict[str, Any]) -> None:
         """Handle permission request from CLI."""
         if not self._transport:
             return
@@ -232,16 +278,42 @@ class CodeBuddySDKClient:
             }
             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.
+        Receive messages until and including a ResultMessage or ErrorMessage.
 
         Yields each message as it's received and terminates after
-        yielding a ResultMessage.
+        yielding a ResultMessage or ErrorMessage.
+        Raises ExecutionError if ResultMessage indicates an error.
         """
         async for message in self.receive_messages():
-            yield 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
+                return
+
+            yield message
+
+            if isinstance(message, ErrorMessage):
                 return
 
     async def interrupt(self) -> None:
@@ -287,12 +359,17 @@ class CodeBuddySDKClient:
             self._transport = None
         self._connected = False
 
-    async def __aenter__(self) -> "CodeBuddySDKClient":
+    async def __aenter__(self) -> CodeBuddySDKClient:
         """Enter async context - automatically connects."""
         await self.connect()
         return self
 
-    async def __aexit__(self, *args: Any) -> bool:
+    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)