ripperdoc 0.2.9__py3-none-any.whl → 0.3.0__py3-none-any.whl

ripperdoc/core/query.py

@@ -43,7 +43,12 @@ from ripperdoc.core.query_utils import (
 from ripperdoc.core.tool import Tool, ToolProgress, ToolResult, ToolUseContext
 from ripperdoc.utils.coerce import parse_optional_int
 from ripperdoc.utils.context_length_errors import detect_context_length_error
-from ripperdoc.utils.file_watch import ChangedFileNotice, FileSnapshot, detect_changed_files
+from ripperdoc.utils.file_watch import (
+    BoundedFileCache,
+    ChangedFileNotice,
+    detect_changed_files,
+)
+from ripperdoc.utils.pending_messages import PendingMessageQueue
 from ripperdoc.utils.log import get_logger
 from ripperdoc.utils.messages import (
     AssistantMessage,
@@ -63,6 +68,10 @@ logger = get_logger()
 
 DEFAULT_REQUEST_TIMEOUT_SEC = float(os.getenv("RIPPERDOC_API_TIMEOUT", "120"))
 MAX_LLM_RETRIES = int(os.getenv("RIPPERDOC_MAX_RETRIES", "10"))
+# Timeout for individual tool execution (can be overridden per tool if needed)
+DEFAULT_TOOL_TIMEOUT_SEC = float(os.getenv("RIPPERDOC_TOOL_TIMEOUT", "300"))  # 5 minutes
+# Timeout for concurrent tool execution (total for all tools)
+DEFAULT_CONCURRENT_TOOL_TIMEOUT_SEC = float(os.getenv("RIPPERDOC_CONCURRENT_TOOL_TIMEOUT", "600"))  # 10 minutes
 
 
 def infer_thinking_mode(model_profile: ModelProfile) -> Optional[str]:
@@ -81,6 +90,9 @@ def infer_thinking_mode(model_profile: ModelProfile) -> Optional[str]:
     # Use explicit config if set
     explicit_mode = model_profile.thinking_mode
     if explicit_mode:
+        # "none", "disabled", "off" means thinking is explicitly disabled
+        if explicit_mode.lower() in ("disabled", "off"):
+            return None
         return explicit_mode
 
     # Auto-detect based on API base and model name
@@ -131,7 +143,7 @@ async def _check_tool_permissions(
     parsed_input: Any,
     query_context: "QueryContext",
     can_use_tool_fn: Optional[ToolPermissionCallable],
-) -> tuple[bool, Optional[str]]:
+) -> tuple[bool, Optional[str], Optional[Any]]:
     """Evaluate whether a tool call is allowed."""
     try:
         if can_use_tool_fn is not None:
@@ -139,12 +151,16 @@ async def _check_tool_permissions(
             if inspect.isawaitable(decision):
                 decision = await decision
             if isinstance(decision, PermissionResult):
-                return decision.result, decision.message
+                return decision.result, decision.message, decision.updated_input
             if isinstance(decision, dict) and "result" in decision:
-                return bool(decision.get("result")), decision.get("message")
+                return (
+                    bool(decision.get("result")),
+                    decision.get("message"),
+                    decision.get("updated_input"),
+                )
             if isinstance(decision, tuple) and len(decision) == 2:
-                return bool(decision[0]), decision[1]
-            return bool(decision), None
+                return bool(decision[0]), decision[1], None
+            return bool(decision), None, None
 
         if not query_context.yolo_mode and tool.needs_permissions(parsed_input):
             loop = asyncio.get_running_loop()
@@ -155,15 +171,15 @@ async def _check_tool_permissions(
             )
             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 response.strip().lower() in ("y", "yes"), None, None
 
-        return True, None
+        return True, None, None
     except (TypeError, AttributeError, ValueError) as exc:
         logger.warning(
             f"Error checking permissions for tool '{tool.name}': {type(exc).__name__}: {exc}",
             extra={"tool": getattr(tool, "name", None), "error_type": type(exc).__name__},
         )
-        return False, None
+        return False, None, None
 
 
 def _format_changed_file_notice(notices: List[ChangedFileNotice]) -> str:
@@ -182,6 +198,18 @@ def _format_changed_file_notice(notices: List[ChangedFileNotice]) -> str:
     return "\n".join(lines)
 
 
+def _append_hook_context(context: Dict[str, str], label: str, payload: Optional[str]) -> None:
+    """Append hook-supplied context to the shared context dict."""
+    if not payload:
+        return
+    key = f"Hook:{label}"
+    existing = context.get(key)
+    if existing:
+        context[key] = f"{existing}\n{payload}"
+    else:
+        context[key] = payload
+
+
 async def _run_tool_use_generator(
     tool: Tool[Any, Any],
     tool_use_id: str,
@@ -189,8 +217,14 @@ async def _run_tool_use_generator(
     parsed_input: Any,
     sibling_ids: set[str],
     tool_context: ToolUseContext,
+    context: Dict[str, str],
 ) -> AsyncGenerator[Union[UserMessage, ProgressMessage], None]:
     """Execute a single tool_use and yield progress/results."""
+    logger.debug(
+        "[query] _run_tool_use_generator ENTER: tool='%s' tool_use_id=%s",
+        tool_name,
+        tool_use_id,
+    )
     # Get tool input as dict for hooks
     tool_input_dict = (
         parsed_input.model_dump()
@@ -221,8 +255,14 @@ async def _run_tool_use_generator(
         )
         # Re-parse the input with the updated values
         try:
-            parsed_input = tool.input_schema(**pre_result.updated_input)
-            tool_input_dict = pre_result.updated_input
+            # Ensure updated_input is a dict, not a Pydantic model
+            updated_input = pre_result.updated_input
+            if hasattr(updated_input, "model_dump"):
+                updated_input = updated_input.model_dump()
+            elif not isinstance(updated_input, dict):
+                updated_input = {"value": str(updated_input)}
+            parsed_input = tool.input_schema(**updated_input)
+            tool_input_dict = updated_input
         except (ValueError, TypeError) as exc:
             logger.warning(
                 f"[query] Failed to apply updated input from hook: {exc}",
@@ -235,30 +275,58 @@ async def _run_tool_use_generator(
             f"[query] PreToolUse hook added context for {tool_name}",
             extra={"context": pre_result.additional_context[:100]},
         )
+        _append_hook_context(context, f"PreToolUse:{tool_name}", pre_result.additional_context)
+    if pre_result.system_message:
+        _append_hook_context(context, f"PreToolUse:{tool_name}:system", pre_result.system_message)
 
     tool_output = None
 
     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):
-                tool_output = output.data
-                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)}"
-                )
+        logger.debug("[query] _run_tool_use_generator: BEFORE tool.call() for '%s'", tool_name)
+        # Wrap tool execution with timeout to prevent hangs
+        try:
+            async with asyncio.timeout(DEFAULT_TOOL_TIMEOUT_SEC):
+                async for output in tool.call(parsed_input, tool_context):
+                    logger.debug(
+                        "[query] _run_tool_use_generator: tool='%s' yielded output type=%s",
+                        tool_name,
+                        type(output).__name__,
+                    )
+                    if isinstance(output, ToolProgress):
+                        yield create_progress_message(
+                            tool_use_id=tool_use_id,
+                            sibling_tool_use_ids=sibling_ids,
+                            content=output.content,
+                            is_subagent_message=getattr(output, 'is_subagent_message', False),
+                        )
+                        logger.debug(
+                            f"[query] Progress from tool_use_id={tool_use_id}: {output.content}"
+                        )
+                    elif isinstance(output, ToolResult):
+                        tool_output = output.data
+                        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 asyncio.TimeoutError:
+            logger.error(
+                f"[query] Tool '{tool_name}' timed out after {DEFAULT_TOOL_TIMEOUT_SEC}s",
+                extra={"tool": tool_name, "tool_use_id": tool_use_id},
+            )
+            yield tool_result_message(
+                tool_use_id,
+                f"Tool '{tool_name}' timed out after {DEFAULT_TOOL_TIMEOUT_SEC:.0f} seconds",
+                is_error=True,
+            )
+            return  # Exit early on timeout
+        logger.debug("[query] _run_tool_use_generator: AFTER tool.call() loop for '%s'", tool_name)
     except CancelledError:
+        logger.debug("[query] _run_tool_use_generator: tool='%s' CANCELLED", tool_name)
         raise  # Don't suppress task cancellation
     except (RuntimeError, ValueError, TypeError, OSError, IOError, AttributeError, KeyError) as exc:
         logger.warning(
@@ -271,9 +339,20 @@ async def _run_tool_use_generator(
         yield tool_result_message(tool_use_id, f"Error executing tool: {str(exc)}", is_error=True)
 
     # Run PostToolUse hooks
-    await hook_manager.run_post_tool_use_async(
+    post_result = await hook_manager.run_post_tool_use_async(
         tool_name, tool_input_dict, tool_response=tool_output, tool_use_id=tool_use_id
     )
+    if post_result.additional_context:
+        _append_hook_context(context, f"PostToolUse:{tool_name}", post_result.additional_context)
+    if post_result.system_message:
+        _append_hook_context(context, f"PostToolUse:{tool_name}:system", post_result.system_message)
+    if post_result.should_block:
+        reason = post_result.block_reason or post_result.stop_reason or "Blocked by hook."
+        yield create_user_message(f"PostToolUse hook blocked: {reason}")
+
+    logger.debug(
+        "[query] _run_tool_use_generator DONE: tool='%s' tool_use_id=%s", tool_name, tool_use_id
+    )
 
 
 def _group_tool_calls_by_concurrency(prepared_calls: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
@@ -306,9 +385,18 @@ 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):
+    logger.debug("[query] _execute_tools_in_parallel ENTER: %d items", len(items))
+    valid_items = [call for call in items if call.get("generator")]
+    generators = [call["generator"] for call in valid_items]
+    tool_names = [call.get("tool_name", "unknown") for call in valid_items]
+    logger.debug(
+        "[query] _execute_tools_in_parallel: %d valid generators, tools=%s",
+        len(generators),
+        tool_names,
+    )
+    async for message in _run_concurrent_tool_uses(generators, tool_names, tool_results):
         yield message
+    logger.debug("[query] _execute_tools_in_parallel DONE")
 
 
 async def _run_tools_concurrently(
@@ -340,45 +428,164 @@ async def _run_tools_serially(
 
 async def _run_concurrent_tool_uses(
     generators: List[AsyncGenerator[Union[UserMessage, ProgressMessage], None]],
+    tool_names: List[str],
     tool_results: List[UserMessage],
 ) -> AsyncGenerator[Union[UserMessage, ProgressMessage], None]:
-    """Drain multiple tool generators concurrently and stream outputs."""
+    """Drain multiple tool generators concurrently and stream outputs with overall timeout."""
+    logger.debug(
+        "[query] _run_concurrent_tool_uses ENTER: %d generators, tools=%s, timeout=%s",
+        len(generators),
+        tool_names,
+        DEFAULT_CONCURRENT_TOOL_TIMEOUT_SEC,
+    )
     if not generators:
+        logger.debug("[query] _run_concurrent_tool_uses: no generators, returning")
         return
+        yield  # Make this a proper async generator that yields nothing (unreachable but required)
 
     queue: asyncio.Queue[Optional[Union[UserMessage, ProgressMessage]]] = asyncio.Queue()
 
-    async def _consume(gen: AsyncGenerator[Union[UserMessage, ProgressMessage], None]) -> None:
+    async def _consume(
+        gen: AsyncGenerator[Union[UserMessage, ProgressMessage], None],
+        gen_index: int,
+        tool_name: str,
+    ) -> Optional[Exception]:
+        """Consume a tool generator and return any exception that occurred."""
+        logger.debug(
+            "[query] _consume START: tool='%s' index=%d gen=%s",
+            tool_name,
+            gen_index,
+            type(gen).__name__,
+        )
+        captured_exception: Optional[Exception] = None
+        message_count = 0
         try:
+            logger.debug("[query] _consume: entering async for loop for '%s'", tool_name)
             async for message in gen:
+                message_count += 1
+                msg_type = type(message).__name__
+                logger.debug(
+                    "[query] _consume: tool='%s' received message #%d type=%s",
+                    tool_name,
+                    message_count,
+                    msg_type,
+                )
                 await queue.put(message)
+                logger.debug("[query] _consume: tool='%s' put message to queue", tool_name)
+            logger.debug(
+                "[query] _consume: tool='%s' async for loop finished, total messages=%d",
+                tool_name,
+                message_count,
+            )
         except asyncio.CancelledError:
+            logger.debug("[query] _consume: tool='%s' was CANCELLED", tool_name)
             raise  # Don't suppress cancellation
         except (StopAsyncIteration, GeneratorExit):
+            logger.debug("[query] _consume: tool='%s' StopAsyncIteration/GeneratorExit", tool_name)
             pass  # Normal generator termination
-        except (RuntimeError, ValueError, TypeError) as exc:
+        except Exception as exc:
+            # Capture exception for reporting to caller
+            captured_exception = exc
             logger.warning(
-                "[query] Error while consuming tool generator: %s: %s",
+                "[query] Error while consuming tool '%s' (task %d): %s: %s",
+                tool_name,
+                gen_index,
                 type(exc).__name__,
                 exc,
             )
         finally:
+            logger.debug("[query] _consume FINALLY: tool='%s' putting None to queue", tool_name)
             await queue.put(None)
+            logger.debug("[query] _consume DONE: tool='%s' messages=%d", tool_name, message_count)
+        return captured_exception
 
-    tasks = [asyncio.create_task(_consume(gen)) for gen in generators]
+    logger.debug("[query] _run_concurrent_tool_uses: creating %d tasks", len(generators))
+    tasks = [
+        asyncio.create_task(_consume(gen, i, tool_names[i])) for i, gen in enumerate(generators)
+    ]
     active = len(tasks)
+    logger.debug("[query] _run_concurrent_tool_uses: %d tasks created, entering while loop", active)
 
     try:
-        while active:
-            message = await queue.get()
-            if message is None:
-                active -= 1
-                continue
-            if isinstance(message, UserMessage):
-                tool_results.append(message)
-            yield message
+        # Add overall timeout for entire concurrent execution
+        async with asyncio.timeout(DEFAULT_CONCURRENT_TOOL_TIMEOUT_SEC):
+            while active:
+                logger.debug(
+                    "[query] _run_concurrent_tool_uses: waiting for queue.get(), active=%d", active
+                )
+                try:
+                    message = await asyncio.wait_for(
+                        queue.get(), timeout=DEFAULT_CONCURRENT_TOOL_TIMEOUT_SEC
+                    )
+                except asyncio.TimeoutError:
+                    logger.error(
+                        "[query] Concurrent tool execution timed out waiting for messages"
+                    )
+                    # Cancel all remaining tasks
+                    for task in tasks:
+                        if not task.done():
+                            task.cancel()
+                    raise
+
+                logger.debug(
+                    "[query] _run_concurrent_tool_uses: got message type=%s, active=%d",
+                    type(message).__name__ if message else "None",
+                    active,
+                )
+                if message is None:
+                    active -= 1
+                    logger.debug(
+                        "[query] _run_concurrent_tool_uses: None received, active now=%d", active
+                    )
+                    continue
+                if isinstance(message, UserMessage):
+                    tool_results.append(message)
+                yield message
+            logger.debug("[query] _run_concurrent_tool_uses: while loop finished, all tools done")
+    except asyncio.TimeoutError:
+        logger.error(
+            f"[query] Concurrent tool execution timed out after {DEFAULT_CONCURRENT_TOOL_TIMEOUT_SEC}s",
+            extra={"tool_names": tool_names},
+        )
+        # Ensure all tasks are cancelled
+        for task in tasks:
+            if not task.done():
+                task.cancel()
+        raise
     finally:
-        await asyncio.gather(*tasks, return_exceptions=True)
+        # Wait for all tasks and collect any exceptions
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        exceptions_found: List[tuple[int, str, BaseException]] = []
+        for i, result in enumerate(results):
+            if isinstance(result, asyncio.CancelledError):
+                continue
+            elif isinstance(result, Exception):
+                # Exception from gather itself (shouldn't happen with return_exceptions=True)
+                exceptions_found.append((i, tool_names[i], result))
+            elif result is not None:
+                # Exception returned by _consume
+                exceptions_found.append((i, tool_names[i], result))
+
+        # Log all exceptions for debugging
+        for i, name, exc in exceptions_found:
+            logger.warning(
+                "[query] Concurrent tool '%s' (task %d) failed: %s: %s",
+                name,
+                i,
+                type(exc).__name__,
+                exc,
+            )
+
+        # Re-raise first exception if any occurred, so caller knows something failed
+        if exceptions_found:
+            first_name = exceptions_found[0][1]
+            first_exc = exceptions_found[0][2]
+            logger.error(
+                "[query] %d tool(s) failed during concurrent execution, first error in '%s': %s",
+                len(exceptions_found),
+                first_name,
+                first_exc,
+            )
 
 
 class ToolRegistry:
@@ -451,6 +658,9 @@ class ToolRegistry:
         """Activate deferred tools by name."""
         activated: List[str] = []
         missing: List[str] = []
+
+        # First pass: collect tools to activate (no mutations)
+        to_activate: List[str] = []
         for raw_name in names:
             name = (raw_name or "").strip()
             if not name:
@@ -459,12 +669,17 @@ class ToolRegistry:
                 continue
             tool = self._tool_map.get(name)
             if tool:
-                self._active.append(name)
-                self._active_set.add(name)
-                self._deferred.discard(name)
-                activated.append(name)
+                to_activate.append(name)
             else:
                 missing.append(name)
+
+        # Second pass: atomically update all data structures
+        if to_activate:
+            self._active.extend(to_activate)
+            self._active_set.update(to_activate)
+            self._deferred.difference_update(to_activate)
+            activated.extend(to_activate)
+
         return activated, missing
 
     def iter_named_tools(self) -> Iterable[tuple[str, Tool[Any, Any]]]:
@@ -537,6 +752,12 @@ class QueryContext:
         verbose: bool = False,
         pause_ui: Optional[Callable[[], None]] = None,
         resume_ui: Optional[Callable[[], None]] = None,
+        stop_hook: str = "stop",
+        file_cache_max_entries: int = 500,
+        file_cache_max_memory_mb: float = 50.0,
+        pending_message_queue: Optional[PendingMessageQueue] = None,
+        max_turns: Optional[int] = None,
+        permission_mode: str = "default",
     ) -> None:
         self.tool_registry = ToolRegistry(tools)
         self.max_thinking_tokens = max_thinking_tokens
@@ -544,9 +765,20 @@ class QueryContext:
         self.model = model
         self.verbose = verbose
         self.abort_controller = asyncio.Event()
-        self.file_state_cache: Dict[str, FileSnapshot] = {}
+        self.pending_message_queue: PendingMessageQueue = (
+            pending_message_queue if pending_message_queue is not None else PendingMessageQueue()
+        )
+        # Use BoundedFileCache instead of plain Dict to prevent unbounded growth
+        self.file_state_cache: BoundedFileCache = BoundedFileCache(
+            max_entries=file_cache_max_entries,
+            max_memory_mb=file_cache_max_memory_mb,
+        )
         self.pause_ui = pause_ui
         self.resume_ui = resume_ui
+        self.stop_hook = stop_hook
+        self.stop_hook_active = False
+        self.max_turns = max_turns
+        self.permission_mode = permission_mode
 
     @property
     def tools(self) -> List[Tool[Any, Any]]:
@@ -566,6 +798,22 @@ class QueryContext:
         """Return all known tools (active + deferred)."""
         return self.tool_registry.all_tools
 
+    def get_memory_stats(self) -> Dict[str, Any]:
+        """Return memory usage statistics for monitoring."""
+        return {
+            "file_cache": self.file_state_cache.stats(),
+            "tool_count": len(self.tool_registry.all_tools),
+            "active_tool_count": len(self.tool_registry.active_tools),
+        }
+
+    def drain_pending_messages(self) -> List[UserMessage]:
+        """Drain queued messages waiting to be injected into the conversation."""
+        return self.pending_message_queue.drain()
+
+    def enqueue_user_message(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> None:
+        """Queue a user-style message to inject once the current loop finishes."""
+        self.pending_message_queue.enqueue_text(text, metadata=metadata)
+
 
 async def query_llm(
     messages: List[Union[UserMessage, AssistantMessage, ProgressMessage]],
@@ -598,7 +846,6 @@ async def query_llm(
         AssistantMessage with the model's response
     """
     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)
@@ -657,15 +904,29 @@ async def query_llm(
     start_time = time.time()
 
     try:
-        client: Optional[ProviderClient] = get_provider_client(model_profile.provider)
+        try:
+            client: Optional[ProviderClient] = get_provider_client(model_profile.provider)
+        except RuntimeError as exc:
+            duration_ms = (time.time() - start_time) * 1000
+            error_msg = create_assistant_message(
+                content=str(exc),
+                duration_ms=duration_ms,
+                model=model_profile.model,
+            )
+            error_msg.is_api_error_message = True
+            return error_msg
         if client is None:
             duration_ms = (time.time() - start_time) * 1000
+            provider_label = getattr(model_profile.provider, "value", None) or str(
+                model_profile.provider
+            )
             error_msg = create_assistant_message(
                 content=(
-                    "Gemini protocol is not supported yet in Ripperdoc. "
-                    "Please configure an Anthropic or OpenAI-compatible model."
+                    f"No provider client available for '{provider_label}'. "
+                    "Check your model configuration and provider dependencies."
                 ),
                 duration_ms=duration_ms,
+                model=model_profile.model,
             )
             error_msg.is_api_error_message = True
             return error_msg
@@ -706,6 +967,7 @@ async def query_llm(
                 content=provider_response.content_blocks,
                 duration_ms=provider_response.duration_ms,
                 metadata=metadata,
+                model=model_profile.model,
             )
             error_msg.is_api_error_message = True
             return error_msg
@@ -715,6 +977,13 @@ async def query_llm(
             cost_usd=provider_response.cost_usd,
             duration_ms=provider_response.duration_ms,
             metadata=provider_response.metadata,
+            model=model_profile.model,
+            input_tokens=provider_response.usage_tokens.get("input_tokens", 0),
+            output_tokens=provider_response.usage_tokens.get("output_tokens", 0),
+            cache_read_tokens=provider_response.usage_tokens.get("cache_read_input_tokens", 0),
+            cache_creation_tokens=provider_response.usage_tokens.get(
+                "cache_creation_input_tokens", 0
+            ),
         )
 
     except CancelledError:
@@ -756,7 +1025,10 @@ async def query_llm(
             )
 
         error_msg = create_assistant_message(
-            content=content, duration_ms=duration_ms, metadata=error_metadata
+            content=content,
+            duration_ms=duration_ms,
+            metadata=error_metadata,
+            model=model_profile.model,
         )
         error_msg.is_api_error_message = True
         return error_msg
@@ -806,7 +1078,7 @@ async def _run_query_iteration(
     Yields:
         Messages (progress, assistant, tool results) as they are generated
     """
-    logger.debug(f"[query] Iteration {iteration}/{MAX_QUERY_ITERATIONS}")
+    logger.info(f"[query] Starting iteration {iteration}/{MAX_QUERY_ITERATIONS}")
 
     # Check for file changes at the start of each iteration
     change_notices = detect_changed_files(query_context.file_state_cache)
@@ -830,21 +1102,25 @@ async def _run_query_iteration(
     )
 
     # Stream LLM response
-    progress_queue: asyncio.Queue[Optional[ProgressMessage]] = asyncio.Queue()
+    progress_queue: asyncio.Queue[Optional[ProgressMessage]] = asyncio.Queue(maxsize=1000)
 
     async def _stream_progress(chunk: str) -> None:
         if not chunk:
             return
         try:
-            await progress_queue.put(
-                create_progress_message(
-                    tool_use_id="stream",
-                    sibling_tool_use_ids=set(),
-                    content=chunk,
-                )
+            msg = create_progress_message(
+                tool_use_id="stream",
+                sibling_tool_use_ids=set(),
+                content=chunk,
             )
-        except asyncio.QueueFull:
-            logger.warning("[query] Progress queue full, dropping chunk")
+            try:
+                progress_queue.put_nowait(msg)
+            except asyncio.QueueFull:
+                # Queue full - wait with timeout instead of dropping immediately
+                try:
+                    await asyncio.wait_for(progress_queue.put(msg), timeout=0.5)
+                except asyncio.TimeoutError:
+                    logger.warning("[query] Progress queue full after timeout, dropping chunk")
         except (RuntimeError, ValueError) as exc:
             logger.warning("[query] Failed to enqueue stream progress chunk: %s", exc)
 
@@ -863,6 +1139,8 @@ async def _run_query_iteration(
         )
     )
 
+    logger.debug("[query] Created query_llm task, waiting for response...")
+
     assistant_message: Optional[AssistantMessage] = None
 
     # Wait for LLM response while yielding progress
@@ -873,7 +1151,7 @@ async def _run_query_iteration(
                 await assistant_task
             except CancelledError:
                 pass
-            yield create_assistant_message(INTERRUPT_MESSAGE)
+            yield create_assistant_message(INTERRUPT_MESSAGE, model=model_profile.model)
             result.should_stop = True
             return
         if assistant_task.done():
@@ -883,23 +1161,23 @@ async def _run_query_iteration(
             progress = progress_queue.get_nowait()
         except asyncio.QueueEmpty:
             waiter = asyncio.create_task(progress_queue.get())
-            # Use timeout to periodically check abort_controller during LLM request
+            abort_waiter = asyncio.create_task(query_context.abort_controller.wait())
             done, pending = await asyncio.wait(
-                {assistant_task, waiter},
+                {assistant_task, waiter, abort_waiter},
                 return_when=asyncio.FIRST_COMPLETED,
-                timeout=0.1,  # Check abort_controller every 100ms
             )
-            if not done:
-                # Timeout - cancel waiter and continue loop to check abort_controller
-                waiter.cancel()
-                try:
-                    await waiter
-                except asyncio.CancelledError:
-                    pass
+            for task in pending:
+                # Don't cancel assistant_task here - it should only be cancelled
+                # through abort_controller in the main loop
+                if task is not assistant_task:
+                    task.cancel()
+                    try:
+                        await task
+                    except asyncio.CancelledError:
+                        pass
+            if abort_waiter in done:
                 continue
             if assistant_task in done:
-                for task in pending:
-                    task.cancel()
                 assistant_message = await assistant_task
                 break
             progress = waiter.result()
@@ -912,12 +1190,13 @@ async def _run_query_iteration(
         if residual:
             yield residual
 
-    assert assistant_message is not None
+    if assistant_message is None:
+        raise RuntimeError("assistant_message was unexpectedly None after LLM query")
     result.assistant_message = assistant_message
 
     # Check for abort
     if query_context.abort_controller.is_set():
-        yield create_assistant_message(INTERRUPT_MESSAGE)
+        yield create_assistant_message(INTERRUPT_MESSAGE, model=model_profile.model)
         result.should_stop = True
         return
 
@@ -936,7 +1215,39 @@ async def _run_query_iteration(
     )
 
     if not tool_use_blocks:
-        logger.debug("[query] No tool_use blocks; returning response to user.")
+        logger.debug(
+            "[query] No tool_use blocks; running stop hook and returning response to user."
+        )
+        stop_hook = query_context.stop_hook
+        logger.debug(
+            f"[query] stop_hook={stop_hook}, stop_hook_active={query_context.stop_hook_active}"
+        )
+        logger.debug("[query] BEFORE calling hook_manager.run_stop_async")
+        stop_result = (
+            await hook_manager.run_subagent_stop_async(
+                stop_hook_active=query_context.stop_hook_active
+            )
+            if stop_hook == "subagent"
+            else await hook_manager.run_stop_async(stop_hook_active=query_context.stop_hook_active)
+        )
+        logger.debug("[query] AFTER calling hook_manager.run_stop_async")
+        logger.debug("[query] Checking additional_context")
+        if stop_result.additional_context:
+            _append_hook_context(context, f"{stop_hook}:context", stop_result.additional_context)
+        logger.debug("[query] Checking system_message")
+        if stop_result.system_message:
+            _append_hook_context(context, f"{stop_hook}:system", stop_result.system_message)
+        logger.debug("[query] Checking should_block")
+        if stop_result.should_block:
+            reason = stop_result.block_reason or stop_result.stop_reason or "Blocked by hook."
+            result.tool_results = [create_user_message(f"{stop_hook} hook blocked: {reason}")]
+            for msg in result.tool_results:
+                yield msg
+            query_context.stop_hook_active = True
+            result.should_stop = False
+            return
+        logger.debug("[query] Setting should_stop=True and returning")
+        query_context.stop_hook_active = False
         result.should_stop = True
         return
 
@@ -956,13 +1267,25 @@ async def _run_query_iteration(
         tool_use_id = getattr(tool_use, "tool_use_id", None) or getattr(tool_use, "id", None) or ""
         tool_input = getattr(tool_use, "input", {}) or {}
 
+        # Handle case where input is a Pydantic model instead of a dict
+        # This can happen when the API response contains structured tool input objects
+        # Always try to convert if it has model_dump or dict methods
+        if tool_input and hasattr(tool_input, "model_dump"):
+            tool_input = tool_input.model_dump()
+        elif tool_input and hasattr(tool_input, "dict") and callable(getattr(tool_input, "dict")):
+            tool_input = tool_input.dict()
+        elif tool_input and not isinstance(tool_input, dict):
+            # Last resort: convert unknown type to string representation
+            tool_input = {"value": str(tool_input)}
+
         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
-        assert tool is not None
+        if tool is None:
+            raise RuntimeError(f"Tool '{tool_name}' resolved to None unexpectedly")
 
         try:
             parsed_input = tool.input_schema(**tool_input)
@@ -972,14 +1295,17 @@ async def _run_query_iteration(
             )
 
             tool_context = ToolUseContext(
+                message_id=tool_use_id,  # Set message_id for parent_tool_use_id tracking
                 yolo_mode=query_context.yolo_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,
+                conversation_messages=messages,
                 abort_signal=query_context.abort_controller,
                 pause_ui=query_context.pause_ui,
                 resume_ui=query_context.resume_ui,
+                pending_message_queue=query_context.pending_message_queue,
             )
 
             validation = await tool.validate_input(parsed_input, tool_context)
@@ -997,7 +1323,7 @@ async def _run_query_iteration(
                 continue
 
             if not query_context.yolo_mode or can_use_tool_fn is not None:
-                allowed, denial_message = await _check_tool_permissions(
+                allowed, denial_message, updated_input = await _check_tool_permissions(
                     tool, parsed_input, query_context, can_use_tool_fn
                 )
                 if not allowed:
@@ -1010,9 +1336,39 @@ async def _run_query_iteration(
                     yield denial_msg
                     permission_denied = True
                     break
+                if updated_input:
+                    try:
+                        # Ensure updated_input is a dict, not a Pydantic model
+                        normalized_input = updated_input
+                        if hasattr(normalized_input, "model_dump"):
+                            normalized_input = normalized_input.model_dump()
+                        elif not isinstance(normalized_input, dict):
+                            normalized_input = {"value": str(normalized_input)}
+                        parsed_input = tool.input_schema(**normalized_input)
+                    except ValidationError as ve:
+                        detail_text = format_pydantic_errors(ve)
+                        error_msg = tool_result_message(
+                            tool_use_id,
+                            f"Invalid permission-updated input for tool '{tool_name}': {detail_text}",
+                            is_error=True,
+                        )
+                        tool_results.append(error_msg)
+                        yield error_msg
+                        continue
+                    validation = await tool.validate_input(parsed_input, tool_context)
+                    if not validation.result:
+                        error_msg = tool_result_message(
+                            tool_use_id,
+                            validation.message or "Tool input validation failed.",
+                            is_error=True,
+                        )
+                        tool_results.append(error_msg)
+                        yield error_msg
+                        continue
 
             prepared_calls.append(
                 {
+                    "tool_name": tool_name,
                     "is_concurrency_safe": tool.is_concurrency_safe(),
                     "generator": _run_tool_use_generator(
                         tool,
@@ -1021,6 +1377,7 @@ async def _run_query_iteration(
                         parsed_input,
                         sibling_ids,
                         tool_context,
+                        context,
                     ),
                 }
             )
@@ -1075,7 +1432,7 @@ async def _run_query_iteration(
 
     # Check for abort after tools
     if query_context.abort_controller.is_set():
-        yield create_assistant_message(INTERRUPT_MESSAGE_FOR_TOOL_USE)
+        yield create_assistant_message(INTERRUPT_MESSAGE_FOR_TOOL_USE, model=model_profile.model)
         result.tool_results = tool_results
         result.should_stop = True
         return
@@ -1099,6 +1456,26 @@ async def query(
     3. Executes tools
     4. Continues the conversation in a loop until no more tool calls
 
+    Args:
+        messages: Conversation history
+        system_prompt: Base system prompt
+        context: Additional context dictionary
+        query_context: Query configuration
+        can_use_tool_fn: Optional function to check tool permissions
+
+    Yields:
+        Messages (user, assistant, progress) as they are generated
+    """
+    # Resolve model once for use in messages (e.g., max iterations, errors)
+    model_profile = resolve_model_profile(query_context.model)
+    """Execute a query with tool support.
+
+    This is the main query loop that:
+    1. Sends messages to the AI
+    2. Handles tool use responses
+    3. Executes tools
+    4. Continues the conversation in a loop until no more tool calls
+
     Args:
         messages: Conversation history
         system_prompt: Base system prompt
@@ -1116,6 +1493,8 @@ async def query(
             "tool_count": len(query_context.tools),
             "yolo_mode": query_context.yolo_mode,
             "model_pointer": query_context.model,
+            "max_turns": query_context.max_turns,
+            "permission_mode": query_context.permission_mode,
         },
     )
     # Work on a copy so external mutations (e.g., UI appending messages while consuming)
@@ -1123,6 +1502,13 @@ async def query(
     messages = list(messages)
 
     for iteration in range(1, MAX_QUERY_ITERATIONS + 1):
+        # Inject any pending messages queued by background events or user interjections
+        pending_messages = query_context.drain_pending_messages()
+        if pending_messages:
+            messages.extend(pending_messages)
+            for pending in pending_messages:
+                yield pending
+
         result = IterationResult()
 
         async for msg in _run_query_iteration(
@@ -1137,6 +1523,20 @@ async def query(
             yield msg
 
         if result.should_stop:
+            # Before stopping, check if new pending messages arrived during this iteration.
+            trailing_pending = query_context.drain_pending_messages()
+            if trailing_pending:
+                # type: ignore[operator,list-item]
+                next_messages = (
+                    messages + [result.assistant_message] + result.tool_results
+                    if result.assistant_message is not None
+                    else messages + result.tool_results  # type: ignore[operator]
+                )  # type: ignore[operator]
+                next_messages = next_messages + trailing_pending  # type: ignore[operator,list-item]
+                for pending in trailing_pending:
+                    yield pending
+                messages = next_messages
+                continue
             return
 
         # Update messages for next iteration
@@ -1144,6 +1544,7 @@ async def query(
             messages = messages + [result.assistant_message] + result.tool_results  # type: ignore[operator]
         else:
             messages = messages + result.tool_results  # type: ignore[operator]
+
         logger.debug(
             f"[query] Continuing loop with {len(messages)} messages after tools; "
             f"tool_results_count={len(result.tool_results)}"
@@ -1155,5 +1556,6 @@ async def query(
     )
     yield create_assistant_message(
         f"Reached maximum query iterations ({MAX_QUERY_ITERATIONS}). "
-        "Please continue the conversation to proceed."
+        "Please continue the conversation to proceed.",
+        model=model_profile.model,
     )