klaude-code 1.2.11__py3-none-any.whl → 1.2.13__py3-none-any.whl

klaude_code/command/thinking_cmd.py

@@ -0,0 +1,229 @@
+import asyncio
+from typing import TYPE_CHECKING
+
+import questionary
+
+from klaude_code.command.command_abc import CommandABC, CommandResult
+from klaude_code.protocol import commands, events, llm_param, model
+
+if TYPE_CHECKING:
+    from klaude_code.core.agent import Agent
+
+
+# Thinking level options for different protocols
+RESPONSES_LEVELS = ["minimal", "low", "medium", "high"]
+RESPONSES_GPT51_LEVELS = ["none", "minimal", "low", "medium", "high"]
+RESPONSES_CODEX_MAX_LEVELS = ["medium", "high", "xhigh"]
+
+ANTHROPIC_LEVELS: list[tuple[str, int | None]] = [
+    ("off", 0),
+    ("low (2048 tokens)", 2048),
+    ("medium (8192 tokens)", 8192),
+    ("high (31999 tokens)", 31999),
+]
+
+
+def _is_openrouter_model_with_reasoning_effort(model_name: str | None) -> bool:
+    """Check if the model is GPT series, Grok or Gemini 3."""
+    if not model_name:
+        return False
+    model_lower = model_name.lower()
+    return model_lower.startswith(("openai/gpt-", "x-ai/grok-", "google/gemini-3"))
+
+
+def _is_gpt51_model(model_name: str | None) -> bool:
+    """Check if the model is GPT-5.1."""
+    if not model_name:
+        return False
+    return model_name.lower() in ["gpt5.1", "openai/gpt-5.1", "gpt-5.1-codex-2025-11-13"]
+
+
+def _is_codex_max_model(model_name: str | None) -> bool:
+    """Check if the model is GPT-5.1-codex-max."""
+    if not model_name:
+        return False
+    return "codex-max" in model_name.lower()
+
+
+def _get_levels_for_responses(model_name: str | None) -> list[str]:
+    """Get thinking levels for responses protocol."""
+    if _is_codex_max_model(model_name):
+        return RESPONSES_CODEX_MAX_LEVELS
+    if _is_gpt51_model(model_name):
+        return RESPONSES_GPT51_LEVELS
+    return RESPONSES_LEVELS
+
+
+def _format_current_thinking(config: llm_param.LLMConfigParameter) -> str:
+    """Format the current thinking configuration for display."""
+    thinking = config.thinking
+    if not thinking:
+        return "not configured"
+
+    protocol = config.protocol
+
+    if protocol in (llm_param.LLMClientProtocol.RESPONSES, llm_param.LLMClientProtocol.CODEX):
+        if thinking.reasoning_effort:
+            return f"reasoning_effort={thinking.reasoning_effort}"
+        return "not set"
+
+    if protocol == llm_param.LLMClientProtocol.ANTHROPIC:
+        if thinking.type == "disabled":
+            return "off"
+        if thinking.type == "enabled":
+            return f"enabled (budget_tokens={thinking.budget_tokens})"
+        return "not set"
+
+    if protocol == llm_param.LLMClientProtocol.OPENROUTER:
+        if _is_openrouter_model_with_reasoning_effort(config.model):
+            if thinking.reasoning_effort:
+                return f"reasoning_effort={thinking.reasoning_effort}"
+        else:
+            if thinking.type == "disabled":
+                return "off"
+            if thinking.type == "enabled":
+                return f"enabled (budget_tokens={thinking.budget_tokens})"
+        return "not set"
+
+    if protocol == llm_param.LLMClientProtocol.OPENAI:
+        if thinking.type == "disabled":
+            return "off"
+        if thinking.type == "enabled":
+            return f"enabled (budget_tokens={thinking.budget_tokens})"
+        return "not set"
+
+    return "unknown protocol"
+
+
+SELECT_STYLE = questionary.Style(
+    [
+        ("instruction", "ansibrightblack"),
+        ("pointer", "ansicyan"),
+        ("highlighted", "ansicyan"),
+        ("text", "ansibrightblack"),
+    ]
+)
+
+
+def _select_responses_thinking_sync(model_name: str | None) -> llm_param.Thinking | None:
+    """Select thinking level for responses/codex protocol (sync version)."""
+    levels = _get_levels_for_responses(model_name)
+    choices: list[questionary.Choice] = [questionary.Choice(title=level, value=level) for level in levels]
+
+    try:
+        result = questionary.select(
+            message="Select reasoning effort:",
+            choices=choices,
+            pointer="→",
+            instruction="Use arrow keys to move, Enter to select",
+            use_jk_keys=False,
+            style=SELECT_STYLE,
+        ).ask()
+
+        if result is None:
+            return None
+        return llm_param.Thinking(reasoning_effort=result)
+    except KeyboardInterrupt:
+        return None
+
+
+def _select_anthropic_thinking_sync() -> llm_param.Thinking | None:
+    """Select thinking level for anthropic/openai_compatible protocol (sync version)."""
+    choices: list[questionary.Choice] = [
+        questionary.Choice(title=label, value=tokens) for label, tokens in ANTHROPIC_LEVELS
+    ]
+
+    try:
+        result = questionary.select(
+            message="Select thinking level:",
+            choices=choices,
+            pointer="→",
+            instruction="Use arrow keys to move, Enter to select",
+            use_jk_keys=False,
+            style=SELECT_STYLE,
+        ).ask()
+        if result is None:
+            return llm_param.Thinking(type="disabled", budget_tokens=0)
+        return llm_param.Thinking(type="enabled", budget_tokens=result or 0)
+    except KeyboardInterrupt:
+        return None
+
+
+class ThinkingCommand(CommandABC):
+    """Configure model thinking/reasoning level."""
+
+    @property
+    def name(self) -> commands.CommandName:
+        return commands.CommandName.THINKING
+
+    @property
+    def summary(self) -> str:
+        return "Configure model thinking/reasoning level"
+
+    @property
+    def is_interactive(self) -> bool:
+        return True
+
+    async def run(self, raw: str, agent: "Agent") -> CommandResult:
+        if not agent.profile:
+            return self._no_change_result(agent, "No profile configured")
+
+        config = agent.profile.llm_client.get_llm_config()
+        protocol = config.protocol
+        model_name = config.model
+
+        current = _format_current_thinking(config)
+
+        # Select new thinking configuration based on protocol
+        new_thinking: llm_param.Thinking | None = None
+
+        if protocol in (llm_param.LLMClientProtocol.RESPONSES, llm_param.LLMClientProtocol.CODEX):
+            new_thinking = await asyncio.to_thread(_select_responses_thinking_sync, model_name)
+
+        elif protocol == llm_param.LLMClientProtocol.ANTHROPIC:
+            new_thinking = await asyncio.to_thread(_select_anthropic_thinking_sync)
+
+        elif protocol == llm_param.LLMClientProtocol.OPENROUTER:
+            if _is_openrouter_model_with_reasoning_effort(model_name):
+                new_thinking = await asyncio.to_thread(_select_responses_thinking_sync, model_name)
+            else:
+                new_thinking = await asyncio.to_thread(_select_anthropic_thinking_sync)
+
+        elif protocol == llm_param.LLMClientProtocol.OPENAI:
+            # openai_compatible uses anthropic style
+            new_thinking = await asyncio.to_thread(_select_anthropic_thinking_sync)
+
+        else:
+            return self._no_change_result(agent, f"Unsupported protocol: {protocol}")
+
+        if new_thinking is None:
+            return self._no_change_result(agent, "(no change)")
+
+        # Apply the new thinking configuration
+        config.thinking = new_thinking
+        new_status = _format_current_thinking(config)
+
+        return CommandResult(
+            events=[
+                events.DeveloperMessageEvent(
+                    session_id=agent.session.id,
+                    item=model.DeveloperMessageItem(
+                        content=f"Thinking changed: {current} -> {new_status}",
+                        command_output=model.CommandOutput(command_name=self.name),
+                    ),
+                )
+            ]
+        )
+
+    def _no_change_result(self, agent: "Agent", message: str) -> CommandResult:
+        return CommandResult(
+            events=[
+                events.DeveloperMessageEvent(
+                    session_id=agent.session.id,
+                    item=model.DeveloperMessageItem(
+                        content=message,
+                        command_output=model.CommandOutput(command_name=self.name),
+                    ),
+                )
+            ]
+        )

