klaude-code 2.2.0__py3-none-any.whl → 2.3.0__py3-none-any.whl

klaude_code/config/sub_agent_model_helper.py

@@ -0,0 +1,217 @@
+"""Helper for sub-agent model availability and selection logic."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from klaude_code.protocol.sub_agent import (
+    AVAILABILITY_IMAGE_MODEL,
+    SubAgentProfile,
+    get_sub_agent_profile,
+    get_sub_agent_profile_by_tool,
+    iter_sub_agent_profiles,
+    sub_agent_tool_names,
+)
+from klaude_code.protocol.tools import SubAgentType
+
+if TYPE_CHECKING:
+    from klaude_code.config.config import Config, ModelEntry
+
+
+@dataclass
+class SubAgentModelInfo:
+    """Sub-agent and its current model configuration."""
+
+    profile: SubAgentProfile
+    # Explicitly configured model selector (from config), if any.
+    configured_model: str | None
+
+    # Effective model name used by this sub-agent.
+    # - When configured_model is set: equals configured_model.
+    # - When requirement-based default applies (e.g. ImageGen): resolved model.
+    # - When inheriting from main agent: None.
+    effective_model: str | None
+
+
+@dataclass(frozen=True, slots=True)
+class EmptySubAgentModelBehavior:
+    """Human-facing description for an unset (empty) sub-agent model config."""
+
+    # Summary text for UI (kept UI-framework agnostic).
+    description: str
+
+    # Best-effort resolved model name (if any). For ImageGen this is usually the
+    # first available image model; for other sub-agents it's the main model.
+    resolved_model_name: str | None
+
+
+class SubAgentModelHelper:
+    """Centralized logic for sub-agent availability and model selection."""
+
+    def __init__(self, config: Config) -> None:
+        self._config = config
+
+    def check_availability_requirement(self, requirement: str | None) -> bool:
+        """Check if a sub-agent's availability requirement is met.
+
+        Args:
+            requirement: The availability requirement constant (e.g., AVAILABILITY_IMAGE_MODEL).
+
+        Returns:
+            True if the requirement is met or if there's no requirement.
+        """
+        if requirement is None:
+            return True
+
+        if requirement == AVAILABILITY_IMAGE_MODEL:
+            return self._config.has_available_image_model()
+
+        return True
+
+    def resolve_model_for_requirement(self, requirement: str | None) -> str | None:
+        """Resolve the model name for a given availability requirement.
+
+        Args:
+            requirement: The availability requirement constant.
+
+        Returns:
+            The model name if found, None otherwise.
+        """
+        if requirement == AVAILABILITY_IMAGE_MODEL:
+            return self._config.get_first_available_image_model()
+        return None
+
+    def resolve_default_model_override(self, sub_agent_type: str) -> str | None:
+        """Resolve the default model override for a sub-agent when unset.
+
+        Returns:
+            - None for sub-agents that default to inheriting the main agent.
+            - A model name for sub-agents that require a dedicated model (e.g. ImageGen).
+
+        Note: This intentionally ignores any explicit user config; callers use this
+        when they want the *unset* behavior.
+        """
+
+        profile = get_sub_agent_profile(sub_agent_type)
+        if profile.availability_requirement is None:
+            return None
+        return self.resolve_model_for_requirement(profile.availability_requirement)
+
+    def describe_empty_model_config_behavior(
+        self,
+        sub_agent_type: str,
+        *,
+        main_model_name: str,
+    ) -> EmptySubAgentModelBehavior:
+        """Describe what happens when a sub-agent model is not configured.
+
+        Most sub-agents default to inheriting the main model.
+
+        Sub-agents with an availability requirement (e.g. ImageGen) do NOT
+        inherit from the main model; instead they auto-resolve a suitable model
+        (currently: the first available image model).
+        """
+
+        profile = get_sub_agent_profile(sub_agent_type)
+
+        requirement = profile.availability_requirement
+        if requirement is None:
+            return EmptySubAgentModelBehavior(
+                description=f"inherit from main agent: {main_model_name}",
+                resolved_model_name=main_model_name,
+            )
+
+        resolved = self.resolve_model_for_requirement(requirement)
+        if requirement == AVAILABILITY_IMAGE_MODEL:
+            if resolved:
+                return EmptySubAgentModelBehavior(
+                    description=f"auto-select first available image model: {resolved}",
+                    resolved_model_name=resolved,
+                )
+            return EmptySubAgentModelBehavior(
+                description="auto-select first available image model",
+                resolved_model_name=None,
+            )
+
+        if resolved:
+            return EmptySubAgentModelBehavior(
+                description=f"auto-select model for requirement '{requirement}': {resolved}",
+                resolved_model_name=resolved,
+            )
+        return EmptySubAgentModelBehavior(
+            description=f"auto-select model for requirement '{requirement}'",
+            resolved_model_name=None,
+        )
+
+    def get_available_sub_agents(self) -> list[SubAgentModelInfo]:
+        """Return all available sub-agents with their current model config.
+
+        Only returns sub-agents that:
+        1. Are enabled by default
+        2. Have their availability requirements met
+
+        For sub-agents without explicit config, resolves model based on availability_requirement.
+        """
+        result: list[SubAgentModelInfo] = []
+        for profile in iter_sub_agent_profiles(enabled_only=True):
+            if not self.check_availability_requirement(profile.availability_requirement):
+                continue
+            configured_model = self._config.sub_agent_models.get(profile.name)
+            effective_model = configured_model
+            if not effective_model and profile.availability_requirement:
+                effective_model = self.resolve_model_for_requirement(profile.availability_requirement)
+            result.append(
+                SubAgentModelInfo(
+                    profile=profile,
+                    configured_model=configured_model,
+                    effective_model=effective_model,
+                )
+            )
+        return result
+
+    def get_selectable_models(self, sub_agent_type: str) -> list[ModelEntry]:
+        """Return selectable models for a specific sub-agent type.
+
+        For sub-agents with availability_requirement (e.g., ImageGen):
+        - Only returns models matching the requirement (e.g., image models)
+
+        For other sub-agents:
+        - Returns all available models
+        """
+        profile = get_sub_agent_profile(sub_agent_type)
+        all_models = self._config.iter_model_entries(only_available=True)
+
+        if profile.availability_requirement == AVAILABILITY_IMAGE_MODEL:
+            return [m for m in all_models if m.model_params.modalities and "image" in m.model_params.modalities]
+
+        return all_models
+
+    def get_enabled_sub_agent_tool_names(self) -> list[str]:
+        """Return sub-agent tool names that should be added to main agent's tool list."""
+        result: list[str] = []
+        for name in sub_agent_tool_names(enabled_only=True):
+            profile = get_sub_agent_profile_by_tool(name)
+            if profile is not None and self.check_availability_requirement(profile.availability_requirement):
+                result.append(name)
+        return result
+
+    def get_enabled_sub_agent_types(self) -> set[SubAgentType]:
+        """Return set of sub-agent types that are enabled and available."""
+        enabled: set[SubAgentType] = set()
+        for name in sub_agent_tool_names(enabled_only=True):
+            profile = get_sub_agent_profile_by_tool(name)
+            if profile is not None and self.check_availability_requirement(profile.availability_requirement):
+                enabled.add(profile.name)
+        return enabled
+
+    def build_sub_agent_client_configs(self) -> dict[SubAgentType, str]:
+        """Return model names for each sub-agent that needs a dedicated client."""
+        result: dict[SubAgentType, str] = {}
+        for profile in iter_sub_agent_profiles():
+            model_name = self._config.sub_agent_models.get(profile.name)
+            if not model_name and profile.availability_requirement:
+                model_name = self.resolve_model_for_requirement(profile.availability_requirement)
+            if model_name:
+                result[profile.name] = model_name
+        return result

