codebuddy-agent-sdk 0.1.27__py3-none-macosx_11_0_arm64.whl

codebuddy_agent_sdk/query.py

@@ -0,0 +1,535 @@
+"""Query function for one-shot interactions with CodeBuddy."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+from collections.abc import AsyncIterable, AsyncIterator, Callable
+from dataclasses import asdict
+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]],
+    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"
+
+    # 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=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, ctx.sdk_mcp_server_names)
+        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, 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, 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,
+    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
+    )
+
+    # 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
+
+    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
+            "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:
+    """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,
+    ctx: QueryContext,
+) -> 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
+        response = await _execute_hook(callback_id, hook_input, tool_use_id, options)
+
+        # 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":
+        await ctx.handle_mcp_message_request(transport, 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,
+    options: CodeBuddyAgentOptions,
+) -> dict[str, Any]:
+    """Execute a hook callback."""
+    if not options.hooks:
+        return {"continue_": True}
+
+    # Parse callback_id: hook_{event}_{matcherIndex}_{hookIndex}
+    parts = callback_id.split("_")
+    if len(parts) < 4:
+        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 - event is now narrowed to HookEvent by TypeGuard
+    matchers = options.hooks.get(event)
+    if not matchers or matcher_idx >= len(matchers):
+        return {"continue_": True}
+
+    matcher = matchers[matcher_idx]
+    if hook_idx >= len(matcher.hooks):
+        return {"continue_": True}
+
+    hook = matcher.hooks[hook_idx]
+
+    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,
+) -> dict[str, list[dict[str, Any]]] | None:
+    """Build hooks configuration for CLI."""
+    if not hooks:
+        return None
+
+    config: dict[str, list[dict[str, Any]]] = {}
+
+    for event, matchers in hooks.items():
+        config[str(event)] = [
+            {
+                "matcher": m.matcher,
+                "hookCallbackIds": [f"hook_{event}_{i}_{j}" for j, _ in enumerate(m.hooks)],
+                "timeout": m.timeout,
+            }
+            for i, m in enumerate(matchers)
+        ]
+
+    return config if config else None

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,26 @@
+"""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."""