klaude_code/config/__init__.py

@@ -4,8 +4,8 @@ from .select_model import select_model_from_config
 
 __all__ = [
     "Config",
-    "load_config",
     "config_path",
     "display_models_and_providers",
+    "load_config",
     "select_model_from_config",
 ]

klaude_code/config/list_model.py

@@ -34,7 +34,7 @@ def _display_codex_status(console: Console) -> None:
             )
         )
     else:
-        expires_dt = datetime.datetime.fromtimestamp(state.expires_at, tz=datetime.timezone.utc)
+        expires_dt = datetime.datetime.fromtimestamp(state.expires_at, tz=datetime.UTC)
         console.print(
             Text.assemble(
                 ("Codex Status: ", "bold"),

klaude_code/config/select_model.py

@@ -1,6 +1,5 @@
 from klaude_code.config.config import load_config
 from klaude_code.trace import log
-from klaude_code.ui.rich.searchable_text import SearchableFormattedList
 
 
 def select_model_from_config(preferred: str | None = None) -> str | None:
@@ -16,9 +15,6 @@ def select_model_from_config(preferred: str | None = None) -> str | None:
         raise ValueError("No models configured. Please update your config.yaml")
 
     names: list[str] = [m.model_name for m in models]
-    default_name: str | None = (
-        preferred if preferred in names else (config.main_model if config.main_model in names else None)
-    )
 
     try:
         import questionary
@@ -28,29 +24,23 @@ def select_model_from_config(preferred: str | None = None) -> str | None:
         max_model_name_length = max(len(m.model_name) for m in models)
         for m in models:
             star = "★ " if m.model_name == config.main_model else "  "
-            fragments = [
-                ("class:t", f"{star}{m.model_name:<{max_model_name_length}}   → "),
-                ("class:b", m.model_params.model or "N/A"),
-                ("class:d", f" {m.provider}"),
-            ]
-            # Provide a formatted title for display and a plain text for search.
-            title = SearchableFormattedList(fragments)
+            title = f"{star}{m.model_name:<{max_model_name_length}}   →  {m.model_params.model or 'N/A'} @ {m.provider}"
             choices.append(questionary.Choice(title=title, value=m.model_name))
 
         try:
             result = questionary.select(
                 message="Select a model:",
                 choices=choices,
-                default=default_name,
                 pointer="→",
                 instruction="↑↓ to move • Enter to select",
                 use_jk_keys=False,
                 use_search_filter=True,
                 style=questionary.Style(
                     [
-                        ("t", ""),
-                        ("b", "bold"),
-                        ("d", "dim"),
+                        ("instruction", "ansibrightblack"),
+                        ("pointer", "ansicyan"),
+                        ("highlighted", "ansicyan"),
+                        ("text", "ansibrightblack"),
                         # search filter colors at the bottom
                         ("search_success", "noinherit fg:ansigreen"),
                         ("search_none", "noinherit fg:ansired"),

klaude_code/const/__init__.py

@@ -62,7 +62,7 @@ BASH_DEFAULT_TIMEOUT_MS = 120000
 
 # -- Tool Output --
 # Maximum length for tool output before truncation
-TOOL_OUTPUT_MAX_LENGTH = 50000
+TOOL_OUTPUT_MAX_LENGTH = 40000
 
 # Characters to show from the beginning of truncated output
 TOOL_OUTPUT_DISPLAY_HEAD = 10000

klaude_code/core/agent.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-from collections.abc import AsyncGenerator, Callable, Iterable
+from collections.abc import AsyncGenerator, Iterable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Protocol
+from typing import Protocol
 
-from klaude_code.core.prompt import get_system_prompt as load_system_prompt
+from klaude_code.core.prompt import load_system_prompt
 from klaude_code.core.reminders import Reminder, load_agent_reminders
 from klaude_code.core.task import SessionContext, TaskExecutionContext, TaskExecutor
 from klaude_code.core.tool import build_todo_context, get_registry, load_agent_tools
@@ -14,38 +14,21 @@ from klaude_code.protocol.model import UserInputPayload
 from klaude_code.session import Session
 from klaude_code.trace import DebugType, log_debug
 
-if TYPE_CHECKING:
-    from klaude_code.core.manager.llm_clients import LLMClients
-
 
 @dataclass(frozen=True)
 class AgentProfile:
     """Encapsulates the active LLM client plus prompts/tools/reminders."""
 
-    llm_client_factory: Callable[[], LLMClientABC]
+    llm_client: LLMClientABC
     system_prompt: str | None
     tools: list[llm_param.ToolSchema]
     reminders: list[Reminder]
 
-    _llm_client: LLMClientABC | None = None
-
-    @property
-    def llm_client(self) -> LLMClientABC:
-        if self._llm_client is None:
-            object.__setattr__(self, "_llm_client", self.llm_client_factory())
-        return self._llm_client  # type: ignore[return-value]
-
 
 class ModelProfileProvider(Protocol):
     """Strategy interface for constructing agent profiles."""
 
     def build_profile(
-        self,
-        llm_clients: LLMClients,
-        sub_agent_type: tools.SubAgentType | None = None,
-    ) -> AgentProfile: ...
-
-    def build_profile_eager(
         self,
         llm_client: LLMClientABC,
         sub_agent_type: tools.SubAgentType | None = None,
@@ -56,26 +39,13 @@ class DefaultModelProfileProvider(ModelProfileProvider):
     """Default provider backed by global prompts/tool/reminder registries."""
 
     def build_profile(
-        self,
-        llm_clients: LLMClients,
-        sub_agent_type: tools.SubAgentType | None = None,
-    ) -> AgentProfile:
-        model_name = llm_clients.main_model_name
-        return AgentProfile(
-            llm_client_factory=lambda: llm_clients.main,
-            system_prompt=load_system_prompt(model_name, sub_agent_type),
-            tools=load_agent_tools(model_name, sub_agent_type),
-            reminders=load_agent_reminders(model_name, sub_agent_type),
-        )
-
-    def build_profile_eager(
         self,
         llm_client: LLMClientABC,
         sub_agent_type: tools.SubAgentType | None = None,
     ) -> AgentProfile:
         model_name = llm_client.model_name
         return AgentProfile(
-            llm_client_factory=lambda: llm_client,
+            llm_client=llm_client,
             system_prompt=load_system_prompt(model_name, sub_agent_type),
             tools=load_agent_tools(model_name, sub_agent_type),
             reminders=load_agent_reminders(model_name, sub_agent_type),
@@ -86,26 +56,13 @@ class VanillaModelProfileProvider(ModelProfileProvider):
     """Provider that strips prompts, reminders, and tools for vanilla mode."""
 
     def build_profile(
-        self,
-        llm_clients: LLMClients,
-        sub_agent_type: tools.SubAgentType | None = None,
-    ) -> AgentProfile:
-        model_name = llm_clients.main_model_name
-        return AgentProfile(
-            llm_client_factory=lambda: llm_clients.main,
-            system_prompt=None,
-            tools=load_agent_tools(model_name, vanilla=True),
-            reminders=load_agent_reminders(model_name, vanilla=True),
-        )
-
-    def build_profile_eager(
         self,
         llm_client: LLMClientABC,
         sub_agent_type: tools.SubAgentType | None = None,
     ) -> AgentProfile:
         model_name = llm_client.model_name
         return AgentProfile(
-            llm_client_factory=lambda: llm_client,
+            llm_client=llm_client,
             system_prompt=None,
             tools=load_agent_tools(model_name, vanilla=True),
             reminders=load_agent_reminders(model_name, vanilla=True),
@@ -117,14 +74,12 @@ class Agent:
         self,
         session: Session,
         profile: AgentProfile,
-        model_name: str | None = None,
     ):
         self.session: Session = session
         self.profile: AgentProfile = profile
         self._current_task: TaskExecutor | None = None
-        self._prev_context_token: int = 0  # Track context size from previous task for delta calculation
-        if not self.session.model_name and model_name:
-            self.session.model_name = model_name
+        if not self.session.model_name:
+            self.session.model_name = profile.llm_client.model_name
 
     def cancel(self) -> Iterable[events.Event]:
         """Handle agent cancellation and persist an interrupt marker and tool cancellations.
@@ -137,8 +92,7 @@ class Agent:
         """
         # First, cancel any running task so it stops emitting events.
         if self._current_task is not None:
-            for ui_event in self._current_task.cancel():
-                yield ui_event
+            yield from self._current_task.cancel()
             self._current_task = None
 
         # Record an interrupt marker in the session history
@@ -149,7 +103,7 @@ class Agent:
             debug_type=DebugType.EXECUTION,
         )
 
-    async def run_task(self, user_input: UserInputPayload) -> AsyncGenerator[events.Event, None]:
+    async def run_task(self, user_input: UserInputPayload) -> AsyncGenerator[events.Event]:
         session_ctx = SessionContext(
             session_id=self.session.id,
             get_conversation_history=lambda: self.session.conversation_history,
@@ -170,17 +124,11 @@ class Agent:
 
         try:
             async for event in task.run(user_input):
-                # Compute context_delta for TaskMetadataEvent
-                if isinstance(event, events.TaskMetadataEvent):
-                    usage = event.metadata.main.usage
-                    if usage is not None and usage.context_token is not None:
-                        usage.context_delta = usage.context_token - self._prev_context_token
-                        self._prev_context_token = usage.context_token
                 yield event
         finally:
             self._current_task = None
 
-    async def replay_history(self) -> AsyncGenerator[events.Event, None]:
+    async def replay_history(self) -> AsyncGenerator[events.Event]:
         """Yield UI events reconstructed from saved conversation history."""
 
         if len(self.session.conversation_history) == 0:
@@ -192,21 +140,18 @@ class Agent:
             session_id=self.session.id,
         )
 
-    async def _process_reminder(self, reminder: Reminder) -> AsyncGenerator[events.DeveloperMessageEvent, None]:
+    async def _process_reminder(self, reminder: Reminder) -> AsyncGenerator[events.DeveloperMessageEvent]:
         """Process a single reminder and yield events if it produces output."""
         item = await reminder(self.session)
         if item is not None:
             self.session.append_history([item])
             yield events.DeveloperMessageEvent(session_id=self.session.id, item=item)
 
-    def set_model_profile(self, profile: AgentProfile, model_name: str | None = None) -> None:
+    def set_model_profile(self, profile: AgentProfile) -> None:
         """Apply a fully constructed profile to the agent."""
 
         self.profile = profile
-        if model_name:
-            self.session.model_name = model_name
-        elif not self.session.model_name:
-            self.session.model_name = profile.llm_client.model_name
+        self.session.model_name = profile.llm_client.model_name
 
     def get_llm_client(self) -> LLMClientABC:
         return self.profile.llm_client

klaude_code/core/executor.py

@@ -264,14 +264,14 @@ class ExecutorContext:
             import traceback
 
             log_debug(
-                f"Agent task {task_id} failed: {str(e)}",
+                f"Agent task {task_id} failed: {e!s}",
                 style="red",
                 debug_type=DebugType.EXECUTION,
             )
             log_debug(traceback.format_exc(), style="red", debug_type=DebugType.EXECUTION)
             await self.emit_event(
                 events.ErrorEvent(
-                    error_message=f"Agent task failed: [{e.__class__.__name__}] {str(e)}",
+                    error_message=f"Agent task failed: [{e.__class__.__name__}] {e!s}",
                     can_retry=False,
                 )
             )
@@ -317,6 +317,7 @@ class Executor:
         self.submission_queue: asyncio.Queue[op.Submission] = asyncio.Queue()
         # Track completion events for all submissions (not just those with ActiveTask)
         self._completion_events: dict[str, asyncio.Event] = {}
+        self._background_tasks: set[asyncio.Task[None]] = set()
 
     async def submit(self, operation: op.Operation) -> str:
         """
@@ -388,12 +389,12 @@ class Executor:
             except Exception as e:
                 # Handle unexpected errors
                 log_debug(
-                    f"Executor error: {str(e)}",
+                    f"Executor error: {e!s}",
                     style="red",
                     debug_type=DebugType.EXECUTION,
                 )
                 await self.context.emit_event(
-                    events.ErrorEvent(error_message=f"Executor error: {str(e)}", can_retry=False)
+                    events.ErrorEvent(error_message=f"Executor error: {e!s}", can_retry=False)
                 )
 
     async def stop(self) -> None:
@@ -420,7 +421,7 @@ class Executor:
             await self.submission_queue.put(submission)
         except Exception as e:
             log_debug(
-                f"Failed to send EndOperation: {str(e)}",
+                f"Failed to send EndOperation: {e!s}",
                 style="red",
                 debug_type=DebugType.EXECUTION,
             )
@@ -460,17 +461,17 @@ class Executor:
                     event.set()
             else:
                 # Run in background so the submission loop can continue (e.g., to handle interrupts)
-                asyncio.create_task(_await_agent_and_complete(task))
+                background_task = asyncio.create_task(_await_agent_and_complete(task))
+                self._background_tasks.add(background_task)
+                background_task.add_done_callback(self._background_tasks.discard)
 
         except Exception as e:
             log_debug(
-                f"Failed to handle submission {submission.id}: {str(e)}",
+                f"Failed to handle submission {submission.id}: {e!s}",
                 style="red",
                 debug_type=DebugType.EXECUTION,
             )
-            await self.context.emit_event(
-                events.ErrorEvent(error_message=f"Operation failed: {str(e)}", can_retry=False)
-            )
+            await self.context.emit_event(events.ErrorEvent(error_message=f"Operation failed: {e!s}", can_retry=False))
             # Set completion event even on error to prevent wait_for_completion from hanging
             event = self._completion_events.get(submission.id)
             if event is not None:

klaude_code/core/manager/agent_manager.py

@@ -51,8 +51,8 @@ class AgentManager:
             if agent is not None:
                 return agent
             session = Session.load(session_id)
-        profile = self._model_profile_provider.build_profile(self._llm_clients)
-        agent = Agent(session=session, profile=profile, model_name=self._llm_clients.main_model_name)
+        profile = self._model_profile_provider.build_profile(self._llm_clients.main)
+        agent = Agent(session=session, profile=profile)
 
         async for evt in agent.replay_history():
             await self.emit_event(evt)
@@ -60,7 +60,7 @@ class AgentManager:
         await self.emit_event(
             events.WelcomeEvent(
                 work_dir=str(session.work_dir),
-                llm_config=self._llm_clients.get_llm_config(),
+                llm_config=self._llm_clients.main.get_llm_config(),
             )
         )
 
@@ -81,7 +81,7 @@ class AgentManager:
 
         llm_config = config.get_model_config(model_name)
         llm_client = create_llm_client(llm_config)
-        agent.set_model_profile(self._model_profile_provider.build_profile_eager(llm_client), model_name=model_name)
+        agent.set_model_profile(self._model_profile_provider.build_profile(llm_client))
 
         developer_item = model.DeveloperMessageItem(
             content=f"switched to model: {model_name}",

klaude_code/core/manager/llm_clients.py

@@ -2,66 +2,27 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
+from dataclasses import dataclass
+from dataclasses import field as dataclass_field
 
 from klaude_code.llm.client import LLMClientABC
-from klaude_code.protocol import llm_param
 from klaude_code.protocol.tools import SubAgentType
 
 
-class LLMClients:
-    """Container for LLM clients used by main agent and sub-agents."""
-
-    def __init__(
-        self,
-        main_factory: Callable[[], LLMClientABC],
-        main_model_name: str,
-        main_llm_config: llm_param.LLMConfigParameter,
-    ) -> None:
-        self._main_factory: Callable[[], LLMClientABC] | None = main_factory
-        self._main_client: LLMClientABC | None = None
-        self._main_model_name: str = main_model_name
-        self._main_llm_config: llm_param.LLMConfigParameter = main_llm_config
-        self._sub_clients: dict[SubAgentType, LLMClientABC] = {}
-        self._sub_factories: dict[SubAgentType, Callable[[], LLMClientABC]] = {}
+def _default_sub_clients() -> dict[SubAgentType, LLMClientABC]:
+    return {}
 
-    @property
-    def main_model_name(self) -> str:
-        return self._main_model_name
 
-    def get_llm_config(self) -> llm_param.LLMConfigParameter:
-        return self._main_llm_config
-
-    @property
-    def main(self) -> LLMClientABC:
-        if self._main_client is None:
-            if self._main_factory is None:
-                raise RuntimeError("Main client factory not set")
-            self._main_client = self._main_factory()
-            self._main_factory = None
-        return self._main_client
+@dataclass
+class LLMClients:
+    """Container for LLM clients used by main agent and sub-agents."""
 
-    def register_sub_client_factory(
-        self,
-        sub_agent_type: SubAgentType,
-        factory: Callable[[], LLMClientABC],
-    ) -> None:
-        self._sub_factories[sub_agent_type] = factory
+    main: LLMClientABC
+    sub_clients: dict[SubAgentType, LLMClientABC] = dataclass_field(default_factory=_default_sub_clients)
 
     def get_client(self, sub_agent_type: SubAgentType | None = None) -> LLMClientABC:
         """Return client for a sub-agent type or the main client."""
 
         if sub_agent_type is None:
             return self.main
-
-        existing = self._sub_clients.get(sub_agent_type)
-        if existing is not None:
-            return existing
-
-        factory = self._sub_factories.get(sub_agent_type)
-        if factory is None:
-            return self.main
-
-        client = factory()
-        self._sub_clients[sub_agent_type] = client
-        return client
+        return self.sub_clients.get(sub_agent_type) or self.main