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

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,13 +15,20 @@ 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
     if model_name is None:
-        raise ValueError("No model specified. Use --model or --select-model to specify a model.")
+        raise ValueError("No model specified. Set main_model in the config or pass --model.")
     llm_config = config.get_model_config(model_name)
 
     log_debug(
@@ -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

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

klaude_code/core/tool/shell/command_safety.py

@@ -1,5 +1,4 @@
 import os
-import re
 import shlex
 
 
@@ -11,76 +10,6 @@ class SafetyCheckResult:
         self.error_msg = error_msg
 
 
-def _is_valid_sed_n_arg(s: str | None) -> bool:
-    if not s:
-        return False
-    # Matches: Np or M,Np where M,N are positive integers
-    return bool(re.fullmatch(r"\d+(,\d+)?p", s))
-
-
-def _is_safe_awk_program(program: str) -> SafetyCheckResult:
-    lowered = program.lower()
-
-    if "`" in program:
-        return SafetyCheckResult(False, "awk: backticks not allowed in program")
-    if "$(" in program:
-        return SafetyCheckResult(False, "awk: command substitution not allowed in program")
-    if "|&" in program:
-        return SafetyCheckResult(False, "awk: background pipeline not allowed in program")
-
-    if "system(" in lowered:
-        return SafetyCheckResult(False, "awk: system() call not allowed in program")
-
-    if re.search(r"(?<![|&>])\bprint\s*\|", program, re.IGNORECASE):
-        return SafetyCheckResult(False, "awk: piping output to external command not allowed")
-    if re.search(r"\bprintf\s*\|", program, re.IGNORECASE):
-        return SafetyCheckResult(False, "awk: piping output to external command not allowed")
-
-    return SafetyCheckResult(True)
-
-
-def _is_safe_awk_argv(argv: list[str]) -> SafetyCheckResult:
-    if len(argv) < 2:
-        return SafetyCheckResult(False, "awk: Missing program")
-
-    program: str | None = None
-
-    i = 1
-    while i < len(argv):
-        arg = argv[i]
-
-        if arg in {"-f", "--file", "--source"} or arg.startswith("-f"):
-            return SafetyCheckResult(False, "awk: -f/--file not allowed")
-
-        if arg in {"-e", "--exec"}:
-            if i + 1 >= len(argv):
-                return SafetyCheckResult(False, "awk: Missing program for -e")
-            script = argv[i + 1]
-            program_check = _is_safe_awk_program(script)
-            if not program_check.is_safe:
-                return program_check
-            if program is None:
-                program = script
-            i += 2
-            continue
-
-        if arg.startswith("-"):
-            i += 1
-            continue
-
-        if program is None:
-            program_check = _is_safe_awk_program(arg)
-            if not program_check.is_safe:
-                return program_check
-            program = arg
-        i += 1
-
-    if program is None:
-        return SafetyCheckResult(False, "awk: Missing program")
-
-    return SafetyCheckResult(True)
-
-
 def _is_safe_rm_argv(argv: list[str]) -> SafetyCheckResult:
     """Check safety of rm command arguments."""
     # Enforce strict safety rules for rm operands
@@ -217,112 +146,12 @@ def _is_safe_argv(argv: list[str]) -> SafetyCheckResult:
 
     cmd0 = argv[0]
 
-    # if _has_shell_redirection(argv):
-    #     return SafetyCheckResult(False, "Shell redirection and pipelines are not allowed in single commands")
-
-    # Special handling for rm to prevent dangerous operations
     if cmd0 == "rm":
         return _is_safe_rm_argv(argv)
 
-    # Special handling for trash to prevent dangerous operations
     if cmd0 == "trash":
         return _is_safe_trash_argv(argv)
 
-    if cmd0 == "find":
-        unsafe_opts = {
-            "-exec": "command execution",
-            "-execdir": "command execution",
-            "-ok": "interactive command execution",
-            "-okdir": "interactive command execution",
-            "-delete": "file deletion",
-            "-fls": "file output",
-            "-fprint": "file output",
-            "-fprint0": "file output",
-            "-fprintf": "formatted file output",
-        }
-        for arg in argv[1:]:
-            if arg in unsafe_opts:
-                return SafetyCheckResult(False, f"find: {unsafe_opts[arg]} option '{arg}' not allowed")
-        return SafetyCheckResult(True)
-
-    if cmd0 == "git":
-        sub = argv[1] if len(argv) > 1 else None
-        if not sub:
-            return SafetyCheckResult(False, "git: Missing subcommand")
-
-        # Allow most local git operations, but block remote operations
-        allowed_git_cmds = {
-            "add",
-            "branch",
-            "checkout",
-            "commit",
-            "config",
-            "diff",
-            "fetch",
-            "init",
-            "log",
-            "merge",
-            "mv",
-            "rebase",
-            "reset",
-            "restore",
-            "revert",
-            "rm",
-            "show",
-            "stash",
-            "status",
-            "switch",
-            "tag",
-            "clone",
-            "worktree",
-            "push",
-            "pull",
-            "remote",
-        }
-        if sub not in allowed_git_cmds:
-            return SafetyCheckResult(False, f"git: Subcommand '{sub}' not in allow list")
-        return SafetyCheckResult(True)
-
-    # Build tools and linters - allow all subcommands
-    if cmd0 in {
-        "cargo",
-        "uv",
-        "go",
-        "ruff",
-        "pyright",
-        "make",
-        "npm",
-        "pnpm",
-        "bun",
-    }:
-        return SafetyCheckResult(True)
-
-    if cmd0 == "sed":
-        # Allow sed -n patterns (line printing)
-        if len(argv) >= 3 and argv[1] == "-n" and _is_valid_sed_n_arg(argv[2]):
-            return SafetyCheckResult(True)
-        # Allow simple text replacement: sed 's/old/new/g' file
-        # or sed -i 's/old/new/g' file for in-place editing
-        if len(argv) >= 3:
-            # Find the sed script argument (usually starts with 's/')
-            for arg in argv[1:]:
-                if arg.startswith("s/") or arg.startswith("s|"):
-                    # Basic safety check: no command execution in replacement
-                    if ";" in arg:
-                        return SafetyCheckResult(False, f"sed: Command separator ';' not allowed in '{arg}'")
-                    if "`" in arg:
-                        return SafetyCheckResult(False, f"sed: Backticks not allowed in '{arg}'")
-                    if "$(" in arg:
-                        return SafetyCheckResult(False, f"sed: Command substitution not allowed in '{arg}'")
-                    return SafetyCheckResult(True)
-        return SafetyCheckResult(
-            False,
-            "sed: Only text replacement (s/old/new/) or line printing (-n 'Np') is allowed",
-        )
-
-    if cmd0 == "awk":
-        return _is_safe_awk_argv(argv)
-
     # Default allow when command is not explicitly restricted
     return SafetyCheckResult(True)
 
@@ -330,30 +159,16 @@ def _is_safe_argv(argv: list[str]) -> SafetyCheckResult:
 def is_safe_command(command: str) -> SafetyCheckResult:
     """Determine if a command is safe enough to run.
 
-    The check is intentionally lightweight: it blocks only a small set of
-    obviously dangerous patterns (rm/trash/git remotes, unsafe sed/awk,
-    find -exec/-delete, etc.) and otherwise lets the real shell surface
-    syntax errors (for example, unmatched quotes in complex multiline
-    scripts).
+    Only rm and trash commands are checked for safety. All other commands
+    are allowed by default.
     """
-
-    # Try to parse into an argv-style list first. If this fails (e.g. due
-    # to unterminated quotes in a complex heredoc), treat the command as
-    # safe here and let bash itself perform syntax checking instead of
-    # blocking execution pre-emptively.
     try:
         argv = shlex.split(command, posix=True)
     except ValueError:
-        # If we cannot reliably parse the command (e.g. due to unterminated
-        # quotes in a complex heredoc), treat it as safe here and let the
-        # real shell surface any syntax errors instead of blocking execution
-        # pre-emptively.
+        # If we cannot reliably parse the command, treat it as safe here
+        # and let the real shell surface any syntax errors
         return SafetyCheckResult(True)
 
-    # All further safety checks are done directly on the parsed argv via
-    # _is_safe_argv. We intentionally avoid trying to re-interpret complex
-    # shell sequences here and rely on the real shell to handle syntax.
-
     if not argv:
         return SafetyCheckResult(False, "Empty command")
 

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/core/turn.py

@@ -348,7 +348,7 @@ class TurnExecutor:
                         style="red",
                         debug_type=DebugType.RESPONSE,
                     )
