codebuddy-agent-sdk 0.1.8__py3-none-manylinux_2_17_x86_64.whl → 0.2.0__py3-none-manylinux_2_17_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 (
@@ -22,12 +34,15 @@ from .types import (
     CanUseToolOptions,
     CodeBuddyAgentOptions,
     ContentBlock,
+    ErrorMessage,
     HookCallback,
     HookContext,
     HookEvent,
     HookJSONOutput,
     HookMatcher,
+    McpSdkServerConfig,
     McpServerConfig,
+    McpStdioServerConfig,
     Message,
     PermissionMode,
     PermissionResult,
@@ -50,6 +65,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 +84,7 @@ __all__ = [
     "SystemMessage",
     "ResultMessage",
     "StreamEvent",
+    "ErrorMessage",
     # Types - Content blocks
     "ContentBlock",
     "TextBlock",
@@ -88,10 +114,13 @@ __all__ = [
     "HookContext",
     # 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.8"
+__version__ = "0.2.0"

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, Message, ResultMessage
 
 
 class CodeBuddySDKClient:
@@ -47,9 +50,7 @@ class CodeBuddySDKClient:
 
         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
@@ -150,9 +151,7 @@ class CodeBuddySDKClient:
             }
             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
@@ -234,14 +233,23 @@ class CodeBuddySDKClient:
 
     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 +295,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)

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

@@ -1,23 +1,185 @@
 """Query function for one-shot interactions with CodeBuddy."""
 
+from __future__ import annotations
+
+import asyncio
 import json
 import os
-from collections.abc import AsyncIterable, AsyncIterator
+from collections.abc import AsyncIterable, AsyncIterator, Callable
 from dataclasses import asdict
-from typing import Any
+from typing import Any, TypeGuard
 
+from ._errors import ExecutionError
 from ._message_parser import parse_message
+from .mcp.sdk_control_server_transport import SdkControlServerTransport
+from .mcp.types import JSONRPCMessage, SdkMcpServer
 from .transport import SubprocessTransport, Transport
 from .types import (
     AppendSystemPrompt,
     CanUseToolOptions,
     CodeBuddyAgentOptions,
+    ErrorMessage,
+    HookEvent,
     HookMatcher,
+    McpSdkServerConfig,
+    McpServerConfig,
     Message,
     ResultMessage,
 )
 
 
+def _is_sdk_mcp_server(config: McpServerConfig) -> TypeGuard[McpSdkServerConfig]:
+    """Type guard to check if config is an SDK MCP server."""
+    return isinstance(config, dict) and config.get("type") == "sdk"
+
+
+def _is_valid_hook_event(event: str) -> TypeGuard[HookEvent]:
+    """Type guard to check if event is a valid HookEvent."""
+    return event in {
+        "PreToolUse",
+        "PostToolUse",
+        "UserPromptSubmit",
+        "Stop",
+        "SubagentStop",
+        "PreCompact",
+    }
+
+
+class QueryContext:
+    """Context for managing query state including SDK MCP servers."""
+
+    def __init__(self, options: CodeBuddyAgentOptions):
+        self.options = options
+        # SDK MCP Server management
+        self.sdk_mcp_transports: dict[str, SdkControlServerTransport] = {}
+        self.sdk_mcp_servers: dict[str, SdkMcpServer] = {}
+        self.pending_mcp_responses: dict[str, asyncio.Future[JSONRPCMessage]] = {}
+        self.sdk_mcp_server_names: list[str] = []
+
+    def extract_mcp_servers(
+        self,
+    ) -> tuple[dict[str, SdkMcpServer], dict[str, McpServerConfig] | None]:
+        """
+        Extract SDK MCP servers from the mcp_servers config.
+        SDK servers are identified by having type: 'sdk'.
+
+        Returns:
+            Tuple of (sdk_servers dict, regular_servers dict or None)
+        """
+        mcp_servers = self.options.mcp_servers
+
+        if not mcp_servers or not isinstance(mcp_servers, dict):
+            return {}, None
+
+        sdk_servers: dict[str, SdkMcpServer] = {}
+        regular_servers: dict[str, McpServerConfig] = {}
+
+        for name, config in mcp_servers.items():
+            if _is_sdk_mcp_server(config):
+                # SDK MCP server
+                sdk_servers[name] = config["server"]
+                self.sdk_mcp_server_names.append(name)
+            else:
+                # Regular MCP server (stdio)
+                regular_servers[name] = config
+
+        return (
+            sdk_servers,
+            regular_servers if regular_servers else None,
+        )
+
+    def connect_sdk_mcp_server(
+        self,
+        name: str,
+        server: SdkMcpServer,
+        send_callback: Callable[[str, JSONRPCMessage], None],
+    ) -> None:
+        """Connect an SDK MCP server."""
+
+        def _create_message_forwarder(
+            server_name: str,
+        ) -> Callable[[JSONRPCMessage], None]:
+            def forwarder(msg: JSONRPCMessage) -> None:
+                send_callback(server_name, msg)
+
+            return forwarder
+
+        # Create custom transport that forwards to CLI
+        transport = SdkControlServerTransport(_create_message_forwarder(name))
+
+        # Store transport and server
+        self.sdk_mcp_transports[name] = transport
+        self.sdk_mcp_servers[name] = server
+
+        # Connect server to transport
+        server.connect(transport)
+
+    async def handle_mcp_message_request(
+        self,
+        transport: Transport,
+        request_id: str,
+        request: dict[str, Any],
+    ) -> None:
+        """Handle MCP message control request from CLI."""
+        server_name = request.get("server_name", "")
+        message: JSONRPCMessage = request.get("message", {})
+
+        server = self.sdk_mcp_servers.get(server_name)
+
+        if not server:
+            response = {
+                "type": "control_response",
+                "response": {
+                    "subtype": "error",
+                    "request_id": request_id,
+                    "error": f"SDK MCP server not found: {server_name}",
+                },
+            }
+            await transport.write(json.dumps(response))
+            return
+
+        try:
+            # Handle the message with the MCP server
+            mcp_response = await server.handle_message(message)
+
+            response = {
+                "type": "control_response",
+                "response": {
+                    "subtype": "success",
+                    "request_id": request_id,
+                    "response": {
+                        "mcp_response": mcp_response or {"jsonrpc": "2.0", "result": {}, "id": 0},
+                    },
+                },
+            }
+            await transport.write(json.dumps(response))
+
+        except Exception as e:
+            response = {
+                "type": "control_response",
+                "response": {
+                    "subtype": "error",
+                    "request_id": request_id,
+                    "error": str(e),
+                },
+            }
+            await transport.write(json.dumps(response))
+
+    async def cleanup(self) -> None:
+        """Cleanup all resources."""
+        # Cancel pending MCP responses
+        for future in self.pending_mcp_responses.values():
+            if not future.done():
+                future.cancel()
+        self.pending_mcp_responses.clear()
+
+        # Close all SDK MCP transports
+        for transport in self.sdk_mcp_transports.values():
+            await transport.close()
+        self.sdk_mcp_transports.clear()
+        self.sdk_mcp_servers.clear()
+
+
 async def query(
     *,
     prompt: str | AsyncIterable[dict[str, Any]],
@@ -52,13 +214,53 @@ async def query(
 
     os.environ["CODEBUDDY_CODE_ENTRYPOINT"] = "sdk-py"
 
+    # Create query context for managing SDK MCP servers
+    ctx = QueryContext(options)
+
+    # Extract SDK MCP servers and regular MCP servers
+    sdk_servers, regular_servers = ctx.extract_mcp_servers()
+
+    # Create modified options with only regular MCP servers for subprocess
+    modified_options = CodeBuddyAgentOptions(
+        allowed_tools=options.allowed_tools,
+        disallowed_tools=options.disallowed_tools,
+        system_prompt=options.system_prompt,
+        mcp_servers=regular_servers or {},
+        permission_mode=options.permission_mode,
+        continue_conversation=options.continue_conversation,
+        resume=options.resume,
+        max_turns=options.max_turns,
+        model=options.model,
+        fallback_model=options.fallback_model,
+        cwd=options.cwd,
+        codebuddy_code_path=options.codebuddy_code_path,
+        env=options.env,
+        extra_args=options.extra_args,
+        stderr=options.stderr,
+        hooks=options.hooks,
+        include_partial_messages=options.include_partial_messages,
+        fork_session=options.fork_session,
+        agents=options.agents,
+        setting_sources=options.setting_sources,
+        can_use_tool=options.can_use_tool,
+    )
+
     if transport is None:
-        transport = SubprocessTransport(options=options, prompt=prompt)
+        transport = SubprocessTransport(options=modified_options, prompt=prompt)
 
     await transport.connect()
 
+    # Connect SDK MCP servers
+    def send_mcp_message(_server_name: str, _message: JSONRPCMessage) -> None:
+        # This callback sends MCP messages from server to CLI
+        # For now, SDK servers only respond to CLI requests, so this is rarely used
+        pass
+
+    for name, server in sdk_servers.items():
+        ctx.connect_sdk_mcp_server(name, server, send_mcp_message)
+
     try:
-        await _send_initialize(transport, options)
+        await _send_initialize(transport, options, ctx.sdk_mcp_server_names)
         await _send_prompt(transport, prompt)
 
         async for line in transport.read():
@@ -68,35 +270,45 @@ async def query(
             try:
                 data = json.loads(line)
 
-                # Handle control requests (hooks)
+                # Handle control requests (hooks, permissions, MCP messages)
                 if data.get("type") == "control_request":
-                    await _handle_control_request(transport, data, options)
+                    await _handle_control_request(transport, data, options, ctx)
                     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, ResultMessage):
+                    if isinstance(message, ErrorMessage):
                         break
 
             except json.JSONDecodeError:
                 continue  # Ignore non-JSON lines
 
     finally:
+        await ctx.cleanup()
         await transport.close()
 
 
-async def _send_initialize(transport: Transport, options: CodeBuddyAgentOptions) -> None:
+async def _send_initialize(
+    transport: Transport,
+    options: CodeBuddyAgentOptions,
+    sdk_mcp_server_names: list[str],
+) -> None:
     """Send initialization control request."""
     hooks_config = _build_hooks_config(options.hooks) if options.hooks else None
     agents_config = (
-        {name: asdict(agent) for name, agent in options.agents.items()}
-        if options.agents
-        else None
+        {name: asdict(agent) for name, agent in options.agents.items()} if options.agents else None
     )
 
-    # 解析 system_prompt 配置
+    # Parse system_prompt config
     system_prompt: str | None = None
     append_system_prompt: str | None = None
     if isinstance(options.system_prompt, str):
@@ -113,14 +325,14 @@ async def _send_initialize(transport: Transport, options: CodeBuddyAgentOptions)
             "systemPrompt": system_prompt,
             "appendSystemPrompt": append_system_prompt,
             "agents": agents_config,
+            # Include SDK MCP server names
+            "sdkMcpServers": sdk_mcp_server_names if sdk_mcp_server_names else None,
         },
     }
     await transport.write(json.dumps(request))
 
 
-async def _send_prompt(
-    transport: Transport, prompt: str | AsyncIterable[dict[str, Any]]
-) -> None:
+async def _send_prompt(transport: Transport, prompt: str | AsyncIterable[dict[str, Any]]) -> None:
     """Send user prompt."""
     if isinstance(prompt, str):
         message = {
@@ -139,6 +351,7 @@ async def _handle_control_request(
     transport: Transport,
     data: dict[str, Any],
     options: CodeBuddyAgentOptions,
+    ctx: QueryContext,
 ) -> None:
     """Handle control request from CLI."""
     request_id = data.get("request_id", "")
@@ -168,6 +381,9 @@ async def _handle_control_request(
     elif subtype == "can_use_tool":
         await _handle_permission_request(transport, request_id, request, options)
 
+    elif subtype == "mcp_message":
+        await ctx.handle_mcp_message_request(transport, request_id, request)
+
 
 async def _handle_permission_request(
     transport: Transport,
@@ -268,14 +484,19 @@ async def _execute_hook(
         return {"continue_": True}
 
     event = parts[1]
+
+    # Validate event is a known HookEvent using TypeGuard
+    if not _is_valid_hook_event(event):
+        return {"continue_": True}
+
     try:
         matcher_idx = int(parts[2])
         hook_idx = int(parts[3])
     except ValueError:
         return {"continue_": True}
 
-    # Find the hook
-    matchers = options.hooks.get(event)  # type: ignore[arg-type]
+    # Find the hook - event is now narrowed to HookEvent by TypeGuard
+    matchers = options.hooks.get(event)
     if not matchers or matcher_idx >= len(matchers):
         return {"continue_": True}
 
@@ -305,9 +526,7 @@ def _build_hooks_config(
         config[str(event)] = [
             {
                 "matcher": m.matcher,
-                "hookCallbackIds": [
-                    f"hook_{event}_{i}_{j}" for j, _ in enumerate(m.hooks)
-                ],
+                "hookCallbackIds": [f"hook_{event}_{i}_{j}" for j, _ in enumerate(m.hooks)],
                 "timeout": m.timeout,
             }
             for i, m in enumerate(matchers)

codebuddy_agent_sdk/transport/base.py

@@ -1,5 +1,7 @@
 """Transport base class for CLI communication."""
 
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator
 

codebuddy_agent_sdk/transport/subprocess.py

@@ -1,5 +1,7 @@
 """Subprocess transport for CLI communication."""
 
+from __future__ import annotations
+
 import asyncio
 import json
 import os
@@ -81,8 +83,10 @@ class SubprocessTransport(Transport):
         # When setting_sources is explicitly provided (including empty list), use it
         # When not provided (None), default to 'none' for SDK isolation
         if opts.setting_sources is not None:
-            value = "none" if len(opts.setting_sources) == 0 else ",".join(opts.setting_sources)
-            args.extend(["--setting-sources", value])
+            setting_value = (
+                "none" if len(opts.setting_sources) == 0 else ",".join(opts.setting_sources)
+            )
+            args.extend(["--setting-sources", setting_value])
         else:
             # SDK default behavior: no filesystem settings loaded
             args.extend(["--setting-sources", "none"])

codebuddy_agent_sdk/types.py

@@ -1,9 +1,14 @@
 """Type definitions for CodeBuddy Agent SDK."""
 
+from __future__ import annotations
+
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Literal, NotRequired, TypedDict
+from typing import TYPE_CHECKING, Any, Literal, NotRequired, TypedDict
+
+if TYPE_CHECKING:
+    from .mcp.types import SdkMcpServer
 
 # ============= AskUserQuestion Types =============
 
@@ -183,6 +188,7 @@ class ResultMessage:
     total_cost_usd: float | None = None
     usage: dict[str, Any] | None = None
     result: str | None = None
+    errors: list[str] | None = None
 
 
 @dataclass
@@ -195,7 +201,17 @@ class StreamEvent:
     parent_tool_use_id: str | None = None
 
 
-Message = UserMessage | AssistantMessage | SystemMessage | ResultMessage | StreamEvent
+@dataclass
+class ErrorMessage:
+    """Error message from CLI."""
+
+    error: str
+    session_id: str | None = None
+
+
+Message = (
+    UserMessage | AssistantMessage | SystemMessage | ResultMessage | StreamEvent | ErrorMessage
+)
 
 
 # Hook types
@@ -242,7 +258,18 @@ class McpStdioServerConfig(TypedDict):
     env: NotRequired[dict[str, str]]
 
 
-McpServerConfig = McpStdioServerConfig
+class McpSdkServerConfig(TypedDict):
+    """
+    SDK MCP Server configuration - for servers running within the SDK process.
+    Created via create_sdk_mcp_server().
+    """
+
+    type: Literal["sdk"]
+    name: str
+    server: SdkMcpServer
+
+
+McpServerConfig = McpStdioServerConfig | McpSdkServerConfig
 
 
 # System prompt configuration

{codebuddy_agent_sdk-0.1.8.dist-info → codebuddy_agent_sdk-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: codebuddy-agent-sdk
-Version: 0.1.8
+Version: 0.2.0
 Summary: CodeBuddy Code SDK for Python
 Author-email: ninoyi <ninoyi@tencent.com>
 Keywords: agent,ai,codebuddy,sdk

codebuddy_agent_sdk-0.2.0.dist-info/RECORD

@@ -0,0 +1,20 @@
+codebuddy_agent_sdk/__init__.py,sha256=GKrVvH2KKyuA9ht-7uMVHJJn7eWBKOi0_DKc2Ef5RCU,2690
+codebuddy_agent_sdk/_binary.py,sha256=rQFj2B__X7zHMxFazp_0vgXQFsqPi578f8g4jSkLUBc,4338
+codebuddy_agent_sdk/_errors.py,sha256=bLmyebUowNlbUomXyuS1L6J6ZbwHohvzqH_1GmN4_oM,1248
+codebuddy_agent_sdk/_message_parser.py,sha256=vtyeuwIFdl9Wz1WrWEg7G6E3OECDZT4G6bvpH15Q3VY,3504
+codebuddy_agent_sdk/_version.py,sha256=EhZlhQ69654jiXymzw5m6O61qdW_fo3ug1b4ABdOtm4,74
+codebuddy_agent_sdk/client.py,sha256=JfdA0wq4KOOrP4mZXfVuF4uwdrEbVHGIUrzpt4Bt_ZM,10803
+codebuddy_agent_sdk/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+codebuddy_agent_sdk/query.py,sha256=tCfFomCNGS1Sygbk0U9j25JLsF2azWa-3AuckThYaII,17341
+codebuddy_agent_sdk/types.py,sha256=rw5jrIg8LWtZN1Km_7TzY7SP_wFRJwKNGwPxgRcQ47A,7473
+codebuddy_agent_sdk/bin/codebuddy,sha256=pfQZlfnGCrEocOsuQPXb8bozsVAPnWKmV62w-EhGVeE,119690167
+codebuddy_agent_sdk/mcp/__init__.py,sha256=0aCmHN0MnMCCUU0bdTWBRlxAblnxA-vR5Og0y3BA480,764
+codebuddy_agent_sdk/mcp/create_sdk_mcp_server.py,sha256=ig0b0ghf0T8GGjltFBRSmQSTFsaJ6fBI8S1Pjupqhps,4959
+codebuddy_agent_sdk/mcp/sdk_control_server_transport.py,sha256=CIG3h5gULubGcQPHPVd-GG9cLswwVDfgDlLba6ZPQbs,2944
+codebuddy_agent_sdk/mcp/types.py,sha256=CPEBUW5vCt93cBgu3HvAacXMsa5FdRWwWxPeicQ0PWw,8578
+codebuddy_agent_sdk/transport/__init__.py,sha256=zv_8OJHgnWjCInqOiu3GOFby4XVGvDwICHpOsymOlko,166
+codebuddy_agent_sdk/transport/base.py,sha256=XtLquCmt4yzPhTBHcVU0XFb016uG0E-12ETc-N9tzIk,662
+codebuddy_agent_sdk/transport/subprocess.py,sha256=VHoGXoP0qm597bhIE_mTV9xyIDUM-nMk7fwZwntx94c,5701
+codebuddy_agent_sdk-0.2.0.dist-info/METADATA,sha256=vATnYz7KEZ_IgyQqOsWBGLiVxdD-Z-_MV1RvKMK4TS0,2538
+codebuddy_agent_sdk-0.2.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+codebuddy_agent_sdk-0.2.0.dist-info/RECORD,,

codebuddy_agent_sdk-0.1.8.dist-info/RECORD

@@ -1,16 +0,0 @@
-codebuddy_agent_sdk/__init__.py,sha256=MF75KZvSF3PTGk9OwjFEW5R7-vI59lexKn5e9-FY1Dw,2052
-codebuddy_agent_sdk/_binary.py,sha256=7UBOzgooaojNAlb3tFvsN12TBAndcP3YLGD3Ic3LlmI,4302
-codebuddy_agent_sdk/_errors.py,sha256=5ABIym6mazInI33kmt5mR94v6c325jDbi0YCgQjrysY,816
-codebuddy_agent_sdk/_message_parser.py,sha256=U0dy1nqe2kap0_A_VWQN-KME-HEQmDhZQcCbCLVMe4s,3255
-codebuddy_agent_sdk/_version.py,sha256=Pi_yM0znG_C1DcEEzAPpeZEEBjrYZVv-JFkiLHeMm98,74
-codebuddy_agent_sdk/client.py,sha256=Bl6l2KFXkSs4j9usuzM7JMYsEO9Np4p6Csg9fR42LtI,10191
-codebuddy_agent_sdk/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-codebuddy_agent_sdk/query.py,sha256=LNv4LIdzdg-fRunMlIicu0SAEv5AoAPpu3_w7iDGhYw,9544
-codebuddy_agent_sdk/types.py,sha256=PYCgC6rAElM5VAcDcdDDr4MXCnIXbRcSNwiuTf__Jrk,6926
-codebuddy_agent_sdk/bin/codebuddy,sha256=9blxztGO9Uu2AbDmfAmMfXbg85EYDj46vUBfgJbmUwk,122398201
-codebuddy_agent_sdk/transport/__init__.py,sha256=zv_8OJHgnWjCInqOiu3GOFby4XVGvDwICHpOsymOlko,166
-codebuddy_agent_sdk/transport/base.py,sha256=ps-4jWU3jL-IiTIL5Pv70zanxdmrknCHZfpxMGBTLoA,626
-codebuddy_agent_sdk/transport/subprocess.py,sha256=m3o-hx2q6U3yeaOPrMxqi6XetsPZdAjjAxBIDiq-vgU,5617
-codebuddy_agent_sdk-0.1.8.dist-info/METADATA,sha256=HXjZoWayfmxEIBJOX5pP9onqYrVWADHL9h1fEwleyaI,2538
-codebuddy_agent_sdk-0.1.8.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-codebuddy_agent_sdk-0.1.8.dist-info/RECORD,,