klaude_code/const.py

@@ -123,7 +123,7 @@ TAB_EXPAND_WIDTH = 8  # Tab expansion width for text rendering
 DIFF_PREFIX_WIDTH = 4  # Width of line number prefix in diff display
 MAX_DIFF_LINES = 500  # Maximum lines to show in diff output
 INVALID_TOOL_CALL_MAX_LENGTH = 200  # Maximum length for invalid tool call display
-TRUNCATE_DISPLAY_MAX_LINE_LENGTH = 200  # Maximum line length for truncated display
+TRUNCATE_DISPLAY_MAX_LINE_LENGTH = 500  # Maximum line length for truncated display
 TRUNCATE_DISPLAY_MAX_LINES = 4  # Maximum lines for truncated display
 MIN_HIDDEN_LINES_FOR_INDICATOR = 5  # Minimum hidden lines before showing truncation indicator
 SUB_AGENT_RESULT_MAX_LINES = 10  # Maximum lines for sub-agent result display

klaude_code/core/agent_profile.py

@@ -7,8 +7,12 @@ from dataclasses import dataclass
 from functools import cache
 from importlib.resources import files
 from pathlib import Path
-from typing import Any, Protocol
+from typing import TYPE_CHECKING, Any, Protocol
 
+if TYPE_CHECKING:
+    from klaude_code.config.config import Config
+
+from klaude_code.config.sub_agent_model_helper import SubAgentModelHelper
 from klaude_code.core.reminders import (
     at_file_reader_reminder,
     empty_todo_reminder,
@@ -23,7 +27,7 @@ from klaude_code.core.tool.report_back_tool import ReportBackTool
 from klaude_code.core.tool.tool_registry import get_tool_schemas
 from klaude_code.llm import LLMClientABC
 from klaude_code.protocol import llm_param, message, tools
-from klaude_code.protocol.sub_agent import get_sub_agent_profile, sub_agent_tool_names
+from klaude_code.protocol.sub_agent import get_sub_agent_profile
 from klaude_code.session import Session
 
 type Reminder = Callable[[Session], Awaitable[message.DeveloperMessage | None]]
@@ -169,12 +173,14 @@ def load_system_prompt(
 def load_agent_tools(
     model_name: str,
     sub_agent_type: tools.SubAgentType | None = None,
+    config: Config | None = None,
 ) -> list[llm_param.ToolSchema]:
     """Get tools for an agent based on model and agent type.
 
     Args:
         model_name: The model name.
         sub_agent_type: If None, returns main agent tools. Otherwise returns sub-agent tools.
+        config: Config for checking sub-agent availability (e.g., image model availability).
     """
 
     if sub_agent_type is not None:
@@ -183,13 +189,20 @@ def load_agent_tools(
 
     # Main agent tools
     if "gpt-5" in model_name:
-        tool_names = [tools.BASH, tools.READ, tools.APPLY_PATCH, tools.UPDATE_PLAN]
+        tool_names: list[str] = [tools.BASH, tools.READ, tools.APPLY_PATCH, tools.UPDATE_PLAN]
     elif "gemini-3" in model_name:
         tool_names = [tools.BASH, tools.READ, tools.EDIT, tools.WRITE]
     else:
         tool_names = [tools.BASH, tools.READ, tools.EDIT, tools.WRITE, tools.TODO_WRITE]
 
-    tool_names.extend(sub_agent_tool_names(enabled_only=True, model_name=model_name))
+    if config is not None:
+        helper = SubAgentModelHelper(config)
+        tool_names.extend(helper.get_enabled_sub_agent_tool_names())
+    else:
+        from klaude_code.protocol.sub_agent import sub_agent_tool_names
+
+        tool_names.extend(sub_agent_tool_names(enabled_only=True))
+
     tool_names.extend([tools.MERMAID])
     # tool_names.extend([tools.MEMORY])
     return get_tool_schemas(tool_names)
@@ -249,10 +262,17 @@ class ModelProfileProvider(Protocol):
         output_schema: dict[str, Any] | None = None,
     ) -> AgentProfile: ...
 
+    def enabled_sub_agent_types(self) -> set[tools.SubAgentType]:
+        """Return set of sub-agent types enabled for this provider."""
+        ...
+
 
 class DefaultModelProfileProvider(ModelProfileProvider):
     """Default provider backed by global prompts/tool/reminder registries."""
 
+    def __init__(self, config: Config | None = None) -> None:
+        self._config = config
+
     def build_profile(
         self,
         llm_client: LLMClientABC,
@@ -264,13 +284,25 @@ class DefaultModelProfileProvider(ModelProfileProvider):
         profile = AgentProfile(
             llm_client=llm_client,
             system_prompt=load_system_prompt(model_name, llm_client.protocol, sub_agent_type),
-            tools=load_agent_tools(model_name, sub_agent_type),
+            tools=load_agent_tools(model_name, sub_agent_type, config=self._config),
             reminders=load_agent_reminders(model_name, sub_agent_type),
         )
         if output_schema:
             return with_structured_output(profile, output_schema)
         return profile
 
+    def enabled_sub_agent_types(self) -> set[tools.SubAgentType]:
+        if self._config is None:
+            from klaude_code.protocol.sub_agent import get_sub_agent_profile_by_tool, sub_agent_tool_names
+
+            return {
+                profile.name
+                for name in sub_agent_tool_names(enabled_only=True)
+                if (profile := get_sub_agent_profile_by_tool(name)) is not None
+            }
+        helper = SubAgentModelHelper(self._config)
+        return helper.get_enabled_sub_agent_types()
+
 
 class VanillaModelProfileProvider(ModelProfileProvider):
     """Provider that strips prompts, reminders, and tools for vanilla mode."""
@@ -293,6 +325,9 @@ class VanillaModelProfileProvider(ModelProfileProvider):
             return with_structured_output(profile, output_schema)
         return profile
 
+    def enabled_sub_agent_types(self) -> set[tools.SubAgentType]:
+        return set()
+
 
 class NanoBananaModelProfileProvider(ModelProfileProvider):
     """Provider for the Nano Banana image generation model.
@@ -317,3 +352,6 @@ class NanoBananaModelProfileProvider(ModelProfileProvider):
         if output_schema:
             return with_structured_output(profile, output_schema)
         return profile
+
+    def enabled_sub_agent_types(self) -> set[tools.SubAgentType]:
+        return set()

klaude_code/core/executor.py

@@ -15,6 +15,7 @@ from dataclasses import dataclass
 from pathlib import Path
 
 from klaude_code.config import load_config
+from klaude_code.config.sub_agent_model_helper import SubAgentModelHelper
 from klaude_code.core.agent import Agent
 from klaude_code.core.agent_profile import DefaultModelProfileProvider, ModelProfileProvider
 from klaude_code.core.manager import LLMClients, SubAgentManager
@@ -109,6 +110,15 @@ class AgentRuntime:
     def current_agent(self) -> Agent | None:
         return self._agent
 
+    def _get_sub_agent_models(self) -> dict[str, LLMConfigParameter]:
+        """Build a dict of sub-agent type to LLMConfigParameter."""
+        enabled = self._model_profile_provider.enabled_sub_agent_types()
+        return {
+            sub_agent_type: client.get_llm_config()
+            for sub_agent_type, client in self._llm_clients.sub_clients.items()
+            if sub_agent_type in enabled
+        }
+
     async def ensure_agent(self, session_id: str | None = None) -> Agent:
         """Return the active agent, creating or loading a session as needed."""
 
@@ -135,6 +145,7 @@ class AgentRuntime:
                 session_id=session.id,
                 work_dir=str(session.work_dir),
                 llm_config=self._llm_clients.main.get_llm_config(),
+                sub_agent_models=self._get_sub_agent_models(),
             )
         )
 
@@ -200,6 +211,7 @@ class AgentRuntime:
                 session_id=agent.session.id,
                 work_dir=str(agent.session.work_dir),
                 llm_config=self._llm_clients.main.get_llm_config(),
+                sub_agent_models=self._get_sub_agent_models(),
             )
         )
 
@@ -223,6 +235,7 @@ class AgentRuntime:
                 session_id=target_session.id,
                 work_dir=str(target_session.work_dir),
                 llm_config=self._llm_clients.main.get_llm_config(),
+                sub_agent_models=self._get_sub_agent_models(),
             )
         )
 
@@ -406,6 +419,15 @@ class ExecutorContext:
         """Emit an event to the UI display system."""
         await self.event_queue.put(event)
 
+    def _get_sub_agent_models(self) -> dict[str, LLMConfigParameter]:
+        """Build a dict of sub-agent type to LLMConfigParameter."""
+        enabled = self.model_profile_provider.enabled_sub_agent_types()
+        return {
+            sub_agent_type: client.get_llm_config()
+            for sub_agent_type, client in self.llm_clients.sub_clients.items()
+            if sub_agent_type in enabled
+        }
+
     def current_session_id(self) -> str | None:
         """Return the primary active session id, if any.
 
@@ -455,6 +477,7 @@ class ExecutorContext:
                     llm_config=llm_config,
                     work_dir=str(agent.session.work_dir),
                     show_klaude_code_info=False,
+                    show_sub_agent_models=False,
                 )
             )
 
@@ -501,9 +524,61 @@ class ExecutorContext:
                     work_dir=str(agent.session.work_dir),
                     llm_config=agent.profile.llm_client.get_llm_config(),
                     show_klaude_code_info=False,
+                    show_sub_agent_models=False,
                 )
             )
 
