shotgun-sh 0.2.6.dev1__py3-none-any.whl → 0.2.17__py3-none-any.whl

shotgun/agents/export.py

@@ -4,7 +4,6 @@ from functools import partial
 
 from pydantic_ai import (
     Agent,
-    DeferredToolRequests,
 )
 from pydantic_ai.agent import AgentRunResult
 from pydantic_ai.messages import ModelMessage
@@ -19,14 +18,14 @@ from .common import (
     create_usage_limits,
     run_agent,
 )
-from .models import AgentDeps, AgentRuntimeOptions, AgentType
+from .models import AgentDeps, AgentResponse, AgentRuntimeOptions, AgentType
 
 logger = get_logger(__name__)
 
 
-def create_export_agent(
+async def create_export_agent(
     agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
-) -> tuple[Agent[AgentDeps, str | DeferredToolRequests], AgentDeps]:
+) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
     """Create an export agent with file management capabilities.
 
     Args:
@@ -40,7 +39,7 @@ def create_export_agent(
     # Use partial to create system prompt function for export agent
     system_prompt_fn = partial(build_agent_system_prompt, "export")
 
-    agent, deps = create_base_agent(
+    agent, deps = await create_base_agent(
         system_prompt_fn,
         agent_runtime_options,
         provider=provider,
@@ -50,11 +49,11 @@ def create_export_agent(
 
 
 async def run_export_agent(
-    agent: Agent[AgentDeps, str | DeferredToolRequests],
+    agent: Agent[AgentDeps, AgentResponse],
     instruction: str,
     deps: AgentDeps,
     message_history: list[ModelMessage] | None = None,
-) -> AgentRunResult[str | DeferredToolRequests]:
+) -> AgentRunResult[AgentResponse]:
     """Export artifacts based on the given instruction.
 
     Args:

shotgun/agents/history/compaction.py

