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

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(),
             )
         )
 
@@ -151,27 +162,16 @@ class AgentRuntime:
 
     async def run_agent(self, operation: op.RunAgentOperation) -> None:
         agent = await self.ensure_agent(operation.session_id)
-
-        if operation.emit_user_message_event:
-            await self._emit_event(
-                events.UserMessageEvent(
-                    content=operation.input.text,
-                    session_id=agent.session.id,
-                    images=operation.input.images,
-                )
-            )
-
-        if operation.persist_user_input:
-            agent.session.append_history(
-                [
-                    message.UserMessage(
-                        parts=message.parts_from_text_and_images(
-                            operation.input.text,
-                            operation.input.images,
-                        )
+        agent.session.append_history(
+            [
+                message.UserMessage(
+                    parts=message.parts_from_text_and_images(
+                        operation.input.text,
+                        operation.input.images,
                     )
-                ]
-            )
+                )
+            ]
+        )
 
         existing_active = self._task_manager.get(operation.id)
         if existing_active is not None and not existing_active.task.done():
@@ -190,16 +190,19 @@ class AgentRuntime:
         new_session.model_thinking = agent.session.model_thinking
         agent.session = new_session
 
-        developer_item = message.DeveloperMessage(
-            parts=message.text_parts_from_str("started new conversation"),
-            ui_extra=model.build_command_output_extra(commands.CommandName.CLEAR),
+        await self._emit_event(
+            events.CommandOutputEvent(
+                session_id=agent.session.id,
+                command_name=commands.CommandName.CLEAR,
+                content="started new conversation",
+            )
         )
-        await self._emit_event(events.DeveloperMessageEvent(session_id=agent.session.id, item=developer_item))
         await self._emit_event(
             events.WelcomeEvent(
                 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 +226,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 +410,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.
 
@@ -438,12 +451,13 @@ class ExecutorContext:
 
         if operation.emit_switch_message:
             default_note = " (saved as default)" if operation.save_as_default else ""
-            developer_item = message.DeveloperMessage(
-                parts=message.text_parts_from_str(f"Switched to: {llm_config.model}{default_note}"),
-                ui_extra=model.build_command_output_extra(commands.CommandName.MODEL),
+            await self.emit_event(
+                events.CommandOutputEvent(
+                    session_id=agent.session.id,
+                    command_name=commands.CommandName.MODEL,
+                    content=f"Switched to: {llm_config.model_id}{default_note}",
+                )
             )
-            agent.session.append_history([developer_item])
-            await self.emit_event(events.DeveloperMessageEvent(session_id=agent.session.id, item=developer_item))
 
         if self._on_model_change is not None:
             self._on_model_change(llm_client_name)
@@ -455,6 +469,7 @@ class ExecutorContext:
                     llm_config=llm_config,
                     work_dir=str(agent.session.work_dir),
                     show_klaude_code_info=False,
+                    show_sub_agent_models=False,
                 )
             )
 
@@ -487,12 +502,13 @@ class ExecutorContext:
         new_status = _format_thinking_for_display(operation.thinking)
 
         if operation.emit_switch_message:
-            developer_item = message.DeveloperMessage(
-                parts=message.text_parts_from_str(f"Thinking changed: {current} -> {new_status}"),
-                ui_extra=model.build_command_output_extra(commands.CommandName.THINKING),
+            await self.emit_event(
+                events.CommandOutputEvent(
+                    session_id=agent.session.id,
+                    command_name=commands.CommandName.THINKING,
+                    content=f"Thinking changed: {current} -> {new_status}",
+                )
             )
-            agent.session.append_history([developer_item])
-            await self.emit_event(events.DeveloperMessageEvent(session_id=agent.session.id, item=developer_item))
 
         if operation.emit_welcome_event:
             await self.emit_event(
@@ -501,9 +517,62 @@ 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 ""
+        await self.emit_event(
+            events.CommandOutputEvent(
+                session_id=agent.session.id,
+                command_name=commands.CommandName.SUB_AGENT_MODEL,
+                content=f"{sub_agent_type} model: {display_model}{saved_note}",
+            )
+        )
+
     async def handle_clear_session(self, operation: op.ClearSessionOperation) -> None:
         await self._agent_runtime.clear_session(operation.session_id)
 
@@ -518,21 +587,24 @@ class ExecutorContext:
             await asyncio.to_thread(output_path.parent.mkdir, parents=True, exist_ok=True)
             await asyncio.to_thread(output_path.write_text, html_doc, "utf-8")
             await asyncio.to_thread(self._open_file, output_path)
-            developer_item = message.DeveloperMessage(
-                parts=message.text_parts_from_str(f"Session exported and opened: {output_path}"),
-                ui_extra=model.build_command_output_extra(commands.CommandName.EXPORT),
+            await self.emit_event(
+                events.CommandOutputEvent(
+                    session_id=agent.session.id,
+                    command_name=commands.CommandName.EXPORT,
+                    content=f"Session exported and opened: {output_path}",
+                )
             )
-            agent.session.append_history([developer_item])
-            await self.emit_event(events.DeveloperMessageEvent(session_id=agent.session.id, item=developer_item))
         except Exception as exc:  # pragma: no cover
             import traceback
 
-            developer_item = message.DeveloperMessage(
-                parts=message.text_parts_from_str(f"Failed to export session: {exc}\n{traceback.format_exc()}"),
-                ui_extra=model.build_command_output_extra(commands.CommandName.EXPORT, is_error=True),
+            await self.emit_event(
+                events.CommandOutputEvent(
+                    session_id=agent.session.id,
+                    command_name=commands.CommandName.EXPORT,
+                    content=f"Failed to export session: {exc}\n{traceback.format_exc()}",
+                    is_error=True,
+                )
             )
-            agent.session.append_history([developer_item])
-            await self.emit_event(events.DeveloperMessageEvent(session_id=agent.session.id, item=developer_item))
 
     def _resolve_export_output_path(self, raw: str | None, session: Session) -> Path:
         trimmed = (raw or "").strip()
@@ -626,12 +698,15 @@ class Executor:
             Unique submission ID for tracking
         """
 
-        submission = op.Submission(id=operation.id, operation=operation)
-        await self.submission_queue.put(submission)
+        if operation.id in self._completion_events:
+            raise RuntimeError(f"Submission already registered: {operation.id}")
 
-        # Create completion event for tracking
+        # Create completion event before queueing to avoid races.
         self._completion_events[operation.id] = asyncio.Event()
 
+        submission = op.Submission(id=operation.id, operation=operation)
+        await self.submission_queue.put(submission)
+
         log_debug(
             f"Submitted operation {operation.type} with ID {operation.id}",
             style="blue",
@@ -711,9 +786,16 @@ class Executor:
         if tasks_to_await:
             await asyncio.gather(*tasks_to_await, return_exceptions=True)
 
+        if self._background_tasks:
+            await asyncio.gather(*self._background_tasks, return_exceptions=True)
+            self._background_tasks.clear()
+
         # Clear the active task manager
         self.context.task_manager.clear()
 
+        for event in self._completion_events.values():
+            event.set()
+
         # Send EndOperation to wake up the start() loop
         try:
             end_operation = op.EndOperation()

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/file/diff_builder.py

@@ -54,24 +54,31 @@ def _build_file_diff(before: str, after: str, *, file_path: str) -> model.DiffFi
             elif tag == "replace":
                 old_block = before_lines[i1:i2]
                 new_block = after_lines[j1:j2]
-                max_len = max(len(old_block), len(new_block))
-                for idx in range(max_len):
-                    old_line = old_block[idx] if idx < len(old_block) else None
-                    new_line = new_block[idx] if idx < len(new_block) else None
-                    if old_line is not None and new_line is not None:
-                        remove_spans, add_spans = _diff_line_spans(old_line, new_line)
-                        lines.append(_remove_line(remove_spans))
-                        lines.append(_add_line(add_spans, new_line_no))
-                        stats_remove += 1
-                        stats_add += 1
-                        new_line_no += 1
-                    elif old_line is not None:
-                        lines.append(_remove_line([model.DiffSpan(op="equal", text=old_line)]))
-                        stats_remove += 1
-                    elif new_line is not None:
-                        lines.append(_add_line([model.DiffSpan(op="equal", text=new_line)], new_line_no))
-                        stats_add += 1
-                        new_line_no += 1
+
+                # Emit replacement blocks in unified-diff style: all removals first, then all additions.
+                # This matches VSCode's readability (--- then +++), while keeping per-line char spans.
+                remove_block: list[list[model.DiffSpan]] = []
+                add_block: list[list[model.DiffSpan]] = []
+
+                paired_len = min(len(old_block), len(new_block))
+                for idx in range(paired_len):
+                    remove_spans, add_spans = _diff_line_spans(old_block[idx], new_block[idx])
+                    remove_block.append(remove_spans)
+                    add_block.append(add_spans)
+
+                for old_line in old_block[paired_len:]:
+                    remove_block.append([model.DiffSpan(op="equal", text=old_line)])
+                for new_line in new_block[paired_len:]:
+                    add_block.append([model.DiffSpan(op="equal", text=new_line)])
+
+                for spans in remove_block:
+                    lines.append(_remove_line(spans))
+                    stats_remove += 1
+
+                for spans in add_block:
+                    lines.append(_add_line(spans, new_line_no))
+                    stats_add += 1
+                    new_line_no += 1
 
     return model.DiffFileDiff(
         file_path=file_path,

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
 
@@ -64,7 +65,7 @@ def build_payload(
         param: LLM call parameters.
         extra_betas: Additional beta flags to prepend to the betas list.
     """
-    messages = convert_history_to_input(param.input, param.model)
+    messages = convert_history_to_input(param.input, param.model_id)
     tools = convert_tool_schema(param.tools)
     system_messages = [msg for msg in param.input if isinstance(msg, message.SystemMessage)]
     system = convert_system_to_input(param.system, system_messages)
@@ -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,
-        },
+        "model": str(param.model_id),
+        "tool_choice": tool_choice,
         "stream": True,
         "max_tokens": param.max_tokens or DEFAULT_MAX_TOKENS,
         "temperature": param.temperature or DEFAULT_TEMPERATURE,
@@ -183,12 +186,12 @@ async def parse_anthropic_stream(
                 if accumulated_thinking:
                     metadata_tracker.record_token()
                     full_thinking = "".join(accumulated_thinking)
-                    parts.append(message.ThinkingTextPart(text=full_thinking, model_id=str(param.model)))
+                    parts.append(message.ThinkingTextPart(text=full_thinking, model_id=str(param.model_id)))
                     if pending_signature:
                         parts.append(
                             message.ThinkingSignaturePart(
                                 signature=pending_signature,
-                                model_id=str(param.model),
+                                model_id=str(param.model_id),
                                 format="anthropic",
                             )
                         )
@@ -221,7 +224,7 @@ async def parse_anthropic_stream(
                         max_tokens=param.max_tokens,
                     )
                 )
-                metadata_tracker.set_model_name(str(param.model))
+                metadata_tracker.set_model_name(str(param.model_id))
                 metadata_tracker.set_response_id(response_id)
                 raw_stop_reason = getattr(event, "stop_reason", None)
                 if isinstance(raw_stop_reason, str):