ripperdoc 0.2.0__py3-none-any.whl → 0.2.3__py3-none-any.whl

ripperdoc/core/query.py

@@ -6,86 +6,273 @@ the query-response loop including tool execution.
 
 import asyncio
 import inspect
-from typing import AsyncGenerator, List, Optional, Dict, Any, Union, Iterable, Tuple
-from anthropic import AsyncAnthropic
-from openai import AsyncOpenAI
-
-from ripperdoc.core.tool import (
-    Tool,
-    ToolUseContext,
-    ToolResult,
-    ToolProgress,
-    build_tool_description,
-    tool_input_examples,
+import os
+import time
+from asyncio import CancelledError
+from typing import (
+    Any,
+    AsyncGenerator,
+    Awaitable,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+    Union,
+    cast,
+)
+
+from pydantic import ValidationError
+
+from ripperdoc.core.config import provider_protocol
+from ripperdoc.core.providers import ProviderClient, get_provider_client
+from ripperdoc.core.permissions import PermissionResult
+from ripperdoc.core.query_utils import (
+    build_full_system_prompt,
+    determine_tool_mode,
+    extract_tool_use_blocks,
+    format_pydantic_errors,
+    log_openai_messages,
+    resolve_model_profile,
+    text_mode_history,
+    tool_result_message,
 )
+from ripperdoc.core.tool import Tool, ToolProgress, ToolResult, ToolUseContext
+from ripperdoc.utils.file_watch import ChangedFileNotice, FileSnapshot, detect_changed_files
 from ripperdoc.utils.log import get_logger
 from ripperdoc.utils.messages import (
-    MessageContent,
-    UserMessage,
     AssistantMessage,
+    MessageContent,
     ProgressMessage,
-    create_user_message,
+    UserMessage,
     create_assistant_message,
+    create_user_message,
     create_progress_message,
     normalize_messages_for_api,
     INTERRUPT_MESSAGE,
     INTERRUPT_MESSAGE_FOR_TOOL_USE,
 )
-from ripperdoc.core.permissions import PermissionResult
-from ripperdoc.core.config import get_global_config, ProviderType, provider_protocol
-from ripperdoc.utils.session_usage import record_usage
-
-import time
 
 
 logger = get_logger()
 
+DEFAULT_REQUEST_TIMEOUT_SEC = float(os.getenv("RIPPERDOC_API_TIMEOUT", "120"))
+MAX_LLM_RETRIES = 1
+
 
-def _safe_int(value: Any) -> int:
-    """Best-effort int conversion for usage counters."""
+def _resolve_tool(
+    tool_registry: "ToolRegistry", tool_name: str, tool_use_id: str
+) -> tuple[Optional[Tool[Any, Any]], Optional[UserMessage]]:
+    """Find a tool by name and return an error message if missing."""
+    tool = tool_registry.get(tool_name)
+    if tool:
+        tool_registry.activate_tools([tool_name])
+        return tool, None
+    return None, tool_result_message(
+        tool_use_id, f"Error: Tool '{tool_name}' not found", is_error=True
+    )
+
+
+ToolPermissionCallable = Callable[
+    [Tool[Any, Any], Any],
+    Union[
+        PermissionResult,
+        Dict[str, Any],
+        Tuple[bool, Optional[str]],
+        bool,
+        Awaitable[Union[PermissionResult, Dict[str, Any], Tuple[bool, Optional[str]], bool]],
+    ],
+]
+
+
+async def _check_tool_permissions(
+    tool: Tool[Any, Any],
+    parsed_input: Any,
+    query_context: "QueryContext",
+    can_use_tool_fn: Optional[ToolPermissionCallable],
+) -> tuple[bool, Optional[str]]:
+    """Evaluate whether a tool call is allowed."""
     try:
-        if value is None:
-            return 0
-        return int(value)
-    except (TypeError, ValueError):
-        return 0
-
-
-def _get_usage_field(usage: Any, field: str) -> int:
-    """Fetch a usage field from either a dict or object."""
-    if usage is None:
-        return 0
-    if isinstance(usage, dict):
-        return _safe_int(usage.get(field))
-    return _safe_int(getattr(usage, field, 0))
-
-
-def _anthropic_usage_tokens(usage: Any) -> Dict[str, int]:
-    """Extract token counts from an Anthropic response usage payload."""
-    return {
-        "input_tokens": _get_usage_field(usage, "input_tokens"),
-        "output_tokens": _get_usage_field(usage, "output_tokens"),
-        "cache_read_input_tokens": _get_usage_field(usage, "cache_read_input_tokens"),
-        "cache_creation_input_tokens": _get_usage_field(usage, "cache_creation_input_tokens"),
-    }
-
-
-def _openai_usage_tokens(usage: Any) -> Dict[str, int]:
-    """Extract token counts from an OpenAI-compatible response usage payload."""
-    prompt_details = None
-    if isinstance(usage, dict):
-        prompt_details = usage.get("prompt_tokens_details")
-    else:
-        prompt_details = getattr(usage, "prompt_tokens_details", None)
+        if can_use_tool_fn is not None:
+            decision = can_use_tool_fn(tool, parsed_input)
+            if inspect.isawaitable(decision):
+                decision = await decision
+            if isinstance(decision, PermissionResult):
+                return decision.result, decision.message
+            if isinstance(decision, dict) and "result" in decision:
+                return bool(decision.get("result")), decision.get("message")
+            if isinstance(decision, tuple) and len(decision) == 2:
+                return bool(decision[0]), decision[1]
+            return bool(decision), None
+
+        if query_context.safe_mode and tool.needs_permissions(parsed_input):
+            loop = asyncio.get_running_loop()
+            input_preview = (
+                parsed_input.model_dump()
+                if hasattr(parsed_input, "model_dump")
+                else str(parsed_input)
+            )
+            prompt = f"Allow tool '{tool.name}' with input {input_preview}? [y/N]: "
+            response = await loop.run_in_executor(None, lambda: input(prompt))
+            return response.strip().lower() in ("y", "yes"), None
+
+        return True, None
+    except Exception:
+        logger.exception(
+            f"Error checking permissions for tool '{tool.name}'",
+            extra={"tool": getattr(tool, "name", None)},
+        )
+        return False, None
+
+
+def _format_changed_file_notice(notices: List[ChangedFileNotice]) -> str:
+    """Render a system notice about files that changed on disk."""
+    lines: List[str] = [
+        "System notice: Files you previously read have changed on disk.",
+        "Please re-read the affected files before making further edits.",
+        "",
+    ]
+    for notice in notices:
+        lines.append(f"- {notice.file_path}")
+        summary = (notice.summary or "").rstrip()
+        if summary:
+            indented = "\n".join(f"    {line}" for line in summary.splitlines())
+            lines.append(indented)
+    return "\n".join(lines)
+
+
+async def _run_tool_use_generator(
+    tool: Tool[Any, Any],
+    tool_use_id: str,
+    tool_name: str,
+    parsed_input: Any,
+    sibling_ids: set[str],
+    tool_context: ToolUseContext,
+) -> AsyncGenerator[Union[UserMessage, ProgressMessage], None]:
+    """Execute a single tool_use and yield progress/results."""
+    try:
+        async for output in tool.call(parsed_input, tool_context):
+            if isinstance(output, ToolProgress):
+                yield create_progress_message(
+                    tool_use_id=tool_use_id,
+                    sibling_tool_use_ids=sibling_ids,
+                    content=output.content,
+                )
+                logger.debug(f"[query] Progress from tool_use_id={tool_use_id}: {output.content}")
+            elif isinstance(output, ToolResult):
+                result_content = output.result_for_assistant or str(output.data)
+                result_msg = tool_result_message(
+                    tool_use_id, result_content, tool_use_result=output.data
+                )
+                yield result_msg
+                logger.debug(
+                    f"[query] Tool completed tool_use_id={tool_use_id} name={tool_name} "
+                    f"result_len={len(result_content)}"
+                )
+    except Exception as exc:
+        logger.exception(
+            f"Error executing tool '{tool_name}'",
+            extra={"tool": tool_name, "tool_use_id": tool_use_id},
+        )
+        yield tool_result_message(tool_use_id, f"Error executing tool: {str(exc)}", is_error=True)
+
+
+def _group_tool_calls_by_concurrency(prepared_calls: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Group consecutive tool calls by their concurrency safety."""
+    groups: List[Dict[str, Any]] = []
+    for call in prepared_calls:
+        is_safe = bool(call.get("is_concurrency_safe"))
+        if groups and groups[-1]["is_concurrency_safe"] == is_safe:
+            groups[-1]["items"].append(call)
+        else:
+            groups.append({"is_concurrency_safe": is_safe, "items": [call]})
+    return groups
 
-    cache_read_tokens = _get_usage_field(prompt_details, "cached_tokens") if prompt_details else 0
 
-    return {
-        "input_tokens": _get_usage_field(usage, "prompt_tokens"),
-        "output_tokens": _get_usage_field(usage, "completion_tokens"),
-        "cache_read_input_tokens": cache_read_tokens,
-        "cache_creation_input_tokens": 0,
-    }
+async def _execute_tools_sequentially(
+    items: List[Dict[str, Any]], tool_results: List[UserMessage]
+) -> AsyncGenerator[Union[UserMessage, ProgressMessage], None]:
+    """Run tool generators one by one."""
+    for item in items:
+        gen = item.get("generator")
+        if not gen:
+            continue
+        async for message in gen:
+            if isinstance(message, UserMessage):
+                tool_results.append(message)
+            yield message
+
+
+async def _execute_tools_in_parallel(
+    items: List[Dict[str, Any]], tool_results: List[UserMessage]
+) -> AsyncGenerator[Union[UserMessage, ProgressMessage], None]:
+    """Run tool generators concurrently."""
+    generators = [call["generator"] for call in items if call.get("generator")]
+    async for message in _run_concurrent_tool_uses(generators, tool_results):
+        yield message
+
+
+async def _run_tools_concurrently(
+    prepared_calls: List[Dict[str, Any]], tool_results: List[UserMessage]
+) -> AsyncGenerator[Union[UserMessage, ProgressMessage], None]:
+    """Run tools grouped by concurrency safety (parallel for safe groups)."""
+    for group in _group_tool_calls_by_concurrency(prepared_calls):
+        if group["is_concurrency_safe"]:
+            logger.debug(
+                f"[query] Executing {len(group['items'])} concurrency-safe tool(s) in parallel"
+            )
+            async for message in _execute_tools_in_parallel(group["items"], tool_results):
+                yield message
+        else:
+            logger.debug(
+                f"[query] Executing {len(group['items'])} tool(s) sequentially (not concurrency safe)"
+            )
+            async for message in _run_tools_serially(group["items"], tool_results):
+                yield message
+
+
+async def _run_tools_serially(
+    prepared_calls: List[Dict[str, Any]], tool_results: List[UserMessage]
+) -> AsyncGenerator[Union[UserMessage, ProgressMessage], None]:
+    """Run all tools sequentially (helper for clarity)."""
+    async for message in _execute_tools_sequentially(prepared_calls, tool_results):
+        yield message
+
+
+async def _run_concurrent_tool_uses(
+    generators: List[AsyncGenerator[Union[UserMessage, ProgressMessage], None]],
+    tool_results: List[UserMessage],
+) -> AsyncGenerator[Union[UserMessage, ProgressMessage], None]:
+    """Drain multiple tool generators concurrently and stream outputs."""
+    if not generators:
+        return
+
+    queue: asyncio.Queue[Optional[Union[UserMessage, ProgressMessage]]] = asyncio.Queue()
+
+    async def _consume(gen: AsyncGenerator[Union[UserMessage, ProgressMessage], None]) -> None:
+        try:
+            async for message in gen:
+                await queue.put(message)
+        except Exception:
+            logger.exception("[query] Unexpected error while consuming tool generator")
+        finally:
+            await queue.put(None)
+
+    tasks = [asyncio.create_task(_consume(gen)) for gen in generators]
+    active = len(tasks)
+
+    try:
+        while active:
+            message = await queue.get()
+            if message is None:
+                active -= 1
+                continue
+            if isinstance(message, UserMessage):
+                tool_results.append(message)
+            yield message
+    finally:
+        await asyncio.gather(*tasks, return_exceptions=True)
 
 
 class ToolRegistry:
@@ -118,6 +305,10 @@ class ToolRegistry:
             try:
                 deferred = tool.defer_loading()
             except Exception:
+                logger.exception(
+                    "[tool_registry] Tool.defer_loading failed",
+                    extra={"tool": getattr(tool, "name", None)},
+                )
                 deferred = False
             if deferred:
                 self._deferred.add(name)
@@ -193,6 +384,7 @@ class QueryContext:
         self.model = model
         self.verbose = verbose
         self.abort_controller = asyncio.Event()
+        self.file_state_cache: Dict[str, FileSnapshot] = {}
 
     @property
     def tools(self) -> List[Tool[Any, Any]]:
@@ -220,6 +412,11 @@ async def query_llm(
     max_thinking_tokens: int = 0,
     model: str = "main",
     abort_signal: Optional[asyncio.Event] = None,
+    *,
+    progress_callback: Optional[Callable[[str], Awaitable[None]]] = None,
+    request_timeout: Optional[float] = None,
+    max_retries: int = MAX_LLM_RETRIES,
+    stream: bool = True,
 ) -> AssistantMessage:
     """Query the AI model and return the response.
 
@@ -230,47 +427,48 @@ async def query_llm(
         max_thinking_tokens: Maximum tokens for thinking (0 = disabled)
         model: Model pointer to use
         abort_signal: Event to signal abortion
+        progress_callback: Optional async callback invoked with streamed text chunks
+        request_timeout: Max seconds to wait for a provider response before retrying
+        max_retries: Number of retries on timeout/errors (total attempts = retries + 1)
+        stream: Enable streaming for providers that support it (text-only mode)
 
     Returns:
         AssistantMessage with the model's response
     """
-    config = get_global_config()
-
-    # Get the model profile
-    profile_name = getattr(config.model_pointers, model, None)
-    if profile_name is None:
-        profile_name = model
-
-    model_profile = config.model_profiles.get(profile_name)
-    if model_profile is None:
-        fallback_profile = getattr(config.model_pointers, "main", "default")
-        model_profile = config.model_profiles.get(fallback_profile)
-
-    if not model_profile:
-        raise ValueError(f"No model profile found for pointer: {model}")
+    request_timeout = request_timeout or DEFAULT_REQUEST_TIMEOUT_SEC
+    request_timeout = request_timeout or DEFAULT_REQUEST_TIMEOUT_SEC
+    model_profile = resolve_model_profile(model)
 
     # Normalize messages based on protocol family (Anthropic allows tool blocks; OpenAI-style prefers text-only)
     protocol = provider_protocol(model_profile.provider)
-    normalized_messages = normalize_messages_for_api(
-        messages,
-        protocol=protocol,
+    tool_mode = determine_tool_mode(model_profile)
+    messages_for_model: List[Union[UserMessage, AssistantMessage, ProgressMessage]]
+    if tool_mode == "text":
+        messages_for_model = cast(
+            List[Union[UserMessage, AssistantMessage, ProgressMessage]],
+            text_mode_history(messages),
+        )
+    else:
+        messages_for_model = messages
+
+    normalized_messages: List[Dict[str, Any]] = normalize_messages_for_api(
+        messages_for_model, protocol=protocol, tool_mode=tool_mode
+    )
+    logger.info(
+        "[query_llm] Preparing model request",
+        extra={
+            "model_pointer": model,
+            "provider": getattr(model_profile.provider, "value", str(model_profile.provider)),
+            "model": model_profile.model,
+            "normalized_messages": len(normalized_messages),
+            "tool_count": len(tools),
+            "max_thinking_tokens": max_thinking_tokens,
+            "tool_mode": tool_mode,
+        },
     )
 
     if protocol == "openai":
-        summary_parts = []
-        for idx, m in enumerate(normalized_messages):
-            role = m.get("role")
-            tool_calls = m.get("tool_calls")
-            tc_ids = []
-            if tool_calls:
-                tc_ids = [tc.get("id") for tc in tool_calls]
-            tool_call_id = m.get("tool_call_id")
-            summary_parts.append(
-                f"{idx}:{role}"
-                + (f" tool_calls={tc_ids}" if tc_ids else "")
-                + (f" tool_call_id={tool_call_id}" if tool_call_id else "")
-            )
-        logger.debug(f"[query_llm] OpenAI normalized messages: {' | '.join(summary_parts)}")
+        log_openai_messages(normalized_messages)
 
     logger.debug(
         f"[query_llm] Sending {len(normalized_messages)} messages to model pointer "
@@ -282,136 +480,49 @@ async def query_llm(
     start_time = time.time()
 
     try:
-        # Create the appropriate client based on provider
-        if model_profile.provider == ProviderType.ANTHROPIC:
-            async with AsyncAnthropic(api_key=model_profile.api_key) as client:
-                # Build tool schemas
-                tool_schemas = []
-                for tool in tools:
-                    description = await build_tool_description(
-                        tool, include_examples=True, max_examples=2
-                    )
-                    tool_schema = {
-                        "name": tool.name,
-                        "description": description,
-                        "input_schema": tool.input_schema.model_json_schema(),
-                        "defer_loading": bool(getattr(tool, "defer_loading", lambda: False)()),
-                    }
-                    examples = tool_input_examples(tool, limit=5)
-                    if examples:
-                        tool_schema["input_examples"] = examples
-                    tool_schemas.append(tool_schema)
-
-                # Make the API call
-                response = await client.messages.create(
-                    model=model_profile.model,
-                    max_tokens=model_profile.max_tokens,
-                    system=system_prompt,
-                    messages=normalized_messages,  # type: ignore[arg-type]
-                    tools=tool_schemas if tool_schemas else None,  # type: ignore
-                    temperature=model_profile.temperature,
-                )
-
-                duration_ms = (time.time() - start_time) * 1000
-
-                usage_tokens = _anthropic_usage_tokens(getattr(response, "usage", None))
-                record_usage(model_profile.model, duration_ms=duration_ms, **usage_tokens)
-
-                # Calculate cost (simplified, should use actual pricing)
-                cost_usd = 0.0  # TODO: Implement cost calculation
-
-                # Convert response to our format
-                content_blocks = []
-                for block in response.content:
-                    if block.type == "text":
-                        content_blocks.append({"type": "text", "text": block.text})
-                    elif block.type == "tool_use":
-                        content_blocks.append(
-                            {
-                                "type": "tool_use",
-                                "tool_use_id": block.id,
-                                "name": block.name,
-                                "input": block.input,  # type: ignore[dict-item]
-                            }
-                        )
-
-                return create_assistant_message(
-                    content=content_blocks, cost_usd=cost_usd, duration_ms=duration_ms
-                )
-
-        elif model_profile.provider == ProviderType.OPENAI_COMPATIBLE:
-            # OpenAI-compatible APIs (OpenAI, DeepSeek, Mistral, etc.)
-            async with AsyncOpenAI(
-                api_key=model_profile.api_key, base_url=model_profile.api_base
-            ) as client:
-                # Build tool schemas for OpenAI format
-                openai_tools = []
-                for tool in tools:
-                    description = await build_tool_description(
-                        tool, include_examples=True, max_examples=2
-                    )
-                    openai_tools.append(
-                        {
-                            "type": "function",
-                            "function": {
-                                "name": tool.name,
-                                "description": description,
-                                "parameters": tool.input_schema.model_json_schema(),
-                            },
-                        }
-                    )
-
-                # Prepare messages for OpenAI format
-                openai_messages = [
-                    {"role": "system", "content": system_prompt}
-                ] + normalized_messages
-
-                # Make the API call
-                openai_response: Any = await client.chat.completions.create(
-                    model=model_profile.model,
-                    messages=openai_messages,
-                    tools=openai_tools if openai_tools else None,  # type: ignore[arg-type]
-                    temperature=model_profile.temperature,
-                    max_tokens=model_profile.max_tokens,
-                )
-
-                duration_ms = (time.time() - start_time) * 1000
-                usage_tokens = _openai_usage_tokens(getattr(openai_response, "usage", None))
-                record_usage(model_profile.model, duration_ms=duration_ms, **usage_tokens)
-                cost_usd = 0.0  # TODO: Implement cost calculation
-
-                # Convert OpenAI response to our format
-                content_blocks = []
-                choice = openai_response.choices[0]
-
-                if choice.message.content:
-                    content_blocks.append({"type": "text", "text": choice.message.content})
-
-                if choice.message.tool_calls:
-                    for tool_call in choice.message.tool_calls:
-                        import json
-
-                        content_blocks.append(
-                            {
-                                "type": "tool_use",
-                                "tool_use_id": tool_call.id,
-                                "name": tool_call.function.name,
-                                "input": json.loads(tool_call.function.arguments),
-                            }
-                        )
-
-                return create_assistant_message(
-                    content=content_blocks, cost_usd=cost_usd, duration_ms=duration_ms
-                )
+        client: Optional[ProviderClient] = get_provider_client(model_profile.provider)
+        if client is None:
+            duration_ms = (time.time() - start_time) * 1000
+            error_msg = create_assistant_message(
+                content=(
+                    "Gemini protocol is not supported yet in Ripperdoc. "
+                    "Please configure an Anthropic or OpenAI-compatible model."
+                ),
+                duration_ms=duration_ms,
+            )
+            error_msg.is_api_error_message = True
+            return error_msg
+
+        provider_response = await client.call(
+            model_profile=model_profile,
+            system_prompt=system_prompt,
+            normalized_messages=normalized_messages,
+            tools=tools,
+            tool_mode=tool_mode,
+            stream=stream,
+            progress_callback=progress_callback,
+            request_timeout=request_timeout,
+            max_retries=max_retries,
+        )
 
-        elif model_profile.provider == ProviderType.GEMINI:
-            raise NotImplementedError("Gemini protocol is not yet supported.")
-        else:
-            raise NotImplementedError(f"Provider {model_profile.provider} not yet implemented")
+        return create_assistant_message(
+            content=provider_response.content_blocks,
+            cost_usd=provider_response.cost_usd,
+            duration_ms=provider_response.duration_ms,
+        )
 
     except Exception as e:
         # Return error message
-        logger.error(f"Error querying AI model: {e}")
+        logger.exception(
+            "Error querying AI model",
+            extra={
+                "model": getattr(model_profile, "model", None),
+                "model_pointer": model,
+                "provider": (
+                    getattr(model_profile.provider, "value", None) if model_profile else None
+                ),
+            },
+        )
         duration_ms = (time.time() - start_time) * 1000
         error_msg = create_assistant_message(
             content=f"Error querying AI model: {str(e)}", duration_ms=duration_ms
@@ -425,7 +536,7 @@ async def query(
     system_prompt: str,
     context: Dict[str, str],
     query_context: QueryContext,
-    can_use_tool_fn: Optional[Any] = None,
+    can_use_tool_fn: Optional[ToolPermissionCallable] = None,
 ) -> AsyncGenerator[Union[UserMessage, AssistantMessage, ProgressMessage], None]:
     """Execute a query with tool support.
 
@@ -445,59 +556,105 @@ async def query(
     Yields:
         Messages (user, assistant, progress) as they are generated
     """
+    logger.info(
+        "[query] Starting query loop",
+        extra={
+            "message_count": len(messages),
+            "tool_count": len(query_context.tools),
+            "safe_mode": query_context.safe_mode,
+            "model_pointer": query_context.model,
+        },
+    )
     # Work on a copy so external mutations (e.g., UI appending messages while consuming)
     # do not interfere with recursion or normalization.
     messages = list(messages)
+    change_notices = detect_changed_files(query_context.file_state_cache)
+    if change_notices:
+        messages.append(create_user_message(_format_changed_file_notice(change_notices)))
+    model_profile = resolve_model_profile(query_context.model)
+    tool_mode = determine_tool_mode(model_profile)
+    tools_for_model: List[Tool[Any, Any]] = [] if tool_mode == "text" else query_context.all_tools()
+
+    full_system_prompt = build_full_system_prompt(
+        system_prompt, context, tool_mode, query_context.all_tools()
+    )
+    logger.debug(
+        "[query] Built system prompt",
+        extra={
+            "prompt_chars": len(full_system_prompt),
+            "context_entries": len(context),
+            "tool_count": len(tools_for_model),
+        },
+    )
+
+    progress_queue: asyncio.Queue[Optional[ProgressMessage]] = asyncio.Queue()
 
-    async def _check_permissions(
-        tool: Tool[Any, Any], parsed_input: Any
-    ) -> tuple[bool, Optional[str]]:
-        """Check permissions for tool execution."""
+    async def _stream_progress(chunk: str) -> None:
+        if not chunk:
+            return
         try:
-            if can_use_tool_fn is not None:
-                decision = can_use_tool_fn(tool, parsed_input)
-                if inspect.isawaitable(decision):
-                    decision = await decision
-                if isinstance(decision, PermissionResult):
-                    return decision.result, decision.message
-                if isinstance(decision, dict) and "result" in decision:
-                    return bool(decision.get("result")), decision.get("message")
-                if isinstance(decision, tuple) and len(decision) == 2:
-                    return bool(decision[0]), decision[1]
-                return bool(decision), None
-
-            if query_context.safe_mode and tool.needs_permissions(parsed_input):
-                loop = asyncio.get_running_loop()
-                input_preview = (
-                    parsed_input.model_dump()
-                    if hasattr(parsed_input, "model_dump")
-                    else str(parsed_input)
+            await progress_queue.put(
+                create_progress_message(
+                    tool_use_id="stream",
+                    sibling_tool_use_ids=set(),
+                    content=chunk,
                 )
-                prompt = f"Allow tool '{tool.name}' with input {input_preview}? [y/N]: "
-                response = await loop.run_in_executor(None, lambda: input(prompt))
-                return response.strip().lower() in ("y", "yes"), None
-
-            return True, None
-        except Exception as exc:
-            # Fail closed on any errors
-            logger.error(f"Error checking permissions for tool '{tool.name}': {exc}")
-            return False, None
-
-    # Build full system prompt with context
-    full_system_prompt = system_prompt
-    if context:
-        context_str = "\n".join(f"{k}: {v}" for k, v in context.items())
-        full_system_prompt = f"{system_prompt}\n\nContext:\n{context_str}"
-
-    assistant_message = await query_llm(
-        messages,
-        full_system_prompt,
-        query_context.all_tools(),
-        query_context.max_thinking_tokens,
-        query_context.model,
-        query_context.abort_controller,
+            )
+        except Exception:
+            logger.exception("[query] Failed to enqueue stream progress chunk")
+
+    assistant_task = asyncio.create_task(
+        query_llm(
+            messages,
+            full_system_prompt,
+            tools_for_model,
+            query_context.max_thinking_tokens,
+            query_context.model,
+            query_context.abort_controller,
+            progress_callback=_stream_progress,
+            request_timeout=DEFAULT_REQUEST_TIMEOUT_SEC,
+            max_retries=MAX_LLM_RETRIES,
+            stream=True,
+        )
     )
 
+    assistant_message: Optional[AssistantMessage] = None
+
+    while True:
+        if query_context.abort_controller.is_set():
+            assistant_task.cancel()
+            try:
+                await assistant_task
+            except CancelledError:
+                pass
+            yield create_assistant_message(INTERRUPT_MESSAGE)
+            return
+        if assistant_task.done():
+            assistant_message = await assistant_task
+            break
+        try:
+            progress = progress_queue.get_nowait()
+        except asyncio.QueueEmpty:
+            waiter = asyncio.create_task(progress_queue.get())
+            done, pending = await asyncio.wait(
+                {assistant_task, waiter}, return_when=asyncio.FIRST_COMPLETED
+            )
+            if assistant_task in done:
+                for task in pending:
+                    task.cancel()
+                assistant_message = await assistant_task
+                break
+            progress = waiter.result()
+        if progress:
+            yield progress
+
+    while not progress_queue.empty():
+        residual = progress_queue.get_nowait()
+        if residual:
+            yield residual
+
+    assert assistant_message is not None
+
     # Check for abort
     if query_context.abort_controller.is_set():
         yield create_assistant_message(INTERRUPT_MESSAGE)
@@ -505,175 +662,142 @@ async def query(
 
     yield assistant_message
 
-    tool_block_count = 0
-    if isinstance(assistant_message.message.content, list):
-        tool_block_count = sum(
-            1
-            for block in assistant_message.message.content
-            if hasattr(block, "type") and block.type == "tool_use"
-        )
+    tool_use_blocks: List[MessageContent] = extract_tool_use_blocks(assistant_message)
+    text_blocks = (
+        len(assistant_message.message.content)
+        if isinstance(assistant_message.message.content, list)
+        else 1
+    )
     logger.debug(
-        f"[query] Assistant message received: "
-        f"text_blocks={len(assistant_message.message.content) if isinstance(assistant_message.message.content, list) else 1}, "
-        f"tool_use_blocks={tool_block_count}"
+        f"[query] Assistant message received: text_blocks={text_blocks}, "
+        f"tool_use_blocks={len(tool_use_blocks)}"
     )
 
-    # Check for tool use
-    tool_use_blocks = []
-    if isinstance(assistant_message.message.content, list):
-        for block in assistant_message.message.content:
-            normalized_block = MessageContent(**block) if isinstance(block, dict) else block
-            if hasattr(normalized_block, "type") and normalized_block.type == "tool_use":
-                tool_use_blocks.append(normalized_block)
-
-    # If no tool use, we're done
     if not tool_use_blocks:
         logger.debug("[query] No tool_use blocks; returning response to user.")
         return
 
-    # Execute tools
-    tool_results: List[UserMessage] = []
-
     logger.debug(f"[query] Executing {len(tool_use_blocks)} tool_use block(s).")
+    tool_results: List[UserMessage] = []
+    permission_denied = False
+    sibling_ids = set(
+        getattr(t, "tool_use_id", None) or getattr(t, "id", None) or "" for t in tool_use_blocks
+    )
+    prepared_calls: List[Dict[str, Any]] = []
 
     for tool_use in tool_use_blocks:
         tool_name = tool_use.name
         if not tool_name:
             continue
-        tool_id = getattr(tool_use, "tool_use_id", None) or getattr(tool_use, "id", None) or ""
+        tool_use_id = getattr(tool_use, "tool_use_id", None) or getattr(tool_use, "id", None) or ""
         tool_input = getattr(tool_use, "input", {}) or {}
 
-        # Find the tool
-        tool = query_context.tool_registry.get(tool_name)
-        # Auto-activate when used so subsequent rounds list it as active.
-        if tool:
-            query_context.activate_tools([tool_name])  # type: ignore[list-item]
-
-        if not tool:
-            # Tool not found
-            logger.warning(f"[query] Tool '{tool_name}' not found for tool_use_id={tool_id}")
-            result_msg = create_user_message(
-                [
-                    {
-                        "type": "tool_result",
-                        "tool_use_id": tool_id,
-                        "text": f"Error: Tool '{tool_name}' not found",
-                        "is_error": True,
-                    }
-                ]
-            )
-            tool_results.append(result_msg)
-            yield result_msg
+        tool, missing_msg = _resolve_tool(query_context.tool_registry, tool_name, tool_use_id)
+        if missing_msg:
+            logger.warning(f"[query] Tool '{tool_name}' not found for tool_use_id={tool_use_id}")
+            tool_results.append(missing_msg)
+            yield missing_msg
             continue
-
-        # Execute the tool
-        tool_context = ToolUseContext(
-            safe_mode=query_context.safe_mode,
-            verbose=query_context.verbose,
-            permission_checker=can_use_tool_fn,
-            tool_registry=query_context.tool_registry,
-        )
+        assert tool is not None
 
         try:
-            # Parse input using tool's schema
             parsed_input = tool.input_schema(**tool_input)
             logger.debug(
-                f"[query] tool_use_id={tool_id} name={tool_name} parsed_input="
+                f"[query] tool_use_id={tool_use_id} name={tool_name} parsed_input="
                 f"{str(parsed_input)[:500]}"
             )
 
-            # Validate input before execution
+            tool_context = ToolUseContext(
+                safe_mode=query_context.safe_mode,
+                verbose=query_context.verbose,
+                permission_checker=can_use_tool_fn,
+                tool_registry=query_context.tool_registry,
+                file_state_cache=query_context.file_state_cache,
+                abort_signal=query_context.abort_controller,
+            )
+
             validation = await tool.validate_input(parsed_input, tool_context)
             if not validation.result:
                 logger.debug(
-                    f"[query] Validation failed for tool_use_id={tool_id}: {validation.message}"
+                    f"[query] Validation failed for tool_use_id={tool_use_id}: {validation.message}"
                 )
-                result_msg = create_user_message(
-                    [
-                        {
-                            "type": "tool_result",
-                            "tool_use_id": tool_id,
-                            "text": validation.message or "Tool input validation failed.",
-                            "is_error": True,
-                        }
-                    ]
+                result_msg = tool_result_message(
+                    tool_use_id,
+                    validation.message or "Tool input validation failed.",
+                    is_error=True,
                 )
                 tool_results.append(result_msg)
                 yield result_msg
                 continue
 
-            # Permission check (safe mode or custom checker)
             if query_context.safe_mode or can_use_tool_fn is not None:
-                allowed, denial_message = await _check_permissions(tool, parsed_input)
+                allowed, denial_message = await _check_tool_permissions(
+                    tool, parsed_input, query_context, can_use_tool_fn
+                )
                 if not allowed:
                     logger.debug(
-                        f"[query] Permission denied for tool_use_id={tool_id}: {denial_message}"
-                    )
-                    denial_text = denial_message or f"Permission denied for tool '{tool_name}'."
-                    result_msg = create_user_message(
-                        [
-                            {
-                                "type": "tool_result",
-                                "tool_use_id": tool_id,
-                                "text": denial_text,
-                                "is_error": True,
-                            }
-                        ]
-                    )
-                    tool_results.append(result_msg)
-                    yield result_msg
-                    continue
-
-            # Execute tool
-            async for output in tool.call(parsed_input, tool_context):
-                if isinstance(output, ToolProgress):
-                    # Yield progress
-                    progress = create_progress_message(
-                        tool_use_id=tool_id,
-                        sibling_tool_use_ids=set(
-                            getattr(t, "tool_use_id", None) or getattr(t, "id", None) or ""
-                            for t in tool_use_blocks
-                        ),
-                        content=output.content,
-                    )
-                    yield progress
-                    logger.debug(f"[query] Progress from tool_use_id={tool_id}: {output.content}")
-                elif isinstance(output, ToolResult):
-                    # Tool completed
-                    result_content = output.result_for_assistant or str(output.data)
-                    result_msg = create_user_message(
-                        [{"type": "tool_result", "tool_use_id": tool_id, "text": result_content}],
-                        tool_use_result=output.data,
-                    )
-                    tool_results.append(result_msg)
-                    yield result_msg
-                    logger.debug(
-                        f"[query] Tool completed tool_use_id={tool_id} name={tool_name} "
-                        f"result_len={len(result_content)}"
+                        f"[query] Permission denied for tool_use_id={tool_use_id}: {denial_message}"
                     )
+                    denial_text = denial_message or f"User aborted the tool invocation: {tool_name}"
+                    denial_msg = tool_result_message(tool_use_id, denial_text, is_error=True)
+                    tool_results.append(denial_msg)
+                    yield denial_msg
+                    permission_denied = True
+                    break
+
+            prepared_calls.append(
+                {
+                    "is_concurrency_safe": tool.is_concurrency_safe(),
+                    "generator": _run_tool_use_generator(
+                        tool,
+                        tool_use_id,
+                        tool_name,
+                        parsed_input,
+                        sibling_ids,
+                        tool_context,
+                    ),
+                }
+            )
 
+        except ValidationError as ve:
+            detail_text = format_pydantic_errors(ve)
+            error_msg = tool_result_message(
+                tool_use_id,
+                f"Invalid input for tool '{tool_name}': {detail_text}",
+                is_error=True,
+            )
+            tool_results.append(error_msg)
+            yield error_msg
+            continue
         except Exception as e:
-            # Tool execution failed
-            logger.error(f"Error executing tool '{tool_name}': {e}")
-            error_msg = create_user_message(
-                [
-                    {
-                        "type": "tool_result",
-                        "tool_use_id": tool_id,
-                        "text": f"Error executing tool: {str(e)}",
-                        "is_error": True,
-                    }
-                ]
+            logger.exception(
+                f"Error executing tool '{tool_name}'",
+                extra={"tool": tool_name, "tool_use_id": tool_use_id},
+            )
+            error_msg = tool_result_message(
+                tool_use_id, f"Error executing tool: {str(e)}", is_error=True
             )
             tool_results.append(error_msg)
             yield error_msg
 
+        if permission_denied:
+            break
+
+    if permission_denied:
+        return
+
+    if prepared_calls:
+        async for message in _run_tools_concurrently(prepared_calls, tool_results):
+            yield message
+
     # Check for abort after tools
     if query_context.abort_controller.is_set():
         yield create_assistant_message(INTERRUPT_MESSAGE_FOR_TOOL_USE)
         return
 
-    # Continue conversation with tool results
+    if permission_denied:
+        return
+
     new_messages = messages + [assistant_message] + tool_results
     logger.debug(
         f"[query] Recursing with {len(new_messages)} messages after tools; "