@@ -13,7 +13,7 @@ logger = get_logger(__name__)
 
 
 async def apply_persistent_compaction(
-    messages: list[ModelMessage], deps: AgentDeps
+    messages: list[ModelMessage], deps: AgentDeps, force: bool = False
 ) -> list[ModelMessage]:
     """Apply compaction to message history for persistent storage.
 
@@ -23,6 +23,7 @@ async def apply_persistent_compaction(
     Args:
         messages: Full message history from agent run
         deps: Agent dependencies containing model config
+        force: If True, force compaction even if below token threshold
 
     Returns:
         Compacted message history that should be stored as conversation state
@@ -46,7 +47,7 @@ async def apply_persistent_compaction(
                 self.usage = usage
 
         ctx = MockContext(deps, usage)
-        compacted_messages = await token_limit_compactor(ctx, messages)
+        compacted_messages = await token_limit_compactor(ctx, messages, force=force)
 
         # Log the result for monitoring
         original_size = len(messages)
@@ -59,17 +60,21 @@ async def apply_persistent_compaction(
                 f"({reduction_pct:.1f}% reduction)"
             )
 
-            # Track persistent compaction event
+            # Track persistent compaction event with simple metrics (fast, no token counting)
             track_event(
                 "persistent_compaction_applied",
                 {
+                    # Basic compaction metrics
                     "messages_before": original_size,
                     "messages_after": compacted_size,
-                    "tokens_before": estimated_tokens,
                     "reduction_percentage": round(reduction_pct, 2),
                     "agent_mode": deps.agent_mode.value
                     if hasattr(deps, "agent_mode") and deps.agent_mode
                     else "unknown",
+                    # Model and provider info (no computation needed)
+                    "model_name": deps.llm_model.name.value,
+                    "provider": deps.llm_model.provider.value,
+                    "key_provider": deps.llm_model.key_provider.value,
                 },
             )
         else:

shotgun/agents/history/context_extraction.py

@@ -1,5 +1,9 @@
 """Context extraction utilities for history processing."""
 
+import json
+import logging
+import traceback
+
 from pydantic_ai.messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
@@ -16,6 +20,46 @@ from pydantic_ai.messages import (
     UserPromptPart,
 )
 
+logger = logging.getLogger(__name__)
+
+
+def _safely_parse_tool_args(args: dict[str, object] | str | None) -> dict[str, object]:
+    """Safely parse tool call arguments, handling incomplete/invalid JSON.
+
+    Args:
+        args: Tool call arguments (dict, JSON string, or None)
+
+    Returns:
+        Parsed args dict, or empty dict if parsing fails
+    """
+    if args is None:
+        return {}
+
+    if isinstance(args, dict):
+        return args
+
+    if not isinstance(args, str):
+        return {}
+
+    try:
+        parsed = json.loads(args)
+        return parsed if isinstance(parsed, dict) else {}
+    except (json.JSONDecodeError, ValueError) as e:
+        # Only log warning if it looks like JSON (starts with { or [) - incomplete JSON
+        # Plain strings are valid args and shouldn't trigger warnings
+        stripped_args = args.strip()
+        if stripped_args.startswith(("{", "[")):
+            args_preview = args[:100] + "..." if len(args) > 100 else args
+            logger.warning(
+                "Detected incomplete/invalid JSON in tool call args during parsing",
+                extra={
+                    "args_preview": args_preview,
+                    "error": str(e),
+                    "args_length": len(args),
+                },
+            )
+        return {}
+
 
 def extract_context_from_messages(messages: list[ModelMessage]) -> str:
     """Extract context from a list of messages for summarization."""
@@ -87,12 +131,55 @@ def extract_context_from_part(
         return f"<ASSISTANT_TEXT>\n{message_part.content}\n</ASSISTANT_TEXT>"
 
     elif isinstance(message_part, ToolCallPart):
-        if isinstance(message_part.args, dict):
-            args_str = ", ".join(f"{k}={repr(v)}" for k, v in message_part.args.items())
-            tool_call_str = f"{message_part.tool_name}({args_str})"
-        else:
-            tool_call_str = f"{message_part.tool_name}({message_part.args})"
-        return f"<TOOL_CALL>\n{tool_call_str}\n</TOOL_CALL>"
+        # Safely parse args to avoid crashes from incomplete JSON during streaming
+        try:
+            parsed_args = _safely_parse_tool_args(message_part.args)
+            if parsed_args:
+                # Successfully parsed as dict - format nicely
+                args_str = ", ".join(f"{k}={repr(v)}" for k, v in parsed_args.items())
+                tool_call_str = f"{message_part.tool_name}({args_str})"
+            elif isinstance(message_part.args, str) and message_part.args:
+                # Non-empty string that didn't parse as JSON
+                # Check if it looks like JSON (starts with { or [) - if so, it's incomplete
+                stripped_args = message_part.args.strip()
+                if stripped_args.startswith(("{", "[")):
+                    # Looks like incomplete JSON - log warning and show empty parens
+                    args_preview = (
+                        stripped_args[:100] + "..."
+                        if len(stripped_args) > 100
+                        else stripped_args
+                    )
+                    stack_trace = "".join(traceback.format_stack())
+                    logger.warning(
+                        "ToolCallPart with unparseable args encountered during context extraction",
+                        extra={
+                            "tool_name": message_part.tool_name,
+                            "tool_call_id": message_part.tool_call_id,
+                            "args_preview": args_preview,
+                            "args_type": type(message_part.args).__name__,
+                            "stack_trace": stack_trace,
+                        },
+                    )
+                    tool_call_str = f"{message_part.tool_name}()"
+                else:
+                    # Plain string arg - display as-is
+                    tool_call_str = f"{message_part.tool_name}({message_part.args})"
+            else:
+                # No args
+                tool_call_str = f"{message_part.tool_name}()"
+            return f"<TOOL_CALL>\n{tool_call_str}\n</TOOL_CALL>"
+        except Exception as e:  # pragma: no cover - defensive catch-all
+            # If anything goes wrong, log full exception with stack trace
+            logger.error(
+                "Unexpected error processing ToolCallPart",
+                exc_info=True,
+                extra={
+                    "tool_name": message_part.tool_name,
+                    "tool_call_id": message_part.tool_call_id,
+                    "error": str(e),
+                },
+            )
+            return f"<TOOL_CALL>\n{message_part.tool_name}()\n</TOOL_CALL>"
 
     elif isinstance(message_part, BuiltinToolCallPart):
         return f"<BUILTIN_TOOL_CALL>\n{message_part.tool_name}\n</BUILTIN_TOOL_CALL>"

shotgun/agents/history/history_processors.py

@@ -1,7 +1,9 @@
 """History processors for managing conversation history in Shotgun agents."""
 
+from collections.abc import Awaitable, Callable
 from typing import TYPE_CHECKING, Any, Protocol
 
+from anthropic import APIStatusError
 from pydantic_ai import ModelSettings
 from pydantic_ai.messages import (
     ModelMessage,
@@ -14,6 +16,7 @@ from pydantic_ai.messages import (
 from shotgun.agents.llm import shotgun_model_request
 from shotgun.agents.messages import AgentSystemPrompt, SystemStatusPrompt
 from shotgun.agents.models import AgentDeps
+from shotgun.exceptions import ContextSizeLimitExceeded
 from shotgun.logging_config import get_logger
 from shotgun.posthog_telemetry import track_event
 from shotgun.prompts import PromptLoader
@@ -51,6 +54,86 @@ logger = get_logger(__name__)
 prompt_loader = PromptLoader()
 
 
+async def _safe_token_estimation(
+    estimation_func: Callable[..., Awaitable[int]],
+    model_name: str,
+    max_tokens: int,
+    *args: Any,
+    **kwargs: Any,
+) -> int:
+    """Safely estimate tokens with proper error handling.
+
+    Wraps token estimation functions to handle failures gracefully.
+    Only RuntimeError (from token counters) is wrapped in ContextSizeLimitExceeded.
+    Other errors (network, auth) are allowed to bubble up.
+
+    Args:
+        estimation_func: Async function that estimates tokens
+        model_name: Name of the model for error messages
+        max_tokens: Maximum tokens for the model
+        *args: Arguments to pass to estimation_func
+        **kwargs: Keyword arguments to pass to estimation_func
+
+    Returns:
+        Token count from estimation_func
+
+    Raises:
+        ContextSizeLimitExceeded: If token counting fails with RuntimeError
+        Exception: Any other exceptions from estimation_func
+    """
+    try:
+        return await estimation_func(*args, **kwargs)
+    except Exception as e:
+        # Log the error with full context
+        logger.warning(
+            f"Token counting failed for {model_name}",
+            extra={
+                "error_type": type(e).__name__,
+                "error_message": str(e),
+                "model": model_name,
+            },
+        )
+
+        # Token counting behavior with oversized context (verified via testing):
+        #
+        # 1. OpenAI/tiktoken:
+        #    - Successfully counts any size (tested with 752K tokens, no error)
+        #    - Library errors: ValueError, KeyError, AttributeError, SSLError (file/cache issues)
+        #    - Wrapped as: RuntimeError by our counter
+        #
+        # 2. Gemini/SentencePiece:
+        #    - Successfully counts any size (tested with 752K tokens, no error)
+        #    - Library errors: RuntimeError, IOError, TypeError (file/model loading issues)
+        #    - Wrapped as: RuntimeError by our counter
+        #
+        # 3. Anthropic API:
+        #    - Successfully counts large token counts (tested with 752K tokens, no error)
+        #    - Only enforces 32 MB request size limit (not token count)
+        #    - Raises: APIStatusError(413) with error type 'request_too_large' for 32MB+ requests
+        #    - Other API errors: APIConnectionError, RateLimitError, APIStatusError (4xx/5xx)
+        #    - Wrapped as: RuntimeError by our counter
+        #
+        # IMPORTANT: No provider raises errors for "too many tokens" during counting.
+        # Token count validation happens separately by comparing count to max_input_tokens.
+        #
+        # We wrap RuntimeError (library-level failures from tiktoken/sentencepiece).
+        # We also wrap Anthropic's 413 error (request exceeds 32 MB) as it indicates
+        # context is effectively too large and needs user action to reduce it.
+        if isinstance(e, RuntimeError):
+            raise ContextSizeLimitExceeded(
+                model_name=model_name, max_tokens=max_tokens
+            ) from e
+
+        # Check for Anthropic's 32 MB request size limit (APIStatusError with status 413)
+        if isinstance(e, APIStatusError) and e.status_code == 413:
+            raise ContextSizeLimitExceeded(
+                model_name=model_name, max_tokens=max_tokens
+            ) from e
+
+        # Re-raise other exceptions (network errors, auth failures, etc.)
+        raise
+
+
 def is_summary_part(part: Any) -> bool:
     """Check if a message part is a compacted summary."""
     return isinstance(part, TextPart) and part.content.startswith(SUMMARY_MARKER)
@@ -127,6 +210,7 @@ calculate_max_summarization_tokens = _calculate_max_summarization_tokens
 async def token_limit_compactor(
     ctx: ContextProtocol,
     messages: list[ModelMessage],
+    force: bool = False,
 ) -> list[ModelMessage]:
     """Compact message history based on token limits with incremental processing.
 
