"""Internal client implementation.""" import json import os from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator from dataclasses import asdict, replace from typing import Any from ..types import ( ClaudeAgentOptions, HookEvent, HookMatcher, Message, ) from .message_parser import parse_message from .query import Query from .session_resume import ( MaterializedResume, apply_materialized_options, build_mirror_batcher, materialize_resume_session, ) from .session_store_validation import validate_session_store_options from .transport import Transport from .transport.subprocess_cli import SubprocessCLITransport class InternalClient: """Internal client implementation.""" def __init__(self) -> None: """Initialize the internal client.""" def _convert_hooks_to_internal_format( self, hooks: dict[HookEvent, list[HookMatcher]] ) -> dict[str, list[dict[str, Any]]]: """Convert HookMatcher format to internal Query format.""" internal_hooks: dict[str, list[dict[str, Any]]] = {} for event, matchers in hooks.items(): internal_hooks[event] = [] for matcher in matchers: # Convert HookMatcher to internal dict format internal_matcher: dict[str, Any] = { "matcher": matcher.matcher if hasattr(matcher, "matcher") else None, "hooks": matcher.hooks if hasattr(matcher, "hooks") else [], } if hasattr(matcher, "timeout") and matcher.timeout is not None: internal_matcher["timeout"] = matcher.timeout internal_hooks[event].append(internal_matcher) return internal_hooks async def process_query( self, prompt: str | AsyncIterable[dict[str, Any]], options: ClaudeAgentOptions, transport: Transport | None = None, ) -> AsyncIterator[Message]: """Process a query through transport and Query.""" # Fail fast on invalid session_store option combinations before # spawning the subprocess. validate_session_store_options(options) # resume/continue + session_store: load the session from the store # into a temp CLAUDE_CONFIG_DIR for the subprocess to resume from. # Skipped when a custom transport was supplied — the materialized # options never reach a pre-constructed transport, so loading the # store and writing .credentials.json to a temp dir would be wasted. materialized = ( await materialize_resume_session(options) if transport is None else None ) inner = self._process_query_inner(prompt, options, transport, materialized) try: async for msg in inner: yield msg finally: # ``async for`` does NOT close its iterator when the loop body # raises (PEP 533 was deferred). Explicitly aclose the inner # generator first so its ``finally: await query.close()`` runs — # i.e. the subprocess is terminated — *before* we remove the temp # CLAUDE_CONFIG_DIR it is reading/writing. try: await inner.aclose() finally: # The temp dir holds a .credentials.json copy — remove it on # every exit path, including transport spawn failure before # the inner try/finally is reached. if materialized is not None: await materialized.cleanup() async def _process_query_inner( self, prompt: str | AsyncIterable[dict[str, Any]], options: ClaudeAgentOptions, transport: Transport | None, materialized: MaterializedResume | None, ) -> AsyncGenerator[Message, None]: # Validate and configure permission settings (matching TypeScript SDK logic) configured_options = options if options.can_use_tool: # canUseTool callback requires streaming mode (AsyncIterable prompt) if isinstance(prompt, str): raise ValueError( "can_use_tool callback requires streaming mode. " "Please provide prompt as an AsyncIterable instead of a string." ) # canUseTool and permission_prompt_tool_name are mutually exclusive if options.permission_prompt_tool_name: raise ValueError( "can_use_tool callback cannot be used with permission_prompt_tool_name. " "Please use one or the other." ) # Automatically set permission_prompt_tool_name to "stdio" for control protocol configured_options = replace(options, permission_prompt_tool_name="stdio") if materialized is not None: configured_options = apply_materialized_options( configured_options, materialized ) # Use provided transport or create subprocess transport if transport is not None: chosen_transport = transport else: chosen_transport = SubprocessCLITransport( prompt=prompt, options=configured_options, ) # Connect transport await chosen_transport.connect() # Extract SDK MCP servers from configured options sdk_mcp_servers = {} if configured_options.mcp_servers and isinstance( configured_options.mcp_servers, dict ): for name, config in configured_options.mcp_servers.items(): if isinstance(config, dict) and config.get("type") == "sdk": sdk_mcp_servers[name] = config["instance"] # type: ignore[typeddict-item] # Extract exclude_dynamic_sections from preset system prompt for the # initialize request (older CLIs ignore unknown initialize fields). exclude_dynamic_sections: bool | None = None sp = configured_options.system_prompt if isinstance(sp, dict) and sp.get("type") == "preset": eds = sp.get("exclude_dynamic_sections") if isinstance(eds, bool): exclude_dynamic_sections = eds # Convert agents to dict format for initialize request agents_dict = None if configured_options.agents: agents_dict = { name: {k: v for k, v in asdict(agent_def).items() if v is not None} for name, agent_def in configured_options.agents.items() } # Match ClaudeSDKClient.connect() — without this, query() ignores the env var initialize_timeout_ms = int( os.environ.get("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "60000") ) initialize_timeout = max(initialize_timeout_ms / 1000.0, 60.0) # Create Query to handle control protocol # Always use streaming mode internally (matching TypeScript SDK) # This ensures agents are always sent via initialize request query = Query( transport=chosen_transport, is_streaming_mode=True, # Always streaming internally can_use_tool=configured_options.can_use_tool, hooks=self._convert_hooks_to_internal_format(configured_options.hooks) if configured_options.hooks else None, sdk_mcp_servers=sdk_mcp_servers, initialize_timeout=initialize_timeout, agents=agents_dict, exclude_dynamic_sections=exclude_dynamic_sections, skills=configured_options.skills, ) if configured_options.session_store is not None: async def _on_mirror_error(key: Any, error: str) -> None: query.report_mirror_error(key, error) query.set_transcript_mirror_batcher( build_mirror_batcher( store=configured_options.session_store, materialized=materialized, env=configured_options.env, on_error=_on_mirror_error, flush_mode=configured_options.session_store_flush, ) ) try: # Start reading messages await query.start() # Always initialize to send agents via stdin (matching TypeScript SDK) await query.initialize() # Handle prompt input if isinstance(prompt, str): # For string prompts, write user message to stdin after initialize # (matching TypeScript SDK behavior) user_message = { "type": "user", "session_id": "", "message": {"role": "user", "content": prompt}, "parent_tool_use_id": None, } await chosen_transport.write(json.dumps(user_message) + "\n") query.spawn_task(query.wait_for_result_and_end_input()) elif isinstance(prompt, AsyncIterable): # Stream input in background for async iterables query.spawn_task(query.stream_input(prompt)) # Yield parsed messages, skipping unknown message types async for data in query.receive_messages(): message = parse_message(data) if message is not None: yield message finally: await query.close() query.close_receive_stream()