+    async def handle_change_sub_agent_model(self, operation: op.ChangeSubAgentModelOperation) -> None:
+        """Handle a change sub-agent model operation."""
+        agent = await self._agent_runtime.ensure_agent(operation.session_id)
+        config = load_config()
+
+        helper = SubAgentModelHelper(config)
+
+        sub_agent_type = operation.sub_agent_type
+        model_name = operation.model_name
+
+        if model_name is None:
+            # Clear explicit override and revert to sub-agent default behavior.
+            behavior = helper.describe_empty_model_config_behavior(
+                sub_agent_type,
+                main_model_name=self.llm_clients.main.model_name,
+            )
+
+            resolved = helper.resolve_default_model_override(sub_agent_type)
+            if resolved is None:
+                # Default: inherit from main client.
+                self.llm_clients.sub_clients.pop(sub_agent_type, None)
+            else:
+                # Default: use a dedicated model (e.g. first available image model).
+                llm_config = config.get_model_config(resolved)
+                new_client = create_llm_client(llm_config)
+                self.llm_clients.sub_clients[sub_agent_type] = new_client
+
+            display_model = f"({behavior.description})"
+        else:
+            # Create new client for the sub-agent
+            llm_config = config.get_model_config(model_name)
+            new_client = create_llm_client(llm_config)
+            self.llm_clients.sub_clients[sub_agent_type] = new_client
+            display_model = new_client.model_name
+
+        if operation.save_as_default:
+            if model_name is None:
+                # Remove from config to inherit
+                config.sub_agent_models.pop(sub_agent_type, None)
+            else:
+                config.sub_agent_models[sub_agent_type] = model_name
+            await config.save()
+
+        saved_note = " (saved in ~/.klaude/klaude-config.yaml)" if operation.save_as_default else ""
+        developer_item = message.DeveloperMessage(
+            parts=message.text_parts_from_str(f"{sub_agent_type} model: {display_model}{saved_note}"),
+            ui_extra=model.build_command_output_extra(commands.CommandName.SUB_AGENT_MODEL),
+        )
+        agent.session.append_history([developer_item])
+        await self.emit_event(events.DeveloperMessageEvent(session_id=agent.session.id, item=developer_item))
+
     async def handle_clear_session(self, operation: op.ClearSessionOperation) -> None:
         await self._agent_runtime.clear_session(operation.session_id)
 