-                case message.ToolCallStartItem() as msg:
+                case message.ToolCallStartDelta() as msg:
                     if thinking_active:
                         thinking_active = False
                         yield events.ThinkingEndEvent(

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,
@@ -169,7 +172,7 @@ async def parse_anthropic_stream(
                 match event.content_block:
                     case BetaToolUseBlock() as block:
                         metadata_tracker.record_token()
-                        yield message.ToolCallStartItem(
+                        yield message.ToolCallStartDelta(
                             response_id=response_id,
                             call_id=block.id,
                             name=block.name,

klaude_code/llm/anthropic/input.py

@@ -8,10 +8,13 @@ import json
 from typing import Literal, cast
 
 from anthropic.types.beta.beta_base64_image_source_param import BetaBase64ImageSourceParam
+from anthropic.types.beta.beta_content_block_param import BetaContentBlockParam
 from anthropic.types.beta.beta_image_block_param import BetaImageBlockParam
 from anthropic.types.beta.beta_message_param import BetaMessageParam
 from anthropic.types.beta.beta_text_block_param import BetaTextBlockParam
 from anthropic.types.beta.beta_tool_param import BetaToolParam
+from anthropic.types.beta.beta_tool_result_block_param import BetaToolResultBlockParam
+from anthropic.types.beta.beta_tool_use_block_param import BetaToolUseBlockParam
 from anthropic.types.beta.beta_url_image_source_param import BetaURLImageSourceParam
 
 from klaude_code.const import EMPTY_TOOL_OUTPUT_MESSAGE
@@ -60,29 +63,29 @@ def _user_message_to_message(
     blocks: list[BetaTextBlockParam | BetaImageBlockParam] = []
     for part in msg.parts:
         if isinstance(part, message.TextPart):
-            blocks.append({"type": "text", "text": part.text})
+            blocks.append(cast(BetaTextBlockParam, {"type": "text", "text": part.text}))
         elif isinstance(part, message.ImageURLPart):
             blocks.append(_image_part_to_block(part))
     if attachment.text:
-        blocks.append({"type": "text", "text": attachment.text})
+        blocks.append(cast(BetaTextBlockParam, {"type": "text", "text": attachment.text}))
     for image in attachment.images:
         blocks.append(_image_part_to_block(image))
     if not blocks:
-        blocks.append({"type": "text", "text": ""})
+        blocks.append(cast(BetaTextBlockParam, {"type": "text", "text": ""}))
     return {"role": "user", "content": blocks}
 
 
 def _tool_message_to_block(
     msg: message.ToolResultMessage,
     attachment: DeveloperAttachment,
-) -> dict[str, object]:
+) -> BetaToolResultBlockParam:
     """Convert a single tool result message to a tool_result block."""
     tool_content: list[BetaTextBlockParam | BetaImageBlockParam] = []
     merged_text = merge_reminder_text(
         msg.output_text or EMPTY_TOOL_OUTPUT_MESSAGE,
         attachment.text,
     )
-    tool_content.append({"type": "text", "text": merged_text})
+    tool_content.append(cast(BetaTextBlockParam, {"type": "text", "text": merged_text}))
     for image in [part for part in msg.parts if isinstance(part, message.ImageURLPart)]:
         tool_content.append(_image_part_to_block(image))
     for image in attachment.images:
@@ -95,7 +98,7 @@ def _tool_message_to_block(
     }
 
 
-def _tool_blocks_to_message(blocks: list[dict[str, object]]) -> BetaMessageParam:
+def _tool_blocks_to_message(blocks: list[BetaToolResultBlockParam]) -> BetaMessageParam:
     """Convert one or more tool_result blocks to a single user message."""
     return {
         "role": "user",
@@ -104,7 +107,7 @@ def _tool_blocks_to_message(blocks: list[dict[str, object]]) -> BetaMessageParam
 
 
 def _assistant_message_to_message(msg: message.AssistantMessage, model_name: str | None) -> BetaMessageParam:
-    content: list[dict[str, object]] = []
+    content: list[BetaContentBlockParam] = []
     current_thinking_content: str | None = None
     native_thinking_parts, degraded_thinking_texts = split_thinking_parts(msg, model_name)
     native_thinking_ids = {id(part) for part in native_thinking_parts}
@@ -113,7 +116,7 @@ def _assistant_message_to_message(msg: message.AssistantMessage, model_name: str
         nonlocal current_thinking_content
         if current_thinking_content is None:
             return
-        content.append({"type": "thinking", "thinking": current_thinking_content})
+        degraded_thinking_texts.append(current_thinking_content)
         current_thinking_content = None
 
     for part in msg.parts:
@@ -127,33 +130,47 @@ def _assistant_message_to_message(msg: message.AssistantMessage, model_name: str
                 continue
             if part.signature:
                 content.append(
-                    {
-                        "type": "thinking",
-                        "thinking": current_thinking_content or "",
-                        "signature": part.signature,
-                    }
+                    cast(
+                        BetaContentBlockParam,
+                        {
+                            "type": "thinking",
+                            "thinking": current_thinking_content or "",
+                            "signature": part.signature,
+                        },
+                    )
                 )
                 current_thinking_content = None
             continue
 
         _flush_thinking()
         if isinstance(part, message.TextPart):
-            content.append({"type": "text", "text": part.text})
+            content.append(cast(BetaTextBlockParam, {"type": "text", "text": part.text}))
         elif isinstance(part, message.ToolCallPart):
+            tool_input: dict[str, object] = {}
+            if part.arguments_json:
+                try:
+                    parsed = json.loads(part.arguments_json)
+                except json.JSONDecodeError:
+                    parsed = {"_raw": part.arguments_json}
+                tool_input = cast(dict[str, object], parsed) if isinstance(parsed, dict) else {"_value": parsed}
+
             content.append(
-                {
-                    "type": "tool_use",
-                    "id": part.call_id,
-                    "name": part.tool_name,
-                    "input": json.loads(part.arguments_json) if part.arguments_json else None,
-                }
+                cast(
+                    BetaToolUseBlockParam,
+                    {
+                        "type": "tool_use",
+                        "id": part.call_id,
+                        "name": part.tool_name,
+                        "input": tool_input,
+                    },
+                )
             )
 
     _flush_thinking()
 
     if degraded_thinking_texts:
         degraded_text = "<thinking>\n" + "\n".join(degraded_thinking_texts) + "\n</thinking>"
-        content.insert(0, {"type": "text", "text": degraded_text})
+        content.insert(0, cast(BetaTextBlockParam, {"type": "text", "text": degraded_text}))
 
     return {"role": "assistant", "content": content}
 
@@ -174,7 +191,7 @@ def convert_history_to_input(
 ) -> list[BetaMessageParam]:
     """Convert a list of messages to beta message params."""
     messages: list[BetaMessageParam] = []
-    pending_tool_blocks: list[dict[str, object]] = []
+    pending_tool_blocks: list[BetaToolResultBlockParam] = []
 
     def flush_tool_blocks() -> None:
         nonlocal pending_tool_blocks
@@ -213,7 +230,12 @@ def convert_system_to_input(
             parts.append("\n".join(part.text for part in msg.parts))
     if not parts:
         return []
-    return [{"type": "text", "text": "\n".join(parts), "cache_control": {"type": "ephemeral"}}]
+    block: BetaTextBlockParam = {
+        "type": "text",
+        "text": "\n".join(parts),
+        "cache_control": {"type": "ephemeral"},
+    }
+    return [block]
 
 
 def convert_tool_schema(
@@ -222,11 +244,14 @@ def convert_tool_schema(
     if tools is None:
         return []
     return [
-        {
-            "input_schema": tool.parameters,
-            "type": "custom",
-            "name": tool.name,
-            "description": tool.description,
-        }
+        cast(
+            BetaToolParam,
+            {
+                "input_schema": tool.parameters,
+                "type": "custom",
+                "name": tool.name,
+                "description": tool.description,
+            },
+        )
         for tool in tools
     ]