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

ripperdoc/core/query.py

@@ -9,6 +9,7 @@ import inspect
 import os
 import time
 from asyncio import CancelledError
+from dataclasses import dataclass, field
 from typing import (
     Any,
     AsyncGenerator,
@@ -39,6 +40,8 @@ from ripperdoc.core.query_utils import (
     tool_result_message,
 )
 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.log import get_logger
 from ripperdoc.utils.messages import (
@@ -58,7 +61,7 @@ from ripperdoc.utils.messages import (
 logger = get_logger()
 
 DEFAULT_REQUEST_TIMEOUT_SEC = float(os.getenv("RIPPERDOC_API_TIMEOUT", "120"))
-MAX_LLM_RETRIES = 1
+MAX_LLM_RETRIES = int(os.getenv("RIPPERDOC_MAX_RETRIES", "10"))
 
 
 def _resolve_tool(
@@ -118,10 +121,10 @@ async def _check_tool_permissions(
             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)},
+    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
 
@@ -170,9 +173,12 @@ async def _run_tool_use_generator(
                     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}'",
+    except CancelledError:
+        raise  # Don't suppress task cancellation
+    except (RuntimeError, ValueError, TypeError, OSError, IOError, AttributeError, KeyError) as exc:
+        logger.warning(
+            "Error executing tool '%s': %s: %s",
+            tool_name, type(exc).__name__, exc,
             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)
@@ -254,8 +260,15 @@ async def _run_concurrent_tool_uses(
         try:
             async for message in gen:
                 await queue.put(message)
-        except Exception:
-            logger.exception("[query] Unexpected error while consuming tool generator")
+        except asyncio.CancelledError:
+            raise  # Don't suppress cancellation
+        except (StopAsyncIteration, GeneratorExit):
+            pass  # Normal generator termination
+        except (RuntimeError, ValueError, TypeError) as exc:
+            logger.warning(
+                "[query] Error while consuming tool generator: %s: %s",
+                type(exc).__name__, exc,
+            )
         finally:
             await queue.put(None)
 
@@ -304,9 +317,10 @@ class ToolRegistry:
             self._order.append(name)
             try:
                 deferred = tool.defer_loading()
-            except Exception:
-                logger.exception(
-                    "[tool_registry] Tool.defer_loading failed",
+            except (TypeError, AttributeError) as exc:
+                logger.warning(
+                    "[tool_registry] Tool.defer_loading failed: %s: %s",
+                    type(exc).__name__, exc,
                     extra={"tool": getattr(tool, "name", None)},
                 )
                 deferred = False
@@ -367,6 +381,55 @@ class ToolRegistry:
                 yield name, tool
 
 
+def _apply_skill_context_updates(
+    tool_results: List[UserMessage], query_context: "QueryContext"
+) -> None:
+    """Update query context based on Skill tool outputs."""
+    for message in tool_results:
+        data = getattr(message, "tool_use_result", None)
+        if not isinstance(data, dict):
+            continue
+        skill_name = (
+            data.get("skill")
+            or data.get("command_name")
+            or data.get("commandName")
+            or data.get("command")
+        )
+        if not skill_name:
+            continue
+
+        allowed_tools = data.get("allowed_tools") or data.get("allowedTools") or []
+        if allowed_tools and getattr(query_context, "tool_registry", None):
+            try:
+                query_context.tool_registry.activate_tools(
+                    [tool for tool in allowed_tools if isinstance(tool, str) and tool.strip()]
+                )
+            except (KeyError, ValueError, TypeError) as exc:
+                logger.warning(
+                    "[query] Failed to activate tools listed in skill output: %s: %s",
+                    type(exc).__name__, exc,
+                )
+
+        model_hint = data.get("model")
+        if isinstance(model_hint, str) and model_hint.strip():
+            logger.debug(
+                "[query] Applying model hint from skill",
+                extra={"skill": skill_name, "model": model_hint},
+            )
+            query_context.model = model_hint.strip()
+
+        max_tokens = data.get("max_thinking_tokens")
+        if max_tokens is None:
+            max_tokens = data.get("maxThinkingTokens")
+        parsed_max = parse_optional_int(max_tokens)
+        if parsed_max is not None:
+            logger.debug(
+                "[query] Applying max thinking tokens from skill",
+                extra={"skill": skill_name, "max_thinking_tokens": parsed_max},
+            )
+            query_context.max_thinking_tokens = parsed_max
+
+
 class QueryContext:
     """Context for a query session."""
 
@@ -377,6 +440,8 @@ class QueryContext:
         safe_mode: bool = False,
         model: str = "main",
         verbose: bool = False,
+        pause_ui: Optional[Callable[[], None]] = None,
+        resume_ui: Optional[Callable[[], None]] = None,
     ) -> None:
         self.tool_registry = ToolRegistry(tools)
         self.max_thinking_tokens = max_thinking_tokens
@@ -385,6 +450,8 @@ class QueryContext:
         self.verbose = verbose
         self.abort_controller = asyncio.Event()
         self.file_state_cache: Dict[str, FileSnapshot] = {}
+        self.pause_ui = pause_ui
+        self.resume_ui = resume_ui
 
     @property
     def tools(self) -> List[Tool[Any, Any]]:
@@ -411,7 +478,7 @@ async def query_llm(
     tools: List[Tool[Any, Any]],
     max_thinking_tokens: int = 0,
     model: str = "main",
-    abort_signal: Optional[asyncio.Event] = None,
+    _abort_signal: Optional[asyncio.Event] = None,
     *,
     progress_callback: Optional[Callable[[str], Awaitable[None]]] = None,
     request_timeout: Optional[float] = None,
@@ -426,7 +493,7 @@ async def query_llm(
         tools: Available tools
         max_thinking_tokens: Maximum tokens for thinking (0 = disabled)
         model: Model pointer to use
-        abort_signal: Event to signal abortion
+        _abort_signal: Event to signal abortion (currently unused, reserved for future)
         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)
@@ -503,18 +570,50 @@ async def query_llm(
             progress_callback=progress_callback,
             request_timeout=request_timeout,
             max_retries=max_retries,
+            max_thinking_tokens=max_thinking_tokens,
         )
 
+        # Check if provider returned an error response
+        if provider_response.is_error:
+            logger.warning(
+                "[query_llm] Provider returned error response",
+                extra={
+                    "model": model_profile.model,
+                    "error_code": provider_response.error_code,
+                    "error_message": provider_response.error_message,
+                },
+            )
+            metadata: Dict[str, Any] = {
+                "api_error": True,
+                "error_code": provider_response.error_code,
+                "error_message": provider_response.error_message,
+            }
+            # Add context length info if applicable
+            if provider_response.error_code == "context_length_exceeded":
+                metadata["context_length_exceeded"] = True
+
+            error_msg = create_assistant_message(
+                content=provider_response.content_blocks,
+                duration_ms=provider_response.duration_ms,
+                metadata=metadata,
+            )
+            error_msg.is_api_error_message = True
+            return error_msg
+
         return create_assistant_message(
             content=provider_response.content_blocks,
             cost_usd=provider_response.cost_usd,
             duration_ms=provider_response.duration_ms,
+            metadata=provider_response.metadata,
         )
 
-    except Exception as e:
+    except CancelledError:
+        raise  # Don't suppress task cancellation
+    except (RuntimeError, ValueError, TypeError, OSError, ConnectionError, TimeoutError) as e:
         # Return error message
-        logger.exception(
-            "Error querying AI model",
+        logger.warning(
+            "Error querying AI model: %s: %s",
+            type(e).__name__, e,
             extra={
                 "model": getattr(model_profile, "model", None),
                 "model_pointer": model,
@@ -524,56 +623,90 @@ async def query_llm(
             },
         )
         duration_ms = (time.time() - start_time) * 1000
+        context_error = detect_context_length_error(e)
+        metadata = None
+        content = f"Error querying AI model: {str(e)}"
+
+        if context_error:
+            content = f"The request exceeded the model's context window. {context_error.message}"
+            metadata = {
+                "context_length_exceeded": True,
+                "context_length_provider": context_error.provider,
+                "context_length_error_code": context_error.error_code,
+                "context_length_status_code": context_error.status_code,
+            }
+            logger.info(
+                "[query_llm] Detected context-length error; consider compacting history",
+                extra={
+                    "provider": context_error.provider,
+                    "error_code": context_error.error_code,
+                    "status_code": context_error.status_code,
+                },
+            )
+
         error_msg = create_assistant_message(
-            content=f"Error querying AI model: {str(e)}", duration_ms=duration_ms
+            content=content, duration_ms=duration_ms, metadata=metadata
         )
         error_msg.is_api_error_message = True
         return error_msg
 
 
-async def query(
+MAX_QUERY_ITERATIONS = int(os.getenv("RIPPERDOC_MAX_QUERY_ITERATIONS", "1024"))
+
+
+@dataclass
+class IterationResult:
+    """Result of a single query iteration.
+
+    This is used as an "out parameter" to communicate results from
+    _run_query_iteration back to the main query loop.
+    """
+
+    assistant_message: Optional[AssistantMessage] = None
+    tool_results: List[UserMessage] = field(default_factory=list)
+    should_stop: bool = False  # True means exit the query loop entirely
+
+
+async def _run_query_iteration(
     messages: List[Union[UserMessage, AssistantMessage, ProgressMessage]],
     system_prompt: str,
     context: Dict[str, str],
     query_context: QueryContext,
-    can_use_tool_fn: Optional[ToolPermissionCallable] = None,
+    can_use_tool_fn: Optional[ToolPermissionCallable],
+    iteration: int,
+    result: IterationResult,
 ) -> AsyncGenerator[Union[UserMessage, AssistantMessage, ProgressMessage], None]:
-    """Execute a query with tool support.
+    """Run a single iteration of the query loop.
 
-    This is the main query loop that:
-    1. Sends messages to the AI
-    2. Handles tool use responses
-    3. Executes tools
-    4. Recursively continues the conversation
+    This function handles one round of:
+    1. Calling the LLM
+    2. Streaming progress
+    3. Processing tool calls (if any)
 
     Args:
-        messages: Conversation history
+        messages: Current 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
+        iteration: Current iteration number (for logging)
+        result: IterationResult object to store results
 
     Yields:
-        Messages (user, assistant, progress) as they are generated
+        Messages (progress, assistant, tool results) 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)
+    logger.debug(f"[query] 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)
     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()
+    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()
@@ -587,6 +720,7 @@ async def query(
         },
     )
 
+    # Stream LLM response
     progress_queue: asyncio.Queue[Optional[ProgressMessage]] = asyncio.Queue()
 
     async def _stream_progress(chunk: str) -> None:
@@ -600,8 +734,10 @@ async def query(
                     content=chunk,
                 )
             )
-        except Exception:
-            logger.exception("[query] Failed to enqueue stream progress chunk")
+        except asyncio.QueueFull:
+            logger.warning("[query] Progress queue full, dropping chunk")
+        except (RuntimeError, ValueError) as exc:
+            logger.warning("[query] Failed to enqueue stream progress chunk: %s", exc)
 
     assistant_task = asyncio.create_task(
         query_llm(
@@ -620,6 +756,7 @@ async def query(
 
     assistant_message: Optional[AssistantMessage] = None
 
+    # Wait for LLM response while yielding progress
     while True:
         if query_context.abort_controller.is_set():
             assistant_task.cancel()
@@ -628,6 +765,7 @@ async def query(
             except CancelledError:
                 pass
             yield create_assistant_message(INTERRUPT_MESSAGE)
+            result.should_stop = True
             return
         if assistant_task.done():
             assistant_message = await assistant_task
@@ -648,20 +786,24 @@ async def query(
         if progress:
             yield progress
 
+    # Drain remaining progress messages
     while not progress_queue.empty():
         residual = progress_queue.get_nowait()
         if residual:
             yield residual
 
     assert assistant_message is not None
+    result.assistant_message = assistant_message
 
     # Check for abort
     if query_context.abort_controller.is_set():
         yield create_assistant_message(INTERRUPT_MESSAGE)
+        result.should_stop = True
         return
 
     yield assistant_message
 
+    # Extract and process tool calls
     tool_use_blocks: List[MessageContent] = extract_tool_use_blocks(assistant_message)
     text_blocks = (
         len(assistant_message.message.content)
@@ -675,13 +817,16 @@ async def query(
 
     if not tool_use_blocks:
         logger.debug("[query] No tool_use blocks; returning response to user.")
+        result.should_stop = True
         return
 
+    # Process tool calls
     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
+        getattr(t, "tool_use_id", None) or getattr(t, "id", None) or ""
+        for t in tool_use_blocks
     )
     prepared_calls: List[Dict[str, Any]] = []
 
@@ -689,12 +834,18 @@ async def query(
         tool_name = tool_use.name
         if not tool_name:
             continue
-        tool_use_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 {}
 
-        tool, missing_msg = _resolve_tool(query_context.tool_registry, tool_name, tool_use_id)
+        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}")
+            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
@@ -714,12 +865,15 @@ async def query(
                 tool_registry=query_context.tool_registry,
                 file_state_cache=query_context.file_state_cache,
                 abort_signal=query_context.abort_controller,
+                pause_ui=query_context.pause_ui,
+                resume_ui=query_context.resume_ui,
             )
 
             validation = await tool.validate_input(parsed_input, tool_context)
             if not validation.result:
                 logger.debug(
-                    f"[query] Validation failed for tool_use_id={tool_use_id}: {validation.message}"
+                    f"[query] Validation failed for tool_use_id={tool_use_id}: "
+                    f"{validation.message}"
                 )
                 result_msg = tool_result_message(
                     tool_use_id,
@@ -736,9 +890,12 @@ async def query(
                 )
                 if not allowed:
                     logger.debug(
-                        f"[query] Permission denied for tool_use_id={tool_use_id}: {denial_message}"
+                        f"[query] Permission denied for tool_use_id={tool_use_id}: "
+                        f"{denial_message}"
+                    )
+                    denial_text = (
+                        denial_message or f"User aborted the tool invocation: {tool_name}"
                     )
-                    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
@@ -769,9 +926,22 @@ async def query(
             tool_results.append(error_msg)
             yield error_msg
             continue
-        except Exception as e:
-            logger.exception(
-                f"Error executing tool '{tool_name}'",
+        except CancelledError:
+            raise  # Don't suppress task cancellation
+        except (
+            RuntimeError,
+            ValueError,
+            TypeError,
+            OSError,
+            IOError,
+            AttributeError,
+            KeyError,
+        ) as e:
+            logger.warning(
+                "Error executing tool '%s': %s: %s",
+                tool_name,
+                type(e).__name__,
+                e,
                 extra={"tool": tool_name, "tool_use_id": tool_use_id},
             )
             error_msg = tool_result_message(
@@ -784,25 +954,94 @@ async def query(
             break
 
     if permission_denied:
+        result.tool_results = tool_results
+        result.should_stop = True
         return
 
     if prepared_calls:
         async for message in _run_tools_concurrently(prepared_calls, tool_results):
             yield message
 
+    _apply_skill_context_updates(tool_results, query_context)
+
     # Check for abort after tools
     if query_context.abort_controller.is_set():
         yield create_assistant_message(INTERRUPT_MESSAGE_FOR_TOOL_USE)
+        result.tool_results = tool_results
+        result.should_stop = True
         return
 
-    if permission_denied:
-        return
+    result.tool_results = tool_results
+    # should_stop remains False, indicating the loop should continue
 
-    new_messages = messages + [assistant_message] + tool_results
-    logger.debug(
-        f"[query] Recursing with {len(new_messages)} messages after tools; "
-        f"tool_results_count={len(tool_results)}"
+
+async def query(
+    messages: List[Union[UserMessage, AssistantMessage, ProgressMessage]],
+    system_prompt: str,
+    context: Dict[str, str],
+    query_context: QueryContext,
+    can_use_tool_fn: Optional[ToolPermissionCallable] = None,
+) -> AsyncGenerator[Union[UserMessage, AssistantMessage, ProgressMessage], None]:
+    """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
+        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
+    """
+    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 the loop or normalization.
+    messages = list(messages)
+
+    for iteration in range(1, MAX_QUERY_ITERATIONS + 1):
+        result = IterationResult()
 
-    async for msg in query(new_messages, system_prompt, context, query_context, can_use_tool_fn):
-        yield msg
+        async for msg in _run_query_iteration(
+            messages,
+            system_prompt,
+            context,
+            query_context,
+            can_use_tool_fn,
+            iteration,
+            result,
+        ):
+            yield msg
+
+        if result.should_stop:
+            return
+
+        # Update messages for next iteration
+        messages = messages + [result.assistant_message] + result.tool_results
+        logger.debug(
+            f"[query] Continuing loop with {len(messages)} messages after tools; "
+            f"tool_results_count={len(result.tool_results)}"
+        )
+
+    # Reached max iterations
+    logger.warning(
+        f"[query] Reached maximum iterations ({MAX_QUERY_ITERATIONS}), stopping query loop"
+    )
+    yield create_assistant_message(
+        f"Reached maximum query iterations ({MAX_QUERY_ITERATIONS}). "
+        "Please continue the conversation to proceed."
+    )