klaude_code/core/manager/llm_clients_builder.py

@@ -3,11 +3,11 @@
 from __future__ import annotations
 
 from klaude_code.config import Config
+from klaude_code.config.sub_agent_model_helper import SubAgentModelHelper
 from klaude_code.core.manager.llm_clients import LLMClients
 from klaude_code.llm.client import LLMClientABC
 from klaude_code.llm.registry import create_llm_client
 from klaude_code.log import DebugType, log_debug
-from klaude_code.protocol.sub_agent import iter_sub_agent_profiles
 from klaude_code.protocol.tools import SubAgentType
 
 
@@ -15,8 +15,15 @@ def build_llm_clients(
     config: Config,
     *,
     model_override: str | None = None,
+    skip_sub_agents: bool = False,
 ) -> LLMClients:
-    """Create an ``LLMClients`` bundle driven by application config."""
+    """Create an ``LLMClients`` bundle driven by application config.
+
+    Args:
+        config: Application configuration.
+        model_override: Override for the main model name.
+        skip_sub_agents: If True, skip initializing sub-agent clients (e.g., for vanilla/banana modes).
+    """
 
     # Resolve main agent LLM config
     model_name = model_override or config.main_model
@@ -32,17 +39,16 @@ def build_llm_clients(
     )
 
     main_client = create_llm_client(llm_config)
