kubiya-control-plane-api 0.9.15__py3-none-any.whl

control_plane_api/worker/runtimes/claude_code/runtime.py

@@ -0,0 +1,1546 @@
+"""
+Claude Code runtime implementation using Claude Code SDK.
+
+This runtime adapter integrates the Claude Code SDK to power agents with
+advanced coding capabilities, file operations, and specialized tools.
+
+ALL 7 BUGS FIXED:
+- Bug #1: Added metadata = {} initialization
+- Bug #2: Replaced print() with logger.debug()
+- Bug #3: Made MCP fallback patterns explicit
+- Bug #4: Added session_id validation
+- Bug #5: Added explicit disconnect() calls with timeout
+- Bug #6: Added tool name validation
+- Bug #7: Removed debug output
+"""
+
+from typing import Dict, Any, Optional, AsyncIterator, Callable, TYPE_CHECKING
+import structlog
+import asyncio
+import time
+import os
+from temporalio import activity
+
+from ..base import (
+    RuntimeType,
+    RuntimeExecutionResult,
+    RuntimeExecutionContext,
+    RuntimeCapabilities,
+    BaseRuntime,
+    RuntimeRegistry,
+)
+from .config import build_claude_options
+from .utils import (
+    extract_usage_from_result_message,
+    extract_session_id_from_result_message,
+    build_prompt_with_history,
+)
+from .litellm_proxy import clear_execution_context
+from .cleanup import cleanup_sdk_client
+from .client_pool import ClaudeCodeClientPool
+
+if TYPE_CHECKING:
+    from control_plane_client import ControlPlaneClient
+    from services.cancellation_manager import CancellationManager
+
+logger = structlog.get_logger(__name__)
+
+# ⚡ PERFORMANCE: Lazy load Claude SDK at module level (not per-execution)
+# This imports the SDK once when the module loads, making subsequent executions faster
+_CLAUDE_SDK_AVAILABLE = False
+_CLAUDE_SDK_IMPORT_ERROR = None
+_SDK_CLASSES = {}
+
+try:
+    from claude_agent_sdk import (
+        ClaudeSDKClient,
+        AssistantMessage,
+        ResultMessage,
+        TextBlock,
+        ToolUseBlock,
+        ToolResultBlock,
+    )
+    _CLAUDE_SDK_AVAILABLE = True
+    _SDK_CLASSES = {
+        'ClaudeSDKClient': ClaudeSDKClient,
+        'AssistantMessage': AssistantMessage,
+        'ResultMessage': ResultMessage,
+        'TextBlock': TextBlock,
+        'ToolUseBlock': ToolUseBlock,
+        'ToolResultBlock': ToolResultBlock,
+    }
+    logger.info("claude_code_sdk_preloaded", status="success")
+except ImportError as e:
+    _CLAUDE_SDK_IMPORT_ERROR = str(e)
+    logger.warning("claude_code_sdk_not_available", error=str(e))
+
+
+@RuntimeRegistry.register(RuntimeType.CLAUDE_CODE)
+class ClaudeCodeRuntime(BaseRuntime):
+    """
+    Runtime implementation using Claude Code SDK.
+
+    This runtime leverages Claude Code's specialized capabilities for
+    software engineering tasks, file operations, and developer workflows.
+
+    Features:
+    - Streaming execution with real-time updates
+    - Conversation history support via ClaudeSDKClient
+    - Custom tool integration via MCP
+    - Hooks for tool execution monitoring
+    - Cancellation support via interrupt()
+
+    All critical bugs have been fixed in this refactored version.
+    """
+
+    def __init__(
+        self,
+        control_plane_client: "ControlPlaneClient",
+        cancellation_manager: "CancellationManager",
+        **kwargs,
+    ):
+        """
+        Initialize the Claude Code runtime.
+
+        Args:
+            control_plane_client: Client for Control Plane API
+            cancellation_manager: Manager for execution cancellation
+            **kwargs: Additional configuration options
+        """
+        super().__init__(control_plane_client, cancellation_manager, **kwargs)
+
+        # Track active SDK clients for cancellation
+        self._active_clients: Dict[str, Any] = {}
+
+        # Track custom MCP servers
+        self._custom_mcp_servers: Dict[str, Any] = {}  # server_name -> mcp_server
+
+        # Cache MCP discovery results (discovered once, reused per execution)
+        # Format: {server_name: {tools: [...], resources: [...], prompts: [...], connected: bool}}
+        self._mcp_discovery_cache: Dict[str, Any] = {}
+        self._mcp_cache_lock = None  # Will be initialized on first use
+
+    def get_runtime_type(self) -> RuntimeType:
+        """Return RuntimeType.CLAUDE_CODE."""
+        return RuntimeType.CLAUDE_CODE
+
+    def get_capabilities(self) -> RuntimeCapabilities:
+        """Return Claude Code runtime capabilities."""
+        return RuntimeCapabilities(
+            streaming=True,
+            tools=True,
+            mcp=True,
+            hooks=True,
+            cancellation=True,
+            conversation_history=True,
+            custom_tools=True,
+        )
+
+    async def _execute_impl(
+        self, context: RuntimeExecutionContext
+    ) -> RuntimeExecutionResult:
+        """
+        Execute agent using Claude Code SDK (non-streaming).
+
+        Production-grade implementation with:
+        - Comprehensive error handling
+        - Proper resource cleanup
+        - Detailed logging
+        - Timeout management
+        - Graceful degradation
+
+        BUG FIX #1: Added metadata = {} initialization
+        BUG FIX #5: Added explicit disconnect() call
+
+        Args:
+            context: Execution context with prompt, history, config
+
+        Returns:
+            RuntimeExecutionResult with response and metadata
+        """
+        client = None
+        start_time = asyncio.get_event_loop().time()
+
+        try:
+            # ⚡ PERFORMANCE: Use pre-loaded SDK classes (loaded at module level)
+            if not _CLAUDE_SDK_AVAILABLE:
+                return RuntimeExecutionResult(
+                    response="",
+                    usage={},
+                    success=False,
+                    error=f"Claude Code SDK not available: {_CLAUDE_SDK_IMPORT_ERROR}",
+                )
+
+            ClaudeSDKClient = _SDK_CLASSES['ClaudeSDKClient']
+            ResultMessage = _SDK_CLASSES['ResultMessage']
+
+            self.logger.info(
+                "starting_claude_code_non_streaming_execution",
+                execution_id=context.execution_id,
+                model=context.model_id,
+                has_history=bool(context.conversation_history),
+            )
+
+            # Build options and create client
+            options, _, _, _ = await build_claude_options(context, runtime=self)
+
+            # Merge custom MCP servers
+            if self._custom_mcp_servers:
+                if not options.mcp_servers:
+                    options.mcp_servers = {}
+
+                for server_name, mcp_server in self._custom_mcp_servers.items():
+                    options.mcp_servers[server_name] = mcp_server
+
+                    # Add tool names to allowed_tools for permission
+                    if hasattr(mcp_server, 'tools') and mcp_server.tools:
+                        for tool in mcp_server.tools:
+                            if hasattr(tool, 'name'):
+                                tool_name = f"mcp__{server_name}__{tool.name}"
+                                if tool_name not in options.allowed_tools:
+                                    options.allowed_tools.append(tool_name)
+
+                    self.logger.debug(
+                        "custom_mcp_server_added_non_streaming",
+                        server_name=server_name,
+                        execution_id=context.execution_id
+                    )
+
+            # Suppress verbose MCP STDIO parsing error logs
+            # These errors occur when MCP servers incorrectly log to stdout instead of stderr
+            # The connection continues to work, but the SDK logs errors for each non-JSONRPC line
+            import logging
+            mcp_stdio_logger = logging.getLogger("mcp.client.stdio")
+            original_stdio_level = mcp_stdio_logger.level
+            mcp_stdio_logger.setLevel(logging.ERROR)  # Only show critical errors
+
+            # Create and connect client
+            client = ClaudeSDKClient(options=options)
+
+            try:
+                await client.connect()
+            except Exception as e:
+                error_msg = str(e)
+
+                # Build comprehensive error message from all sources
+                # The SDK wraps subprocess errors, so "No conversation found" may be in:
+                # 1. The exception message itself
+                # 2. The stderr attribute (if CalledProcessError)
+                # 3. The output attribute
+                # 4. The __cause__ chain
+                full_error_context = error_msg
+                if hasattr(e, 'stderr') and e.stderr:
+                    stderr_content = e.stderr if isinstance(e.stderr, str) else e.stderr.decode('utf-8', errors='replace')
+                    full_error_context += f" | stderr: {stderr_content}"
+                if hasattr(e, 'output') and e.output:
+                    output_content = e.output if isinstance(e.output, str) else e.output.decode('utf-8', errors='replace')
+                    full_error_context += f" | output: {output_content}"
+                if e.__cause__:
+                    full_error_context += f" | cause: {str(e.__cause__)}"
+
+                # Detect session-related errors:
+                # 1. Explicit "No conversation found" in any error context
+                # 2. "exit code 1" when we were trying to resume (likely session issue)
+                is_session_error = (
+                    "No conversation found" in full_error_context or
+                    "conversation" in full_error_context.lower() or
+                    # Fallback: exit code 1 during resume attempt is likely a session issue
+                    (options.resume and "exit code 1" in error_msg.lower())
+                )
+
+                # Handle session-related errors gracefully by retrying without resume
+                if is_session_error and options.resume:
+                    logger.warning(
+                        "session_resume_failed_retrying_without_resume",
+                        execution_id=context.execution_id[:16],
+                        error=error_msg,
+                        full_error_context=full_error_context[:500],  # Truncate for logging
+                        session_id=options.resume[:16] if options.resume else None,
+                        note="Session doesn't exist or is invalid, creating new conversation"
+                    )
+
+                    # Retry without resume parameter
+                    options.resume = None
+                    session_resume_failed = True  # Track for prompt building
+                    client = ClaudeSDKClient(options=options)
+                    await client.connect()
+                    logger.info(
+                        "client_connected_without_resume",
+                        execution_id=context.execution_id[:16]
+                    )
+                else:
+                    # Different error, re-raise
+                    raise
+            else:
+                session_resume_failed = False
+
+            self._active_clients[context.execution_id] = client
+
+            # Send prompt - include history in prompt if session resume failed
+            # (since we can't rely on SDK's session continuity)
+            if session_resume_failed and context.conversation_history:
+                prompt = build_prompt_with_history(context)
+                logger.info(
+                    "using_prompt_with_history_fallback",
+                    execution_id=context.execution_id[:16],
+                    history_messages=len(context.conversation_history),
+                    note="Session resume failed, including history in prompt"
+                )
+            else:
+                prompt = context.prompt
+
+            self.logger.debug(
+                "sending_query_to_claude_code_sdk",
+                execution_id=context.execution_id,
+                prompt_length=len(prompt),
+                using_session_resume=bool(options.resume),
+            )
+
+            await client.query(prompt)
+
+            # Collect complete response
+            response_text = ""
+            usage = {}
+            tool_messages = []
+            finish_reason = None
+            message_count = 0
+            last_heartbeat = asyncio.get_event_loop().time()  # Track last heartbeat for Temporal activity liveness
+
+            # BUG FIX #1: Initialize metadata before use
+            metadata = {}
+
+            # Use receive_response() to get messages until ResultMessage
+            async for message in client.receive_response():
+                message_count += 1
+
+                # Send heartbeat every 5 seconds or every 10 messages
+                current_time = asyncio.get_event_loop().time()
+                if current_time - last_heartbeat > 5 or message_count % 10 == 0:
+                    try:
+                        activity.heartbeat({
+                            "status": "processing",
+                            "messages_received": message_count,
+                            "response_length": len(response_text),
+                            "elapsed_seconds": int(current_time - last_heartbeat)
+                        })
+                        last_heartbeat = current_time
+                    except Exception as e:
+                        # Non-fatal: heartbeat failure should not break execution
+                        self.logger.warning("heartbeat_failed_non_fatal", execution_id=context.execution_id, error=str(e))
+
+                # Extract content from AssistantMessage
+                if hasattr(message, "content"):
+                    for block in message.content:
+                        if hasattr(block, "text"):
+                            response_text += block.text
+                        elif hasattr(block, "name"):  # ToolUseBlock
+                            tool_messages.append(
+                                {
+                                    "tool": block.name,
+                                    "input": getattr(block, "input", {}),
+                                    "tool_use_id": getattr(block, "id", None),
+                                }
+                            )
+
+                # Extract usage, finish reason, and session_id from ResultMessage
+                if isinstance(message, ResultMessage):
+                    usage = extract_usage_from_result_message(message)
+
+                    if usage:
+                        self.logger.info(
+                            "claude_code_usage_extracted",
+                            execution_id=context.execution_id[:8],
+                            input_tokens=usage["input_tokens"],
+                            output_tokens=usage["output_tokens"],
+                            cache_read=usage["cache_read_tokens"],
+                        )
+
+                    finish_reason = message.subtype  # "success" or "error"
+
+                    # BUG FIX #4: Extract and validate session_id
+                    session_id = extract_session_id_from_result_message(
+                        message, context.execution_id
+                    )
+
+                    if session_id:
+                        # BUG FIX #1: metadata is now properly initialized
+                        metadata["claude_code_session_id"] = session_id
+
+                    self.logger.info(
+                        "claude_code_execution_completed",
+                        execution_id=context.execution_id[:8],
+                        finish_reason=finish_reason,
+                        message_count=message_count,
+                        response_length=len(response_text),
+                        tool_count=len(tool_messages),
+                        tokens=usage.get("total_tokens", 0),
+                        has_session_id=bool(session_id),
+                    )
+                    break
+
+            elapsed_time = asyncio.get_event_loop().time() - start_time
+
+            # Merge metadata with execution stats
+            final_metadata = {
+                **metadata,  # Includes claude_code_session_id if present
+                "elapsed_time": elapsed_time,
+                "message_count": message_count,
+            }
+
+            return RuntimeExecutionResult(
+                response=response_text,
+                usage=usage,
+                success=finish_reason == "success",
+                finish_reason=finish_reason or "stop",
+                tool_execution_messages=tool_messages,  # Use standard field name for analytics
+                tool_messages=tool_messages,  # Keep for backward compatibility
+                model=context.model_id,
+                metadata=final_metadata,
+            )
+
+        except ImportError as e:
+            self.logger.error(
+                "claude_code_sdk_not_installed",
+                execution_id=context.execution_id,
+                error=str(e),
+            )
+
+            # Publish error event
+            try:
+                from control_plane_api.worker.utils.error_publisher import (
+                    ErrorEventPublisher, ErrorSeverity, ErrorCategory
+                )
+                error_publisher = ErrorEventPublisher(self.control_plane)
+                await error_publisher.publish_error(
+                    execution_id=context.execution_id,
+                    exception=e,
+                    severity=ErrorSeverity.CRITICAL,
+                    category=ErrorCategory.RUNTIME_INIT,
+                    stage="initialization",
+                    component="claude_code_runtime",
+                    operation="sdk_import",
+                    recovery_actions=[
+                        "Install Claude Code SDK: pip install claude-agent-sdk",
+                        "Verify SDK version is compatible",
+                        "Check Python environment configuration",
+                    ],
+                )
+            except Exception as publish_error:
+                # Log error publishing failure but don't let it break execution flow
+                self.logger.error(
+                    "error_publish_failed",
+                    error=str(publish_error),
+                    error_type=type(publish_error).__name__,
+                    original_error="Claude Code SDK not available",
+                    execution_id=context.execution_id
+                )
+
+            return RuntimeExecutionResult(
+                response="",
+                usage={},
+                success=False,
+                error=f"Claude Code SDK not available: {str(e)}",
+            )
+
+        except asyncio.TimeoutError:
+            elapsed_time = asyncio.get_event_loop().time() - start_time
+            self.logger.error(
+                "claude_code_execution_timeout",
+                execution_id=context.execution_id,
+                elapsed_time=elapsed_time,
+            )
+
+            # Publish timeout error event
+            try:
+                from control_plane_api.worker.utils.error_publisher import (
+                    ErrorEventPublisher, ErrorSeverity, ErrorCategory
+                )
+                error_publisher = ErrorEventPublisher(self.control_plane)
+                await error_publisher.publish_error(
+                    execution_id=context.execution_id,
+                    exception=asyncio.TimeoutError("Execution timeout exceeded"),
+                    severity=ErrorSeverity.ERROR,
+                    category=ErrorCategory.TIMEOUT,
+                    stage="execution",
+                    component="claude_code_runtime",
+                    operation="agent_execution",
+                    metadata={"elapsed_time": elapsed_time},
+                    recovery_actions=[
+                        "Simplify the prompt or reduce complexity",
+                        "Increase timeout settings if appropriate",
+                        "Check system resources and load",
+                    ],
+                )
+            except Exception as publish_error:
+                # Log error publishing failure but don't let it break execution flow
+                self.logger.error(
+                    "error_publish_failed",
+                    error=str(publish_error),
+                    error_type=type(publish_error).__name__,
+                    original_error="Execution timeout",
+                    execution_id=context.execution_id
+                )
+
+            return RuntimeExecutionResult(
+                response="",
+                usage={},
+                success=False,
+                error="Execution timeout exceeded",
+            )
+
+        except asyncio.CancelledError:
+            self.logger.warning(
+                "claude_code_execution_cancelled_gracefully",
+                execution_id=context.execution_id,
+            )
+            # DURABILITY FIX: Do NOT re-raise! Handle cancellation gracefully
+            # Return partial result to allow workflow to handle interruption
+            return RuntimeExecutionResult(
+                response="",  # No response accumulated in non-streaming mode
+                usage={},
+                success=False,  # Non-streaming cancellation is a failure (no partial state)
+                error="Execution was cancelled",
+                finish_reason="cancelled",
+                metadata={"interrupted": True, "can_resume": False},
+            )
+
+        except Exception as e:
+            self.logger.error(
+                "claude_code_execution_failed",
+                execution_id=context.execution_id,
+                error=str(e),
+                error_type=type(e).__name__,
+                exc_info=True,
+            )
+
+            # Publish generic error event with stack trace
+            try:
+                from control_plane_api.worker.utils.error_publisher import (
+                    ErrorEventPublisher, ErrorSeverity, ErrorCategory
+                )
+                error_publisher = ErrorEventPublisher(self.control_plane)
+                await error_publisher.publish_error(
+                    execution_id=context.execution_id,
+                    exception=e,
+                    severity=ErrorSeverity.CRITICAL,
+                    category=ErrorCategory.UNKNOWN,
+                    stage="execution",
+                    component="claude_code_runtime",
+                    operation="agent_execution",
+                    include_stack_trace=True,
+                )
+            except Exception as publish_error:
+                # Log error publishing failure but don't let it break execution flow
+                self.logger.error(
+                    "error_publish_failed",
+                    error=str(publish_error),
+                    error_type=type(publish_error).__name__,
+                    original_error=f"{type(e).__name__}: {str(e)}",
+                    execution_id=context.execution_id
+                )
+
+            return RuntimeExecutionResult(
+                response="",
+                usage={},
+                success=False,
+                error=f"{type(e).__name__}: {str(e)}",
+            )
+
+        finally:
+            # Clear execution context from proxy (with delay to allow in-flight SDK requests)
+            try:
+                clear_execution_context(
+                    context.execution_id,
+                    immediate=False,  # Use delayed cleanup
+                    delay_seconds=5.0  # Wait for in-flight SDK requests
+                )
+            except Exception as e:
+                self.logger.warning(
+                    "failed_to_clear_proxy_context",
+                    execution_id=context.execution_id,
+                    error=str(e),
+                )
+
+            # Restore MCP STDIO log level
+            try:
+                import logging
+                mcp_stdio_logger = logging.getLogger("mcp.client.stdio")
+                if 'original_stdio_level' in locals():
+                    mcp_stdio_logger.setLevel(original_stdio_level)
+            except Exception as log_level_error:
+                # Log but ignore errors restoring log level - this is non-critical cleanup
+                self.logger.debug(
+                    "failed_to_restore_log_level",
+                    error=str(log_level_error),
+                    execution_id=context.execution_id
+                )
+
+            # CRITICAL: Cleanup SDK client
+            if context.execution_id in self._active_clients:
+                client = self._active_clients.pop(context.execution_id)
+                cleanup_sdk_client(client, context.execution_id, self.logger)
+
+    async def _stream_execute_impl(
+        self,
+        context: RuntimeExecutionContext,
+        event_callback: Optional[Callable[[Dict], None]] = None,
+    ) -> AsyncIterator[RuntimeExecutionResult]:
+        """
+        Production-grade streaming execution with Claude Code SDK.
+
+        This implementation provides:
+        - Comprehensive error handling with specific exception types
+        - Detailed structured logging at each stage
+        - Proper resource cleanup with finally blocks
+        - Real-time event callbacks for tool execution
+        - Accumulated metrics and metadata tracking
+
+        BUG FIX #5: Added explicit disconnect() call
+        BUG FIX #7: Removed all debug output
+
+        Args:
+            context: Execution context with prompt, history, config
+            event_callback: Optional callback for real-time events
+
+        Yields:
+            RuntimeExecutionResult chunks as they arrive, ending with final metadata
+        """
+        client = None
+        start_time = asyncio.get_event_loop().time()
+        chunk_count = 0
+
+        try:
+            # ⚡ PERFORMANCE: Use pre-loaded SDK classes (loaded at module level)
+            if not _CLAUDE_SDK_AVAILABLE:
+                yield RuntimeExecutionResult(
+                    response="",
+                    usage={},
+                    success=False,
+                    error=f"Claude Code SDK not available: {_CLAUDE_SDK_IMPORT_ERROR}",
+                    finish_reason="error",
+                    tool_messages=[],
+                    tool_execution_messages=[],
+                )
+                return
+
+            ClaudeSDKClient = _SDK_CLASSES['ClaudeSDKClient']
+            AssistantMessage = _SDK_CLASSES['AssistantMessage']
+            ResultMessage = _SDK_CLASSES['ResultMessage']
+            TextBlock = _SDK_CLASSES['TextBlock']
+            ToolUseBlock = _SDK_CLASSES['ToolUseBlock']
+            ToolResultBlock = _SDK_CLASSES['ToolResultBlock']
+
+            self.logger.info(
+                "starting_claude_code_streaming_execution",
+                execution_id=context.execution_id,
+                model=context.model_id,
+                has_history=bool(context.conversation_history),
+                has_callback=event_callback is not None,
+            )
+
+            # Build options and create client
+            options, active_tools, started_tools, completed_tools = await build_claude_options(context, event_callback, runtime=self)
+
+            # Merge custom MCP servers
+            if self._custom_mcp_servers:
+                if not options.mcp_servers:
+                    options.mcp_servers = {}
+
+                for server_name, mcp_server in self._custom_mcp_servers.items():
+                    options.mcp_servers[server_name] = mcp_server
+
+                    # Add tool names to allowed_tools for permission
+                    if hasattr(mcp_server, 'tools') and mcp_server.tools:
+                        for tool in mcp_server.tools:
+                            if hasattr(tool, 'name'):
+                                tool_name = f"mcp__{server_name}__{tool.name}"
+                                if tool_name not in options.allowed_tools:
+                                    options.allowed_tools.append(tool_name)
+
+                    self.logger.debug(
+                        "custom_mcp_server_added_streaming",
+                        server_name=server_name,
+                        execution_id=context.execution_id
+                    )
+
+            self.logger.info(
+                "created_claude_code_sdk_options",
+                execution_id=context.execution_id,
+                has_tools=bool(context.skills),
+                has_mcp=(
+                    len(options.mcp_servers) > 0
+                    if hasattr(options, "mcp_servers")
+                    else False
+                ),
+                has_custom_mcp=len(self._custom_mcp_servers) > 0,
+                has_hooks=bool(options.hooks) if hasattr(options, "hooks") else False,
+                has_event_callback=event_callback is not None,
+            )
+
+            # Suppress verbose MCP STDIO parsing error logs
+            # These errors occur when MCP servers incorrectly log to stdout instead of stderr
+            # The connection continues to work, but the SDK logs errors for each non-JSONRPC line
+            import logging
+            mcp_stdio_logger = logging.getLogger("mcp.client.stdio")
+            original_stdio_level = mcp_stdio_logger.level
+            mcp_stdio_logger.setLevel(logging.ERROR)  # Only show critical errors
+
+            # Create and connect client
+            client = ClaudeSDKClient(options=options)
+
+            try:
+                await client.connect()
+            except Exception as e:
+                error_msg = str(e)
+
+                # Build comprehensive error message from all sources
+                # The SDK wraps subprocess errors, so "No conversation found" may be in:
+                # 1. The exception message itself
+                # 2. The stderr attribute (if CalledProcessError)
+                # 3. The output attribute
+                # 4. The __cause__ chain
+                full_error_context = error_msg
+                if hasattr(e, 'stderr') and e.stderr:
+                    stderr_content = e.stderr if isinstance(e.stderr, str) else e.stderr.decode('utf-8', errors='replace')
+                    full_error_context += f" | stderr: {stderr_content}"
+                if hasattr(e, 'output') and e.output:
+                    output_content = e.output if isinstance(e.output, str) else e.output.decode('utf-8', errors='replace')
+                    full_error_context += f" | output: {output_content}"
+                if e.__cause__:
+                    full_error_context += f" | cause: {str(e.__cause__)}"
+
+                # Detect session-related errors:
+                # 1. Explicit "No conversation found" in any error context
+                # 2. "exit code 1" when we were trying to resume (likely session issue)
+                is_session_error = (
+                    "No conversation found" in full_error_context or
+                    "conversation" in full_error_context.lower() or
+                    # Fallback: exit code 1 during resume attempt is likely a session issue
+                    (options.resume and "exit code 1" in error_msg.lower())
+                )
+
+                # Handle session-related errors gracefully by retrying without resume
+                if is_session_error and options.resume:
+                    logger.warning(
+                        "session_resume_failed_retrying_without_resume_streaming",
+                        execution_id=context.execution_id[:16],
+                        error=error_msg,
+                        full_error_context=full_error_context[:500],  # Truncate for logging
+                        session_id=options.resume[:16] if options.resume else None,
+                        note="Session doesn't exist or is invalid, creating new conversation"
+                    )
+
+                    # Retry without resume parameter
+                    options.resume = None
+                    session_resume_failed = True  # Track for prompt building
+                    client = ClaudeSDKClient(options=options)
+                    await client.connect()
+                    logger.info(
+                        "client_connected_without_resume_streaming",
+                        execution_id=context.execution_id[:16]
+                    )
+                else:
+                    # Different error, re-raise
+                    raise
+            else:
+                session_resume_failed = False
+
+            self._active_clients[context.execution_id] = client
+
+            # Cache execution metadata
+            try:
+                self.control_plane.cache_metadata(context.execution_id, "AGENT")
+            except Exception as cache_error:
+                self.logger.warning(
+                    "failed_to_cache_metadata_non_fatal",
+                    execution_id=context.execution_id,
+                    error=str(cache_error),
+                )
+
+            # Send prompt - include history in prompt if session resume failed
+            # (since we can't rely on SDK's session continuity)
+            if session_resume_failed and context.conversation_history:
+                prompt = build_prompt_with_history(context)
+                logger.info(
+                    "using_prompt_with_history_fallback_streaming",
+                    execution_id=context.execution_id[:16],
+                    history_messages=len(context.conversation_history),
+                    note="Session resume failed, including history in prompt"
+                )
+            else:
+                prompt = context.prompt
+
+            self.logger.debug(
+                "sending_streaming_query_to_claude_code_sdk",
+                execution_id=context.execution_id,
+                prompt_length=len(prompt),
+                using_session_resume=bool(options.resume),
+            )
+
+            await client.query(prompt)
+
+            # Stream messages
+            accumulated_response = ""
+            accumulated_usage = {}
+            tool_messages = []
+            message_count = 0
+            received_stream_events = False  # Track if we got streaming events
+            session_id = None  # Initialize to avoid UnboundLocalError in exception handlers
+            last_heartbeat = time.time()  # Track last heartbeat for Temporal activity liveness
+
+            # completed_tools set is now passed from build_claude_options for tracking
+            # which tool_use_ids have published completion events (prevents duplicates and detects missing)
+
+            # Generate unique message_id for this turn
+            message_id = f"{context.execution_id}_{int(time.time() * 1000000)}"
+
+            # Track thinking state for extended thinking support
+            current_thinking_block = None  # Tracks active thinking block index
+            thinking_content_buffer = ""   # Accumulates thinking content for logging
+
+            async for message in client.receive_response():
+                message_count += 1
+                message_type_name = type(message).__name__
+
+                # Handle StreamEvent messages (partial chunks)
+                if message_type_name == "StreamEvent":
+                    if hasattr(message, "event") and message.event:
+                        event_data = message.event
+
+                        # Extract text from event data
+                        content = None
+                        if isinstance(event_data, dict):
+                            event_type = event_data.get("type")
+
+                            # Handle content_block_start events (detect thinking blocks)
+                            if event_type == "content_block_start":
+                                content_block = event_data.get("content_block", {})
+                                block_type = content_block.get("type")
+                                block_index = event_data.get("index", 0)
+
+                                if block_type == "thinking":
+                                    current_thinking_block = block_index
+                                    thinking_content_buffer = ""
+
+                                    # Emit thinking_start event
+                                    if event_callback:
+                                        try:
+                                            event_callback({
+                                                "type": "thinking_start",
+                                                "message_id": message_id,
+                                                "index": block_index,
+                                                "execution_id": context.execution_id,
+                                            })
+                                        except Exception as callback_error:
+                                            self.logger.warning(
+                                                "thinking_start_callback_failed",
+                                                execution_id=context.execution_id,
+                                                error=str(callback_error),
+                                            )
+                                continue  # Don't process content_block_start further
+
+                            # Handle content_block_delta events
+                            if event_type == "content_block_delta":
+                                delta = event_data.get("delta", {})
+                                delta_type = delta.get("type") if isinstance(delta, dict) else None
+
+                                # Handle thinking_delta events
+                                if delta_type == "thinking_delta":
+                                    thinking_text = delta.get("thinking", "")
+                                    if thinking_text:
+                                        thinking_content_buffer += thinking_text
+
+                                        # Emit thinking_delta event
+                                        if event_callback:
+                                            try:
+                                                event_callback({
+                                                    "type": "thinking_delta",
+                                                    "thinking": thinking_text,
+                                                    "message_id": message_id,
+                                                    "index": event_data.get("index", 0),
+                                                    "execution_id": context.execution_id,
+                                                })
+                                            except Exception as callback_error:
+                                                self.logger.warning(
+                                                    "thinking_delta_callback_failed",
+                                                    execution_id=context.execution_id,
+                                                    error=str(callback_error),
+                                                )
+                                    continue  # Don't process thinking as regular content
+
+                                # Handle signature_delta events (end of thinking block)
+                                if delta_type == "signature_delta":
+                                    signature = delta.get("signature", "")
+
+                                    # Emit thinking_complete event
+                                    if event_callback:
+                                        try:
+                                            event_callback({
+                                                "type": "thinking_complete",
+                                                "signature": signature,
+                                                "message_id": message_id,
+                                                "index": event_data.get("index", 0),
+                                                "execution_id": context.execution_id,
+                                            })
+                                        except Exception as callback_error:
+                                            self.logger.warning(
+                                                "thinking_complete_callback_failed",
+                                                execution_id=context.execution_id,
+                                                error=str(callback_error),
+                                            )
+
+                                    # Log thinking summary
+                                    if thinking_content_buffer:
+                                        self.logger.debug(
+                                            "thinking_block_completed",
+                                            execution_id=context.execution_id,
+                                            thinking_length=len(thinking_content_buffer),
+                                            has_signature=bool(signature),
+                                        )
+                                    current_thinking_block = None
+                                    thinking_content_buffer = ""
+                                    continue  # Don't process signature as regular content
+
+                                # Handle text_delta events (regular text content)
+                                if isinstance(delta, dict):
+                                    content = delta.get("text")
+                                elif isinstance(delta, str):
+                                    content = delta
+
+                            # Fallback: try direct text extraction
+                            if not content:
+                                content = event_data.get("text") or event_data.get(
+                                    "content"
+                                )
+
+                        elif isinstance(event_data, str):
+                            content = event_data
+                        elif hasattr(event_data, "content"):
+                            content = event_data.content
+                        elif hasattr(event_data, "text"):
+                            content = event_data.text
+
+                        if content:
+                            received_stream_events = True
+                            chunk_count += 1
+                            accumulated_response += content
+
+                            # Publish event
+                            if event_callback:
+                                try:
+                                    event_callback(
+                                        {
+                                            "type": "content_chunk",
+                                            "content": content,
+                                            "message_id": message_id,
+                                            "execution_id": context.execution_id,
+                                        }
+                                    )
+                                except Exception as callback_error:
+                                    self.logger.warning(
+                                        "stream_event_callback_failed",
+                                        execution_id=context.execution_id,
+                                        error=str(callback_error),
+                                    )
+
+                            # Yield chunk with explicit empty arrays for frontend compatibility
+                            # Frontend expects arrays, not None, to avoid R.map errors
+                            yield RuntimeExecutionResult(
+                                response=content,
+                                usage={},
+                                success=True,
+                                tool_messages=[],
+                                tool_execution_messages=[],
+                            )
+
+                            # Send heartbeat every 10 seconds or every 50 chunks (matches AgnoRuntime pattern)
+                            current_time = time.time()
+                            if current_time - last_heartbeat > 10 or chunk_count % 50 == 0:
+                                try:
+                                    activity.heartbeat({
+                                        "status": "streaming",
+                                        "chunks_received": chunk_count,
+                                        "response_length": len(accumulated_response),
+                                        "elapsed_seconds": int(current_time - last_heartbeat)
+                                    })
+                                    last_heartbeat = current_time
+                                except Exception as e:
+                                    # Non-fatal: heartbeat failure should not break execution
+                                    self.logger.warning("heartbeat_failed_non_fatal", execution_id=context.execution_id, error=str(e))
+
+                    continue  # Skip to next message
+
+                # Handle assistant messages (final complete message)
+                if isinstance(message, AssistantMessage):
+                    for block in message.content:
+                        if isinstance(block, TextBlock):
+                            # Skip if already streamed via StreamEvents
+                            if received_stream_events:
+                                continue
+
+                            # Only send if we didn't receive StreamEvents
+                            chunk_count += 1
+                            accumulated_response += block.text
+
+                            if event_callback:
+                                try:
+                                    event_callback(
+                                        {
+                                            "type": "content_chunk",
+                                            "content": block.text,
+                                            "message_id": message_id,
+                                            "execution_id": context.execution_id,
+                                        }
+                                    )
+                                except Exception as callback_error:
+                                    self.logger.warning(
+                                        "event_callback_failed_non_fatal",
+                                        execution_id=context.execution_id,
+                                        error=str(callback_error),
+                                    )
+
+                            # Frontend expects arrays, not None, to avoid R.map errors
+                            yield RuntimeExecutionResult(
+                                response=block.text,
+                                usage={},
+                                success=True,
+                                tool_messages=[],
+                                tool_execution_messages=[],
+                            )
+
+                            # Send heartbeat every 10 seconds or every 50 chunks (matches AgnoRuntime pattern)
+                            current_time = time.time()
+                            if current_time - last_heartbeat > 10 or chunk_count % 50 == 0:
+                                try:
+                                    activity.heartbeat({
+                                        "status": "streaming",
+                                        "chunks_received": chunk_count,
+                                        "response_length": len(accumulated_response),
+                                        "elapsed_seconds": int(current_time - last_heartbeat)
+                                    })
+                                    last_heartbeat = current_time
+                                except Exception as e:
+                                    # Non-fatal: heartbeat failure should not break execution
+                                    self.logger.warning("heartbeat_failed_non_fatal", execution_id=context.execution_id, error=str(e))
+
+                        elif isinstance(block, ToolUseBlock):
+                            # Tool use event - Store for later lookup
+                            tool_info = {
+                                "tool": block.name,
+                                "input": block.input,
+                                "tool_use_id": block.id,
+                            }
+                            tool_messages.append(tool_info)
+                            active_tools[block.id] = block.name
+
+                            # Publish tool_start event from ToolUseBlock (with deduplication)
+                            # This ensures built-in tools like TodoWrite that skip hooks still get tool_start events
+                            if event_callback and block.id not in started_tools:
+                                try:
+                                    event_callback(
+                                        {
+                                            "type": "tool_start",
+                                            "tool_name": block.name,
+                                            "tool_args": block.input,  # Include tool input for frontend rendering
+                                            "tool_execution_id": block.id,
+                                            "execution_id": context.execution_id,
+                                        }
+                                    )
+                                    started_tools.add(block.id)
+                                    self.logger.debug(
+                                        "tool_start_published_via_stream",
+                                        tool_use_id=block.id,
+                                        tool_name=block.name,
+                                    )
+                                except Exception as callback_error:
+                                    self.logger.error(
+                                        "failed_to_publish_tool_start_from_stream",
+                                        tool_use_id=block.id,
+                                        tool_name=block.name,
+                                        error=str(callback_error),
+                                    )
+
+                        elif isinstance(block, ToolResultBlock):
+                            # Tool result - Look up tool name from active_tools
+                            tool_name = active_tools.get(block.tool_use_id, "unknown")
+                            if tool_name == "unknown":
+                                self.logger.warning(
+                                    "could_not_find_tool_name_for_tool_use_id",
+                                    execution_id=context.execution_id,
+                                    tool_use_id=block.tool_use_id,
+                                    active_tools_keys=list(active_tools.keys()),
+                                )
+
+                            status = "success" if not block.is_error else "failed"
+
+                            # Publish via callback (with deduplication)
+                            if event_callback and block.tool_use_id not in completed_tools:
+                                try:
+                                    event_callback(
+                                        {
+                                            "type": "tool_complete",
+                                            "tool_name": tool_name,
+                                            "tool_execution_id": block.tool_use_id,
+                                            "status": status,
+                                            "output": (
+                                                str(block.content)[:1000]
+                                                if block.content
+                                                else None
+                                            ),
+                                            "error": (
+                                                str(block.content)
+                                                if block.is_error
+                                                else None
+                                            ),
+                                            "execution_id": context.execution_id,
+                                        }
+                                    )
+                                    # Mark as completed to prevent duplicate events
+                                    completed_tools.add(block.tool_use_id)
+                                    self.logger.debug(
+                                        "tool_complete_published_via_stream",
+                                        tool_use_id=block.tool_use_id,
+                                        tool_name=tool_name,
+                                    )
+                                except Exception as callback_error:
+                                    self.logger.error(
+                                        "tool_complete_callback_failed",
+                                        execution_id=context.execution_id,
+                                        tool_name=tool_name,
+                                        error=str(callback_error),
+                                        exc_info=True,
+                                    )
+                            elif block.tool_use_id in completed_tools:
+                                self.logger.debug(
+                                    "tool_complete_already_published_via_hooks",
+                                    tool_use_id=block.tool_use_id,
+                                    tool_name=tool_name,
+                                )
+
+                # Handle result message (final)
+                elif isinstance(message, ResultMessage):
+                    accumulated_usage = extract_usage_from_result_message(message)
+
+                    # BUG FIX #4: Extract and validate session_id
+                    session_id = extract_session_id_from_result_message(
+                        message, context.execution_id
+                    )
+
+                    elapsed_time = asyncio.get_event_loop().time() - start_time
+
+                    # FALLBACK: Detect missing tool completion events
+                    # Check if any tool_use_ids in tool_messages are not in completed_tools
+                    missing_completions = []
+                    for tool_info in tool_messages:
+                        tool_use_id = tool_info.get("tool_use_id")
+                        if tool_use_id and tool_use_id not in completed_tools:
+                            missing_completions.append(tool_info)
+
+                    if missing_completions:
+                        # Categorize missing tools by type for better diagnostics
+                        task_tools = [t for t in missing_completions if t.get("tool") == "Task"]
+                        builtin_tools = [t for t in missing_completions if t.get("tool") in ["TodoWrite", "Bash", "Read", "Write", "Edit", "Glob", "Grep"]]
+                        other_tools = [t for t in missing_completions if t not in task_tools and t not in builtin_tools]
+
+                        # Use warning level only if unexpected tools are missing
+                        # Task and builtin tools are expected to miss hooks sometimes
+                        log_level = "warning" if other_tools else "info"
+
+                        log_message = (
+                            f"Publishing fallback completion events for {len(missing_completions)} tools. "
+                            f"Task tools: {len(task_tools)} (expected - subagents execute in separate contexts), "
+                            f"Built-in tools: {len(builtin_tools)} (expected - may use optimized execution paths), "
+                            f"Other: {len(other_tools)} (unexpected - may indicate hook registration issue)"
+                        )
+
+                        getattr(self.logger, log_level)(
+                            "detected_missing_tool_completion_events",
+                            execution_id=context.execution_id,
+                            missing_count=len(missing_completions),
+                            task_tools_count=len(task_tools),
+                            builtin_tools_count=len(builtin_tools),
+                            other_tools_count=len(other_tools),
+                            missing_tool_names=[t.get("tool") for t in missing_completions],
+                            missing_tool_ids=[t.get("tool_use_id")[:12] for t in missing_completions],
+                            task_tool_ids=[t.get("tool_use_id")[:12] for t in task_tools] if task_tools else [],
+                            builtin_tool_ids=[t.get("tool_use_id")[:12] for t in builtin_tools] if builtin_tools else [],
+                            other_tool_ids=[t.get("tool_use_id")[:12] for t in other_tools] if other_tools else [],
+                            message=log_message,
+                            note=(
+                                "SubagentStop hook should reduce Task tool misses. "
+                                "Built-in tools may skip hooks by design. "
+                                "Investigate 'other' category if count > 0."
+                            )
+                        )
+
+                        # Publish missing completion events
+                        if event_callback:
+                            for tool_info in missing_completions:
+                                # Defensive: Validate tool_info has required fields
+                                if not isinstance(tool_info, dict):
+                                    self.logger.warning(
+                                        "invalid_tool_info_type_in_fallback",
+                                        tool_info_type=type(tool_info).__name__,
+                                        note="Skipping invalid tool_info"
+                                    )
+                                    continue
+
+                                tool_use_id = tool_info.get("tool_use_id")
+                                tool_name = tool_info.get("tool", "unknown")
+
+                                # Defensive: Skip if no tool_use_id
+                                if not tool_use_id:
+                                    self.logger.warning(
+                                        "missing_tool_use_id_in_fallback",
+                                        tool_name=tool_name,
+                                        note="Cannot publish completion without tool_use_id"
+                                    )
+                                    continue
+
+                                try:
+                                    # Defensive: Safe slicing for logging
+                                    tool_use_id_short = (
+                                        tool_use_id[:12] if isinstance(tool_use_id, str) and len(tool_use_id) >= 12
+                                        else str(tool_use_id)
+                                    )
+
+                                    event_callback(
+                                        {
+                                            "type": "tool_complete",
+                                            "tool_name": tool_name,
+                                            "tool_execution_id": tool_use_id,
+                                            "status": "success",  # Assume success if no error was caught
+                                            "output": None,  # No output available in fallback
+                                            "error": None,
+                                            "execution_id": context.execution_id,
+                                        }
+                                    )
+                                    completed_tools.add(tool_use_id)
+                                    self.logger.info(
+                                        "published_fallback_tool_completion",
+                                        tool_use_id=tool_use_id_short,
+                                        tool_name=tool_name,
+                                        note="Fallback completion event published successfully"
+                                    )
+                                except Exception as e:
+                                    # Non-fatal: Log but continue processing other tools
+                                    self.logger.error(
+                                        "failed_to_publish_fallback_completion",
+                                        tool_use_id=str(tool_use_id) if tool_use_id else "unknown",
+                                        tool_name=tool_name,
+                                        error=str(e),
+                                        error_type=type(e).__name__,
+                                        exc_info=True,
+                                        note="Continuing with remaining tools despite failure"
+                                    )
+
+                    self.logger.info(
+                        "claude_code_streaming_completed",
+                        execution_id=context.execution_id,
+                        finish_reason=message.subtype,
+                        chunk_count=chunk_count,
+                        message_count=message_count,
+                        response_length=len(accumulated_response),
+                        tool_count=len(tool_messages),
+                        completed_tool_count=len(completed_tools),
+                        missing_completions=len(missing_completions) if missing_completions else 0,
+                        usage=accumulated_usage,
+                        elapsed_time=f"{elapsed_time:.2f}s",
+                        has_session_id=bool(session_id),
+                    )
+
+                    # Final result message
+                    yield RuntimeExecutionResult(
+                        response="",  # Already streamed
+                        usage=accumulated_usage,
+                        success=message.subtype == "success",
+                        finish_reason=message.subtype,
+                        tool_execution_messages=tool_messages,  # Use standard field name for analytics
+                        tool_messages=tool_messages,  # Keep for backward compatibility
+                        model=context.model_id,
+                        metadata={
+                            "accumulated_response": accumulated_response,
+                            "elapsed_time": elapsed_time,
+                            "chunk_count": chunk_count,
+                            "message_count": message_count,
+                            "claude_code_session_id": session_id,
+                        },
+                    )
+                    break
+
+        except ImportError as e:
+            elapsed_time = asyncio.get_event_loop().time() - start_time
+            self.logger.error(
+                "claude_code_sdk_not_installed",
+                execution_id=context.execution_id,
+                error=str(e),
+                elapsed_time=f"{elapsed_time:.2f}s",
+            )
+            yield RuntimeExecutionResult(
+                response="",
+                usage={},
+                success=False,
+                error=f"Claude Code SDK not available: {str(e)}",
+                tool_messages=[],
+                tool_execution_messages=[],
+            )
+
+        except asyncio.TimeoutError:
+            elapsed_time = asyncio.get_event_loop().time() - start_time
+            self.logger.error(
+                "claude_code_streaming_timeout",
+                execution_id=context.execution_id,
+                elapsed_time=f"{elapsed_time:.2f}s",
+                chunks_before_timeout=chunk_count,
+            )
+            yield RuntimeExecutionResult(
+                response="",
+                usage={},
+                success=False,
+                error="Streaming execution timeout exceeded",
+                tool_messages=[],
+                tool_execution_messages=[],
+            )
+
+        except asyncio.CancelledError:
+            elapsed_time = asyncio.get_event_loop().time() - start_time
+            self.logger.warning(
+                "claude_code_streaming_cancelled_gracefully",
+                execution_id=context.execution_id,
+                elapsed_time=f"{elapsed_time:.2f}s",
+                chunks_before_cancellation=chunk_count,
+                accumulated_response_length=len(accumulated_response),
+                session_id=session_id[:16] if session_id else None,
+            )
+
+            # DURABILITY FIX: Do NOT re-raise! Handle cancellation gracefully
+            # Save partial state and allow workflow to resume from here
+            # The workflow is durable and should handle interruptions
+
+            # Yield partial success result with accumulated state
+            yield RuntimeExecutionResult(
+                response=accumulated_response,  # Return what we accumulated so far
+                usage=accumulated_usage,
+                success=True,  # Partial success, not a failure
+                finish_reason="cancelled",
+                tool_execution_messages=tool_messages,
+                tool_messages=tool_messages,
+                model=context.model_id,
+                metadata={
+                    "accumulated_response": accumulated_response,
+                    "elapsed_time": elapsed_time,
+                    "chunk_count": chunk_count,
+                    "message_count": message_count,
+                    "claude_code_session_id": session_id,
+                    "interrupted": True,  # Flag that this was interrupted
+                    "can_resume": bool(session_id),  # Can resume if we have session_id
+                },
+            )
+            # NOTE: Do NOT re-raise - this would break Temporal durability!
+
+        except Exception as e:
+            elapsed_time = asyncio.get_event_loop().time() - start_time
+            self.logger.error(
+                "claude_code_streaming_failed",
+                execution_id=context.execution_id,
+                error=str(e),
+                error_type=type(e).__name__,
+                elapsed_time=f"{elapsed_time:.2f}s",
+                chunks_before_error=chunk_count,
+                exc_info=True,
+            )
+            yield RuntimeExecutionResult(
+                response="",
+                usage={},
+                success=False,
+                error=f"{type(e).__name__}: {str(e)}",
+                finish_reason="error",  # CRITICAL: Must set finish_reason so caller recognizes this as final result
+                tool_messages=[],
+                tool_execution_messages=[],
+            )
+
+        finally:
+            # Clear execution context from proxy (with delay to allow in-flight SDK requests)
+            try:
+                clear_execution_context(
+                    context.execution_id,
+                    immediate=False,  # Use delayed cleanup
+                    delay_seconds=5.0  # Wait for in-flight SDK requests
+                )
+            except Exception as e:
+                self.logger.warning(
+                    "failed_to_clear_proxy_context_streaming",
+                    execution_id=context.execution_id,
+                    error=str(e),
+                )
+
+            # Restore MCP STDIO log level
+            try:
+                import logging
+                mcp_stdio_logger = logging.getLogger("mcp.client.stdio")
+                if 'original_stdio_level' in locals():
+                    mcp_stdio_logger.setLevel(original_stdio_level)
+            except Exception as log_level_error:
+                # Log but ignore errors restoring log level - this is non-critical cleanup
+                self.logger.debug(
+                    "failed_to_restore_log_level",
+                    error=str(log_level_error),
+                    execution_id=context.execution_id
+                )
+
+            # CRITICAL: Cleanup SDK client
+            if context.execution_id in self._active_clients:
+                client = self._active_clients.pop(context.execution_id)
+                cleanup_sdk_client(client, context.execution_id, self.logger)
+
+    async def cancel(self, execution_id: str) -> bool:
+        """
+        Cancel an in-progress execution via Claude SDK interrupt.
+
+        Args:
+            execution_id: ID of execution to cancel
+
+        Returns:
+            True if cancellation succeeded
+        """
+        if execution_id in self._active_clients:
+            try:
+                client = self._active_clients[execution_id]
+                await client.interrupt()
+                self.logger.info(
+                    "claude_code_execution_interrupted", execution_id=execution_id
+                )
+                return True
+            except Exception as e:
+                self.logger.error(
+                    "failed_to_interrupt_claude_code_execution",
+                    execution_id=execution_id,
+                    error=str(e),
+                )
+                return False
+        return False
+
+    # ==================== Custom Tool Extension API ====================
+
+    def get_custom_tool_requirements(self) -> Dict[str, Any]:
+        """
+        Get requirements for creating custom MCP servers for Claude Code runtime.
+
+        Returns:
+            Dictionary with format, examples, and documentation for MCP servers
+        """
+        return {
+            "format": "mcp_server",
+            "description": "MCP server created with @tool decorator and create_sdk_mcp_server()",
+            "example_code": '''
+from claude_agent_sdk import tool, create_sdk_mcp_server
+from typing import Any
+
+@tool("my_function", "Description of what this tool does", {"arg": str})
+async def my_function(args: dict[str, Any]) -> dict[str, Any]:
+    """Tool function implementation."""
+    return {
+        "content": [{
+            "type": "text",
+            "text": f"Result: {args['arg']}"
+        }]
+    }
+
+# Create MCP server
+mcp_server = create_sdk_mcp_server(
+    name="my_tools",
+    version="1.0.0",
+    tools=[my_function]
+)
+            ''',
+            "documentation_url": "https://docs.claude.ai/agent-sdk/custom-tools",
+            "required_attributes": ["name", "version"],
+            "schema": {
+                "type": "mcp_server",
+                "required": ["name", "version", "tools"]
+            }
+        }
+
+    def validate_custom_tool(self, mcp_server: Any) -> tuple[bool, Optional[str]]:
+        """
+        Validate an MCP server for Claude Code runtime.
+
+        Args:
+            mcp_server: MCP server instance to validate
+
+        Returns:
+            Tuple of (is_valid, error_message)
+        """
+        # Check required attributes
+        for attr in ['name', 'version']:
+            if not hasattr(mcp_server, attr):
+                return False, f"MCP server must have '{attr}' attribute"
+
+        # Validate name
+        if not isinstance(mcp_server.name, str) or not mcp_server.name:
+            return False, "MCP server name must be non-empty string"
+
+        # Check for tools (optional but recommended)
+        if hasattr(mcp_server, 'tools'):
+            if not mcp_server.tools:
+                self.logger.warning(
+                    "mcp_server_has_no_tools",
+                    server_name=mcp_server.name
+                )
+
+        return True, None
+
+    def register_custom_tool(self, mcp_server: Any, metadata: Optional[Dict] = None) -> str:
+        """
+        Register a custom MCP server with Claude Code runtime.
+
+        Args:
+            mcp_server: MCP server instance
+            metadata: Optional metadata (ignored, server name is used)
+
+        Returns:
+            Server name (identifier for this MCP server)
+
+        Raises:
+            ValueError: If MCP server validation fails or name conflicts
+        """
+        # Validate first
+        is_valid, error = self.validate_custom_tool(mcp_server)
+        if not is_valid:
+            raise ValueError(f"Invalid MCP server: {error}")
+
+        server_name = mcp_server.name
+
+        # Check for name conflicts
+        if server_name in self._custom_mcp_servers:
+            raise ValueError(f"MCP server '{server_name}' already registered")
+
+        # Store MCP server
+        self._custom_mcp_servers[server_name] = mcp_server
+
+        # Extract tool names for logging
+        tool_names = []
+        if hasattr(mcp_server, 'tools') and mcp_server.tools:
+            tool_names = [
+                f"mcp__{server_name}__{t.name}"
+                for t in mcp_server.tools
+                if hasattr(t, 'name')
+            ]
+
+        self.logger.info(
+            "custom_mcp_server_registered",
+            server_name=server_name,
+            tool_count=len(tool_names),
+            tools=tool_names
+        )
+
+        return server_name
+
+    def get_registered_custom_tools(self) -> list[str]:
+        """
+        Get list of registered custom MCP server names.
+
+        Returns:
+            List of server names
+        """
+        return list(self._custom_mcp_servers.keys())