@@ -139,6 +223,7 @@ async def token_limit_compactor(
     Args:
         ctx: Run context with usage information and dependencies
         messages: Current conversation history
+        force: If True, force compaction even if below token threshold
 
     Returns:
         Compacted list of messages within token limits
@@ -155,9 +240,15 @@ async def token_limit_compactor(
 
     if last_summary_index is not None:
         # Check if post-summary conversation exceeds threshold for incremental compaction
-        post_summary_tokens = await estimate_post_summary_tokens(
-            messages, last_summary_index, deps.llm_model
+        post_summary_tokens = await _safe_token_estimation(
+            estimate_post_summary_tokens,
+            deps.llm_model.name,
+            model_max_tokens,
+            messages,
+            last_summary_index,
+            deps.llm_model,
         )
+
         post_summary_percentage = (
             (post_summary_tokens / max_tokens) * 100 if max_tokens > 0 else 0
         )
@@ -169,7 +260,7 @@ async def token_limit_compactor(
         )
 
         # Only do incremental compaction if post-summary conversation exceeds threshold
-        if post_summary_tokens < max_tokens:
+        if post_summary_tokens < max_tokens and not force:
             logger.debug(
                 f"Post-summary conversation under threshold ({post_summary_tokens} < {max_tokens}), "
                 f"keeping all {len(messages)} messages"
@@ -340,6 +431,7 @@ async def token_limit_compactor(
             else 0
         )
 
+        # Track incremental compaction with simple metrics (fast, no token counting)
         track_event(
             "context_compaction_triggered",
             {
@@ -352,6 +444,10 @@ async def token_limit_compactor(
                 "agent_mode": deps.agent_mode.value
                 if hasattr(deps, "agent_mode") and deps.agent_mode
                 else "unknown",
+                # Model and provider info (no computation needed)
+                "model_name": deps.llm_model.name.value,
+                "provider": deps.llm_model.provider.value,
+                "key_provider": deps.llm_model.key_provider.value,
             },
         )
 
@@ -359,7 +455,14 @@ async def token_limit_compactor(
 
     else:
         # Check if total conversation exceeds threshold for full compaction
-        total_tokens = await estimate_tokens_from_messages(messages, deps.llm_model)
+        total_tokens = await _safe_token_estimation(
+            estimate_tokens_from_messages,
+            deps.llm_model.name,
+            model_max_tokens,
+            messages,
+            deps.llm_model,
+        )
+
         total_percentage = (total_tokens / max_tokens) * 100 if max_tokens > 0 else 0
 
         logger.debug(
@@ -368,7 +471,7 @@ async def token_limit_compactor(
         )
 
         # Only do full compaction if total conversation exceeds threshold
-        if total_tokens < max_tokens:
+        if total_tokens < max_tokens and not force:
             logger.debug(
                 f"Total conversation under threshold ({total_tokens} < {max_tokens}), "
                 f"keeping all {len(messages)} messages"
@@ -468,6 +571,7 @@ async def _full_compaction(
     tokens_before = current_tokens  # Already calculated above
     tokens_after = summary_usage.output_tokens if summary_usage else 0
 
+    # Track full compaction with simple metrics (fast, no token counting)
     track_event(
         "context_compaction_triggered",
         {
@@ -480,6 +584,10 @@ async def _full_compaction(
             "agent_mode": deps.agent_mode.value
             if hasattr(deps, "agent_mode") and deps.agent_mode
             else "unknown",
+            # Model and provider info (no computation needed)
+            "model_name": deps.llm_model.name.value,
+            "provider": deps.llm_model.provider.value,
+            "key_provider": deps.llm_model.key_provider.value,
         },
     )
 

shotgun/agents/history/token_counting/anthropic.py

@@ -1,5 +1,6 @@
 """Anthropic token counting using official client."""
 
+import logfire
 from pydantic_ai.messages import ModelMessage
 
 from shotgun.agents.config.models import KeyProvider
@@ -49,7 +50,15 @@ class AnthropicTokenCounter(TokenCounter):
                     f"Initialized async Anthropic token counter for {model_name} via direct API"
                 )
         except Exception as e:
-            raise RuntimeError("Failed to initialize Anthropic async client") from e
+            logfire.exception(
+                f"Failed to initialize Anthropic token counter for {model_name}",
+                model_name=model_name,
+                key_provider=key_provider.value,
+                exception_type=type(e).__name__,
+            )
+            raise RuntimeError(
+                f"Failed to initialize Anthropic async client for {model_name}: {type(e).__name__}: {str(e)}"
+            ) from e
 
     async def count_tokens(self, text: str) -> int:
         """Count tokens using Anthropic's official API (async).
@@ -63,16 +72,39 @@ class AnthropicTokenCounter(TokenCounter):
         Raises:
             RuntimeError: If API call fails
         """
+        # Handle empty text to avoid unnecessary API calls
+        # Anthropic API requires non-empty content, so we need a strict check
+        if not text or not text.strip():
+            return 0
+
+        # Additional validation: ensure the text has actual content
+        # Some edge cases might have only whitespace or control characters
+        cleaned_text = text.strip()
+        if not cleaned_text:
+            return 0
+
         try:
             # Anthropic API expects messages format and model parameter
             # Use await with async client
             result = await self.client.messages.count_tokens(
-                messages=[{"role": "user", "content": text}], model=self.model_name
+                messages=[{"role": "user", "content": cleaned_text}],
+                model=self.model_name,
             )
             return result.input_tokens
         except Exception as e:
+            # Create a preview of the text for logging (truncated to avoid huge logs)
+            text_preview = text[:100] + "..." if len(text) > 100 else text
+
+            logfire.exception(
+                f"Anthropic token counting failed for {self.model_name}",
+                model_name=self.model_name,
+                text_length=len(text),
+                text_preview=text_preview,
+                exception_type=type(e).__name__,
+                exception_message=str(e),
+            )
             raise RuntimeError(
-                f"Anthropic token counting API failed for {self.model_name}"
+                f"Anthropic token counting API failed for {self.model_name}: {type(e).__name__}: {str(e)}"
             ) from e
 
     async def count_message_tokens(self, messages: list[ModelMessage]) -> int:
@@ -87,5 +119,9 @@ class AnthropicTokenCounter(TokenCounter):
         Raises:
             RuntimeError: If token counting fails
         """
+        # Handle empty message list early
+        if not messages:
+            return 0
+
         total_text = extract_text_from_messages(messages)
         return await self.count_tokens(total_text)

shotgun/agents/history/token_counting/base.py

@@ -56,12 +56,23 @@ def extract_text_from_messages(messages: list[ModelMessage]) -> str:
         if hasattr(message, "parts"):
             for part in message.parts:
                 if hasattr(part, "content") and isinstance(part.content, str):
-                    text_parts.append(part.content)
+                    # Only add non-empty content
+                    if part.content.strip():
+                        text_parts.append(part.content)
                 else:
                     # Handle non-text parts (tool calls, etc.)
-                    text_parts.append(str(part))
+                    part_str = str(part)
+                    if part_str.strip():
+                        text_parts.append(part_str)
         else:
             # Handle messages without parts
-            text_parts.append(str(message))
+            msg_str = str(message)
+            if msg_str.strip():
+                text_parts.append(msg_str)
+
+    # If no valid text parts found, return a minimal placeholder
+    # This ensures we never send completely empty content to APIs
+    if not text_parts:
+        return "."
 
     return "\n".join(text_parts)

shotgun/agents/history/token_counting/openai.py

@@ -57,9 +57,15 @@ class OpenAITokenCounter(TokenCounter):
         Raises:
             RuntimeError: If token counting fails
         """
+        # Handle empty text to avoid unnecessary encoding
+        if not text or not text.strip():
+            return 0
+
         try:
             return len(self.encoding.encode(text))
-        except Exception as e:
+        except BaseException as e:
+            # Must catch BaseException to handle PanicException from tiktoken's Rust layer
+            # which can occur with extremely long texts. Regular Exception won't catch it.
             raise RuntimeError(
                 f"Failed to count tokens for OpenAI model {self.model_name}"
             ) from e
@@ -76,5 +82,9 @@ class OpenAITokenCounter(TokenCounter):
         Raises:
             RuntimeError: If token counting fails
         """
+        # Handle empty message list early
+        if not messages:
+            return 0
+
         total_text = extract_text_from_messages(messages)
         return await self.count_tokens(total_text)

shotgun/agents/history/token_counting/sentencepiece_counter.py

@@ -88,6 +88,10 @@ class SentencePieceTokenCounter(TokenCounter):
         Raises:
             RuntimeError: If token counting fails
         """
+        # Handle empty text to avoid unnecessary tokenization
+        if not text or not text.strip():
+            return 0
+
         await self._ensure_tokenizer()
 
         if self.sp is None:
@@ -115,5 +119,9 @@ class SentencePieceTokenCounter(TokenCounter):
         Raises:
             RuntimeError: If token counting fails
         """
+        # Handle empty message list early
+        if not messages:
+            return 0
+
         total_text = extract_text_from_messages(messages)
         return await self.count_tokens(total_text)

shotgun/agents/history/token_counting/tokenizer_cache.py

@@ -3,6 +3,7 @@
 import hashlib
 from pathlib import Path
 
+import aiofiles
 import httpx
 
 from shotgun.logging_config import get_logger
@@ -78,7 +79,8 @@ async def download_gemini_tokenizer() -> Path:
 
         # Atomic write: write to temp file first, then rename
         temp_path = cache_path.with_suffix(".tmp")
-        temp_path.write_bytes(content)
+        async with aiofiles.open(temp_path, "wb") as f:
+            await f.write(content)
         temp_path.rename(cache_path)
 
         logger.info(f"Gemini tokenizer downloaded and cached at {cache_path}")

shotgun/agents/history/token_counting/utils.py

@@ -44,9 +44,6 @@ def get_token_counter(model_config: ModelConfig) -> TokenCounter:
 
     # Return cached instance if available
     if cache_key in _token_counter_cache:
-        logger.debug(
-            f"Reusing cached token counter for {model_config.provider.value}:{model_config.name}"
-        )
         return _token_counter_cache[cache_key]
 
     # Create new instance and cache it

shotgun/agents/models.py

@@ -19,6 +19,30 @@ if TYPE_CHECKING:
     from shotgun.codebase.service import CodebaseService
 
 
+class AgentResponse(BaseModel):
+    """Structured response from an agent with optional clarifying questions.
+
+    This model provides a consistent response format for all agents:
+    - response: The main response text (can be empty if only asking questions)
+    - clarifying_questions: Optional list of questions to ask the user
+
+    When clarifying_questions is provided, the agent expects to receive
+    answers before continuing its work. This replaces the ask_questions tool.
+    """
+
+    response: str = Field(
+        description="The agent's response text. Always respond with some text summarizing what happened, whats next, etc.",
+    )
+    clarifying_questions: list[str] | None = Field(
+        default=None,
+        description="""
+Optional list of clarifying questions to ask the user.
+- Single question: Shown as a non-blocking suggestion (user can answer or continue with other prompts)
+- Multiple questions (2+): Asked sequentially in Q&A mode (blocks input until all answered or cancelled)
+""",
+    )
+
+
 class AgentType(StrEnum):
     """Enumeration for available agent types."""
 
@@ -73,6 +97,30 @@ class UserQuestion(BaseModel):
     )
 
 
+class MultipleUserQuestions(BaseModel):
+    """Multiple questions to ask the user sequentially."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    questions: list[str] = Field(
+        description="List of questions to ask the user",
+    )
+    current_index: int = Field(
+        default=0,
+        description="Current question index being asked",
+    )
+    answers: list[str] = Field(
+        default_factory=list,
+        description="Accumulated answers from the user",
+    )
+    tool_call_id: str = Field(
+        description="Tool call id",
+    )
+    result: Future[UserAnswer] = Field(
+        description="Future that will contain all answers formatted as Q&A pairs"
+    )
+
+
 class AgentRuntimeOptions(BaseModel):
     """User interface options for agents."""
 
@@ -100,9 +148,9 @@ class AgentRuntimeOptions(BaseModel):
         description="Maximum number of iterations for agent loops",
     )
 
-    queue: Queue[UserQuestion] = Field(
+    queue: Queue[UserQuestion | MultipleUserQuestions] = Field(
         default_factory=Queue,
-        description="Queue for storing user responses",
+        description="Queue for storing user questions (single or multiple)",
     )
 
     tasks: list[Future[UserAnswer]] = Field(