-    sub_clients: dict[SubAgentType, LLMClientABC] = {}
 
-    for profile in iter_sub_agent_profiles():
-        model_name = config.sub_agent_models.get(profile.name)
-        if not model_name:
-            continue
+    if skip_sub_agents:
+        return LLMClients(main=main_client)
 
-        if not profile.enabled_for_model(main_client.model_name):
-            continue
+    helper = SubAgentModelHelper(config)
+    sub_agent_configs = helper.build_sub_agent_client_configs()
 
-        sub_llm_config = config.get_model_config(model_name)
-        sub_clients[profile.name] = create_llm_client(sub_llm_config)
+    sub_clients: dict[SubAgentType, LLMClientABC] = {}
+    for sub_agent_type, sub_model_name in sub_agent_configs.items():
+        sub_llm_config = config.get_model_config(sub_model_name)
+        sub_clients[sub_agent_type] = create_llm_client(sub_llm_config)
 
     return LLMClients(main=main_client, sub_clients=sub_clients)

klaude_code/core/prompts/prompt-nano-banana.md

@@ -1 +1 @@
-You're a helpful art assistant
+You're a helpful assistant with capabilities to generate images and edit images.

klaude_code/core/tool/sub_agent_tool.py

@@ -31,11 +31,12 @@ class SubAgentTool(ToolABC):
     @classmethod
     def for_profile(cls, profile: SubAgentProfile) -> type[SubAgentTool]:
         """Create a tool class for a specific sub-agent profile."""
