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

ripperdoc/core/query.py

@@ -48,6 +48,7 @@ from ripperdoc.utils.file_watch import (
     ChangedFileNotice,
     detect_changed_files,
 )
+from ripperdoc.utils.pending_messages import PendingMessageQueue
 from ripperdoc.utils.log import get_logger
 from ripperdoc.utils.messages import (
     AssistantMessage,
@@ -67,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]:
@@ -85,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
@@ -145,8 +153,10 @@ async def _check_tool_permissions(
             if isinstance(decision, PermissionResult):
                 return decision.result, decision.message, decision.updated_input
             if isinstance(decision, dict) and "result" in decision:
-                return bool(decision.get("result")), decision.get("message"), decision.get(
-                    "updated_input"
+                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], None
@@ -210,6 +220,11 @@ async def _run_tool_use_generator(
     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()
@@ -240,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}",
@@ -261,26 +282,51 @@ async def _run_tool_use_generator(
     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(
@@ -299,13 +345,15 @@ async def _run_tool_use_generator(
     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
-        )
+        _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]]:
     """Group consecutive tool calls by their concurrency safety."""
@@ -337,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(
@@ -371,46 +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
+        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:
@@ -483,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:
@@ -491,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]]]:
@@ -560,14 +743,6 @@ def _apply_skill_context_updates(
 class QueryContext:
     """Context for a query session."""
 
-    # Thresholds for memory warnings
-    MESSAGE_COUNT_WARNING_THRESHOLD = int(
-        os.getenv("RIPPERDOC_MESSAGE_WARNING_THRESHOLD", "500")
-    )
-    MESSAGE_COUNT_CRITICAL_THRESHOLD = int(
-        os.getenv("RIPPERDOC_MESSAGE_CRITICAL_THRESHOLD", "1000")
-    )
-
     def __init__(
         self,
         tools: List[Tool[Any, Any]],
@@ -580,6 +755,9 @@ class QueryContext:
         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
@@ -587,6 +765,9 @@ class QueryContext:
         self.model = model
         self.verbose = verbose
         self.abort_controller = asyncio.Event()
+        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,
@@ -596,7 +777,8 @@ class QueryContext:
         self.resume_ui = resume_ui
         self.stop_hook = stop_hook
         self.stop_hook_active = False
-        self._last_message_warning_count = 0
+        self.max_turns = max_turns
+        self.permission_mode = permission_mode
 
     @property
     def tools(self) -> List[Tool[Any, Any]]:
@@ -616,36 +798,6 @@ class QueryContext:
         """Return all known tools (active + deferred)."""
         return self.tool_registry.all_tools
 
-    def check_message_count(self, message_count: int) -> None:
-        """Check message count and log warnings if thresholds are exceeded.
-
-        This helps detect potential memory issues in long sessions.
-        """
-        if message_count >= self.MESSAGE_COUNT_CRITICAL_THRESHOLD:
-            if self._last_message_warning_count < self.MESSAGE_COUNT_CRITICAL_THRESHOLD:
-                logger.warning(
-                    "[query] Critical: Message history is very large. "
-                    "Consider compacting or starting a new session.",
-                    extra={
-                        "message_count": message_count,
-                        "threshold": self.MESSAGE_COUNT_CRITICAL_THRESHOLD,
-                        "file_cache_stats": self.file_state_cache.stats(),
-                    },
-                )
-                self._last_message_warning_count = message_count
-        elif message_count >= self.MESSAGE_COUNT_WARNING_THRESHOLD:
-            # Only warn once per threshold crossing
-            if self._last_message_warning_count < self.MESSAGE_COUNT_WARNING_THRESHOLD:
-                logger.info(
-                    "[query] Message history growing large; automatic compaction may trigger soon",
-                    extra={
-                        "message_count": message_count,
-                        "threshold": self.MESSAGE_COUNT_WARNING_THRESHOLD,
-                        "file_cache_stats": self.file_state_cache.stats(),
-                    },
-                )
-                self._last_message_warning_count = message_count
-
     def get_memory_stats(self) -> Dict[str, Any]:
         """Return memory usage statistics for monitoring."""
         return {
@@ -654,6 +806,14 @@ class QueryContext:
             "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]],
@@ -751,6 +911,7 @@ async def query_llm(
             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
@@ -765,6 +926,7 @@ async def query_llm(
                     "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
@@ -805,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
@@ -862,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
@@ -912,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)
@@ -942,15 +1108,19 @@ async def _run_query_iteration(
         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)
 
@@ -969,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
@@ -979,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():
@@ -1024,7 +1196,7 @@ async def _run_query_iteration(
 
     # 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
 
@@ -1043,8 +1215,14 @@ 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
@@ -1052,10 +1230,14 @@ async def _run_query_iteration(
             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}")]
@@ -1064,6 +1246,7 @@ async def _run_query_iteration(
             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
@@ -1084,6 +1267,17 @@ 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}")
@@ -1101,6 +1295,7 @@ 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,
@@ -1110,6 +1305,7 @@ async def _run_query_iteration(
                 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)
@@ -1142,7 +1338,13 @@ async def _run_query_iteration(
                     break
                 if updated_input:
                     try:
-                        parsed_input = tool.input_schema(**updated_input)
+                        # 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(
@@ -1166,6 +1368,7 @@ async def _run_query_iteration(
 
             prepared_calls.append(
                 {
+                    "tool_name": tool_name,
                     "is_concurrency_safe": tool.is_concurrency_safe(),
                     "generator": _run_tool_use_generator(
                         tool,
@@ -1229,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
@@ -1253,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
@@ -1270,16 +1493,22 @@ 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)
     # do not interfere with the loop or normalization.
     messages = list(messages)
 
-    # Check initial message count for memory warnings
-    query_context.check_message_count(len(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(
@@ -1294,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
@@ -1302,9 +1545,6 @@ async def query(
         else:
             messages = messages + result.tool_results  # type: ignore[operator]
 
-        # Check message count after each iteration for memory warnings
-        query_context.check_message_count(len(messages))
-
         logger.debug(
             f"[query] Continuing loop with {len(messages)} messages after tools; "
             f"tool_results_count={len(result.tool_results)}"
@@ -1316,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,
     )