-        return type(
+        tool_cls = type(
             f"{profile.name}Tool",
             (SubAgentTool,),
             {"_profile": profile},
         )
+        return cast(type[SubAgentTool], tool_cls)
 
     @classmethod
     def metadata(cls) -> ToolMetadata:

klaude_code/llm/anthropic/client.py

@@ -16,6 +16,7 @@ from anthropic.types.beta.beta_raw_message_start_event import BetaRawMessageStar
 from anthropic.types.beta.beta_signature_delta import BetaSignatureDelta
 from anthropic.types.beta.beta_text_delta import BetaTextDelta
 from anthropic.types.beta.beta_thinking_delta import BetaThinkingDelta
+from anthropic.types.beta.beta_tool_choice_auto_param import BetaToolChoiceAutoParam
 from anthropic.types.beta.beta_tool_use_block import BetaToolUseBlock
 from anthropic.types.beta.message_create_params import MessageCreateParamsStreaming
 
@@ -82,12 +83,14 @@ def build_payload(
         # Prepend extra betas, avoiding duplicates
         betas = [b for b in extra_betas if b not in betas] + betas
 
+    tool_choice: BetaToolChoiceAutoParam = {
+        "type": "auto",
+        "disable_parallel_tool_use": False,
+    }
+
     payload: MessageCreateParamsStreaming = {
         "model": str(param.model),
-        "tool_choice": {
-            "type": "auto",
-            "disable_parallel_tool_use": False,
-        },
+        "tool_choice": tool_choice,
         "stream": True,
         "max_tokens": param.max_tokens or DEFAULT_MAX_TOKENS,
         "temperature": param.temperature or DEFAULT_TEMPERATURE,