klaude-code 1.2.8__py3-none-any.whl → 1.2.10__py3-none-any.whl

klaude_code/command/status_cmd.py

@@ -1,51 +1,73 @@
+from typing import TYPE_CHECKING
+
 from klaude_code.command.command_abc import CommandABC, CommandResult
 from klaude_code.command.registry import register_command
-from klaude_code.core.agent import Agent
 from klaude_code.protocol import commands, events, model
 from klaude_code.session.session import Session
 
+if TYPE_CHECKING:
+    from klaude_code.core.agent import Agent
+
+
+class AggregatedUsage(model.BaseModel):
+    """Aggregated usage statistics including per-model breakdown."""
+
+    total: model.Usage
+    by_model: list[model.TaskMetadata]
+    task_count: int
+
 
-def accumulate_session_usage(session: Session) -> tuple[model.Usage, int]:
-    """Accumulate usage statistics from all ResponseMetadataItems in session history.
+def accumulate_session_usage(session: Session) -> AggregatedUsage:
+    """Accumulate usage statistics from all TaskMetadataItems in session history.
 
-    Returns:
-        A tuple of (accumulated_usage, task_count)
+    Includes both main agent and sub-agent task metadata, grouped by model+provider.
     """
-    total = model.Usage()
+    all_metadata: list[model.TaskMetadata] = []
     task_count = 0
-    first_currency_set = False
 
     for item in session.conversation_history:
-        if isinstance(item, model.ResponseMetadataItem) and item.usage:
+        if isinstance(item, model.TaskMetadataItem):
             task_count += 1
-            usage = item.usage
-            
-            # Set currency from first usage item
-            if not first_currency_set and usage.currency:
-                total.currency = usage.currency
-                first_currency_set = True
-            
-            total.input_tokens += usage.input_tokens
-            total.cached_tokens += usage.cached_tokens
-            total.reasoning_tokens += usage.reasoning_tokens
-            total.output_tokens += usage.output_tokens
-            total.total_tokens += usage.total_tokens
-
-            # Accumulate costs
-            if usage.input_cost is not None:
-                total.input_cost = (total.input_cost or 0.0) + usage.input_cost
-            if usage.output_cost is not None:
-                total.output_cost = (total.output_cost or 0.0) + usage.output_cost
-            if usage.cache_read_cost is not None:
-                total.cache_read_cost = (total.cache_read_cost or 0.0) + usage.cache_read_cost
-            if usage.total_cost is not None:
-                total.total_cost = (total.total_cost or 0.0) + usage.total_cost
-
-            # Keep the latest context_usage_percent
-            if usage.context_usage_percent is not None:
-                total.context_usage_percent = usage.context_usage_percent
-
-    return total, task_count
+            all_metadata.append(item.main)
+            all_metadata.extend(item.sub_agent_task_metadata)
+
+    # Aggregate by model+provider
+    by_model = model.TaskMetadata.aggregate_by_model(all_metadata)
+
+    # Calculate total from aggregated results
+    total = model.Usage()
+    for meta in by_model:
+        if not meta.usage:
+            continue
+        usage = meta.usage
+
+        # Set currency from first
+        if total.currency == "USD" and usage.currency:
+            total.currency = usage.currency
+
+        # Accumulate primary token fields (total_tokens is computed)
+        total.input_tokens += usage.input_tokens
+        total.cached_tokens += usage.cached_tokens
+        total.reasoning_tokens += usage.reasoning_tokens
+        total.output_tokens += usage.output_tokens
+
+        # Accumulate cost components (total_cost is computed)
+        if usage.input_cost is not None:
+            total.input_cost = (total.input_cost or 0.0) + usage.input_cost
+        if usage.output_cost is not None:
+            total.output_cost = (total.output_cost or 0.0) + usage.output_cost
+        if usage.cache_read_cost is not None:
+            total.cache_read_cost = (total.cache_read_cost or 0.0) + usage.cache_read_cost
+
+        # Track peak context window size (max across all tasks)
+        if usage.context_token is not None:
+            total.context_token = usage.context_token
+
+        # Keep the latest context_limit for computed context_usage_percent
+        if usage.context_limit is not None:
+            total.context_limit = usage.context_limit
+
+    return AggregatedUsage(total=total, by_model=by_model, task_count=task_count)
 
 
 def _format_tokens(tokens: int) -> str:
@@ -67,20 +89,42 @@ def _format_cost(cost: float | None, currency: str = "USD") -> str:
     return f"{symbol}{cost:.2f}"
 
 
-def format_status_content(usage: model.Usage) -> str:
-    """Format session status as comma-separated text."""
-    parts: list[str] = []
+def _format_model_usage_line(meta: model.TaskMetadata) -> str:
+    """Format a single model's usage as a line."""
+    model_label = meta.model_name
+    if meta.provider:
+        model_label = f"{meta.model_name} ({meta.provider})"
+
+    usage = meta.usage
+    if not usage:
+        return f"      {model_label}: no usage data"
+
+    cost_str = _format_cost(usage.total_cost, usage.currency)
+    return (
+        f"      {model_label}: "
+        f"{_format_tokens(usage.input_tokens)} input, "
+        f"{_format_tokens(usage.output_tokens)} output, "
+        f"{_format_tokens(usage.cached_tokens)} cache read, "
+        f"{_format_tokens(usage.reasoning_tokens)} thinking, "
+        f"({cost_str})"
+    )
+
+
+def format_status_content(aggregated: AggregatedUsage) -> str:
+    """Format session status with per-model breakdown."""
+    lines: list[str] = []
 
-    parts.append(f"Input: {_format_tokens(usage.input_tokens)}")
-    if usage.cached_tokens > 0:
-        parts.append(f"Cached: {_format_tokens(usage.cached_tokens)}")
-    parts.append(f"Output: {_format_tokens(usage.output_tokens)}")
-    parts.append(f"Total: {_format_tokens(usage.total_tokens)}")
+    # Total cost line
+    total_cost_str = _format_cost(aggregated.total.total_cost, aggregated.total.currency)
+    lines.append(f"Total cost: {total_cost_str}")
 
-    if usage.total_cost is not None:
-        parts.append(f"Cost: {_format_cost(usage.total_cost, usage.currency)}")
+    # Per-model breakdown
+    if aggregated.by_model:
+        lines.append("Usage by model:")
+        for stats in aggregated.by_model:
+            lines.append(_format_model_usage_line(stats))
 
-    return ", ".join(parts)
+    return "\n".join(lines)
 
 
 @register_command
@@ -95,22 +139,20 @@ class StatusCommand(CommandABC):
     def summary(self) -> str:
         return "Show session usage statistics"
 
-    async def run(self, raw: str, agent: Agent) -> CommandResult:
+    async def run(self, raw: str, agent: "Agent") -> CommandResult:
         session = agent.session
-        usage, task_count = accumulate_session_usage(session)
+        aggregated = accumulate_session_usage(session)
 
         event = events.DeveloperMessageEvent(
             session_id=session.id,
             item=model.DeveloperMessageItem(
-                content=format_status_content(usage),
+                content=format_status_content(aggregated),
                 command_output=model.CommandOutput(
                     command_name=self.name,
-                    ui_extra=model.ToolResultUIExtra(
-                        type=model.ToolResultUIExtraType.SESSION_STATUS,
-                        session_status=model.SessionStatusUIExtra(
-                            usage=usage,
-                            task_count=task_count,
-                        ),
+                    ui_extra=model.SessionStatusUIExtra(
+                        usage=aggregated.total,
+                        task_count=aggregated.task_count,
+                        by_model=aggregated.by_model,
                     ),
                 ),
             ),

klaude_code/command/terminal_setup_cmd.py

@@ -1,12 +1,15 @@
 import os
 import subprocess
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 from klaude_code.command.command_abc import CommandABC, CommandResult
 from klaude_code.command.registry import register_command
-from klaude_code.core.agent import Agent
 from klaude_code.protocol import commands, events, model
 
+if TYPE_CHECKING:
+    from klaude_code.core.agent import Agent
+
 
 @register_command
 class TerminalSetupCommand(CommandABC):
@@ -24,7 +27,7 @@ class TerminalSetupCommand(CommandABC):
     def is_interactive(self) -> bool:
         return False
 
-    async def run(self, raw: str, agent: Agent) -> CommandResult:
+    async def run(self, raw: str, agent: "Agent") -> CommandResult:
         term_program = os.environ.get("TERM_PROGRAM", "").lower()
 
         try:
@@ -223,7 +226,7 @@ class TerminalSetupCommand(CommandABC):
 
         return message
 
-    def _create_success_result(self, agent: Agent, message: str) -> CommandResult:
+    def _create_success_result(self, agent: "Agent", message: str) -> CommandResult:
         """Create success result"""
         return CommandResult(
             events=[
@@ -237,7 +240,7 @@ class TerminalSetupCommand(CommandABC):
             ]
         )
 
-    def _create_error_result(self, agent: Agent, message: str) -> CommandResult:
+    def _create_error_result(self, agent: "Agent", message: str) -> CommandResult:
         """Create error result"""
         return CommandResult(
             events=[

klaude_code/const/__init__.py

@@ -91,7 +91,7 @@ INVALID_TOOL_CALL_MAX_LENGTH = 500
 TRUNCATE_DISPLAY_MAX_LINE_LENGTH = 1000
 
 # Maximum lines for truncated display output
-TRUNCATE_DISPLAY_MAX_LINES = 10
+TRUNCATE_DISPLAY_MAX_LINES = 20
 
 # Maximum lines for sub-agent result display
 SUB_AGENT_RESULT_MAX_LINES = 12

klaude_code/core/agent.py

@@ -1,34 +1,51 @@
 from __future__ import annotations
 
-from collections.abc import AsyncGenerator, Iterable
+from collections.abc import AsyncGenerator, Callable, Iterable
 from dataclasses import dataclass
-from typing import Protocol
+from typing import TYPE_CHECKING, Protocol
 
 from klaude_code.core.prompt import get_system_prompt as load_system_prompt
 from klaude_code.core.reminders import Reminder, load_agent_reminders
-from klaude_code.core.task import TaskExecutionContext, TaskExecutor
-from klaude_code.core.tool import TodoContext, get_registry, load_agent_tools
+from klaude_code.core.task import SessionContext, TaskExecutionContext, TaskExecutor
+from klaude_code.core.tool import build_todo_context, get_registry, load_agent_tools
 from klaude_code.llm import LLMClientABC
 from klaude_code.protocol import events, llm_param, model, tools
 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: LLMClientABC
+    llm_client_factory: Callable[[], 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,
@@ -39,13 +56,26 @@ 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=llm_client,
+            llm_client_factory=lambda: 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),
@@ -56,13 +86,26 @@ 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=llm_client,
+            llm_client_factory=lambda: llm_client,
             system_prompt=None,
             tools=load_agent_tools(model_name, vanilla=True),
             reminders=load_agent_reminders(model_name, vanilla=True),
@@ -74,13 +117,13 @@ class Agent:
         self,
         session: Session,
         profile: AgentProfile,
+        model_name: str | None = None,
     ):
         self.session: Session = session
-        self.profile: AgentProfile | None = None
-        # Active task executor, if any
+        self.profile: AgentProfile = profile
         self._current_task: TaskExecutor | None = None
-        # Ensure runtime configuration matches the active model on initialization
-        self.set_model_profile(profile)
+        if not self.session.model_name and model_name:
+            self.session.model_name = model_name
 
     def cancel(self) -> Iterable[events.Event]:
         """Handle agent cancellation and persist an interrupt marker and tool cancellations.
@@ -106,17 +149,17 @@ class Agent:
         )
 
     async def run_task(self, user_input: UserInputPayload) -> AsyncGenerator[events.Event, None]:
-        context = TaskExecutionContext(
+        session_ctx = SessionContext(
             session_id=self.session.id,
-            profile=self._require_profile(),
             get_conversation_history=lambda: self.session.conversation_history,
             append_history=self.session.append_history,
-            tool_registry=get_registry(),
             file_tracker=self.session.file_tracker,
-            todo_context=TodoContext(
-                get_todos=lambda: self.session.todos,
-                set_todos=lambda todos: setattr(self.session, "todos", todos),
-            ),
+            todo_context=build_todo_context(self.session),
+        )
+        context = TaskExecutionContext(
+            session_ctx=session_ctx,
+            profile=self.profile,
+            tool_registry=get_registry(),
             process_reminder=self._process_reminder,
             sub_agent_state=self.session.sub_agent_state,
         )
@@ -149,17 +192,14 @@ class Agent:
             self.session.append_history([item])
             yield events.DeveloperMessageEvent(session_id=self.session.id, item=item)
 
-    def set_model_profile(self, profile: AgentProfile) -> None:
+    def set_model_profile(self, profile: AgentProfile, model_name: str | None = None) -> None:
         """Apply a fully constructed profile to the agent."""
 
         self.profile = profile
-        if not self.session.model_name:
+        if model_name:
+            self.session.model_name = model_name
+        elif not self.session.model_name:
             self.session.model_name = profile.llm_client.model_name
 
     def get_llm_client(self) -> LLMClientABC:
-        return self._require_profile().llm_client
-
-    def _require_profile(self) -> AgentProfile:
-        if self.profile is None:
-            raise RuntimeError("Agent profile is not initialized")
-        return self.profile
+        return self.profile.llm_client

klaude_code/core/executor.py

@@ -117,7 +117,7 @@ class ExecutorContext:
         if operation.session_id is None:
             raise ValueError("session_id cannot be None")
 
-        await self.agent_manager.ensure_agent(operation.session_id)
+        await self.agent_manager.ensure_agent(operation.session_id, is_new_session=operation.is_new_session)
 
     async def handle_user_input(self, operation: op.UserInputOperation) -> None:
         """Handle a user input operation by running it through an agent."""
@@ -482,4 +482,4 @@ class Executor:
 
 # Static type check: ExecutorContext must satisfy OperationHandler protocol.
 # If this line causes a type error, ExecutorContext is missing required methods.
-_: type[OperationHandler] = ExecutorContext  # pyright: ignore[reportUnusedVariable]
+_: type[OperationHandler] = ExecutorContext

klaude_code/core/manager/agent_manager.py

@@ -38,16 +38,15 @@ class AgentManager:
 
         await self._event_queue.put(event)
 
-    async def ensure_agent(self, session_id: str) -> Agent:
+    async def ensure_agent(self, session_id: str, *, is_new_session: bool = False) -> Agent:
         """Return an existing agent for the session or create a new one."""
-
         agent = self._active_agents.get(session_id)
         if agent is not None:
             return agent
 
-        session = Session.load(session_id)
-        profile = self._model_profile_provider.build_profile(self._llm_clients.main)
-        agent = Agent(session=session, profile=profile)
+        session = Session.load(session_id, skip_if_missing=is_new_session)
+        profile = self._model_profile_provider.build_profile(self._llm_clients)
+        agent = Agent(session=session, profile=profile, model_name=self._llm_clients.main_model_name)
 
         async for evt in agent.replay_history():
             await self.emit_event(evt)
@@ -55,7 +54,7 @@ class AgentManager:
         await self.emit_event(
             events.WelcomeEvent(
                 work_dir=str(session.work_dir),
-                llm_config=self._llm_clients.main.get_llm_config(),
+                llm_config=self._llm_clients.get_llm_config(),
             )
         )
 
@@ -76,7 +75,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(llm_client))
+        agent.set_model_profile(self._model_profile_provider.build_profile_eager(llm_client), model_name=model_name)
 
         developer_item = model.DeveloperMessageItem(
             content=f"switched to model: {model_name}",

klaude_code/core/manager/llm_clients.py

@@ -2,41 +2,66 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
-from dataclasses import field as dataclass_field
+from collections.abc import Callable
 
 from klaude_code.llm.client import LLMClientABC
+from klaude_code.protocol import llm_param
 from klaude_code.protocol.tools import SubAgentType
 
 
-def _default_sub_clients() -> dict[SubAgentType, LLMClientABC]:
-    """Return an empty mapping for sub-agent clients.
+class LLMClients:
+    """Container for LLM clients used by main agent and sub-agents."""
 
-    Defined separately so static type checkers can infer the dictionary
-    key and value types instead of treating them as ``Unknown``.
-    """
+    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]] = {}
 
-    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
 
-@dataclass
-class LLMClients:
-    """Container for LLM clients used by main agent and sub-agents."""
+    @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
 
-    main: LLMClientABC
-    sub_clients: dict[SubAgentType, LLMClientABC] = dataclass_field(default_factory=_default_sub_clients)
+    def register_sub_client_factory(
+        self,
+        sub_agent_type: SubAgentType,
+        factory: Callable[[], LLMClientABC],
+    ) -> None:
+        self._sub_factories[sub_agent_type] = factory
 
     def get_client(self, sub_agent_type: SubAgentType | None = None) -> LLMClientABC:
-        """Return client for a sub-agent type or the main client.
+        """Return client for a sub-agent type or the main client."""
 
-        Args:
-            sub_agent_type: Optional sub-agent type whose client should be returned.
+        if sub_agent_type is None:
+            return self.main
 
-        Returns:
-            The LLM client corresponding to the sub-agent type, or the main client
-            when no specialized client is available.
-        """
+        existing = self._sub_clients.get(sub_agent_type)
+        if existing is not None:
+            return existing
 
-        if sub_agent_type is None:
+        factory = self._sub_factories.get(sub_agent_type)
+        if factory is None:
             return self.main
-        return self.sub_clients.get(sub_agent_type) or self.main
+
+        client = factory()
+        self._sub_clients[sub_agent_type] = client
+        return client

klaude_code/core/manager/llm_clients_builder.py

@@ -32,18 +32,30 @@ def build_llm_clients(
         debug_type=DebugType.LLM_CONFIG,
     )
 
-    main_client = create_llm_client(llm_config)
-    sub_clients: dict[SubAgentType, LLMClientABC] = {}
+    main_model_name = str(llm_config.model)
+
+    def _main_factory() -> LLMClientABC:
+        return create_llm_client(llm_config)
+
+    clients = LLMClients(
+        main_factory=_main_factory,
+        main_model_name=main_model_name,
+        main_llm_config=llm_config,
+    )
 
-    # Initialize sub-agent clients
     for sub_agent_type in enabled_sub_agents or []:
         model_name = config.subagent_models.get(sub_agent_type)
         if not model_name:
             continue
+
         profile = get_sub_agent_profile(sub_agent_type)
-        if not profile.enabled_for_model(main_client.model_name):
+        if not profile.enabled_for_model(main_model_name):
             continue
-        sub_llm_config = config.get_model_config(model_name)
-        sub_clients[sub_agent_type] = create_llm_client(sub_llm_config)
 
-    return LLMClients(main=main_client, sub_clients=sub_clients)
+        def _factory(model_name_for_factory: str = model_name) -> LLMClientABC:
+            sub_llm_config = config.get_model_config(model_name_for_factory)
+            return create_llm_client(sub_llm_config)
+
+        clients.register_sub_client_factory(sub_agent_type, _factory)
+
+    return clients

klaude_code/core/manager/sub_agent_manager.py

@@ -43,7 +43,7 @@ class SubAgentManager:
         child_session = Session(work_dir=parent_session.work_dir)
         child_session.sub_agent_state = state
 
-        child_profile = self._model_profile_provider.build_profile(
+        child_profile = self._model_profile_provider.build_profile_eager(
             self._llm_clients.get_client(state.sub_agent_type),
             state.sub_agent_type,
         )
@@ -58,13 +58,17 @@ class SubAgentManager:
         try:
             # Not emit the subtask's user input since task tool call is already rendered
             result: str = ""
+            task_metadata: model.TaskMetadata | None = None
             sub_agent_input = model.UserInputPayload(text=state.sub_agent_prompt, images=None)
             async for event in child_agent.run_task(sub_agent_input):
                 # Capture TaskFinishEvent content for return
                 if isinstance(event, events.TaskFinishEvent):
                     result = event.task_result
+                # Capture TaskMetadataEvent for metadata propagation
+                elif isinstance(event, events.TaskMetadataEvent):
+                    task_metadata = event.metadata.main
                 await self.emit_event(event)
-            return SubAgentResult(task_result=result, session_id=child_session.id)
+            return SubAgentResult(task_result=result, session_id=child_session.id, task_metadata=task_metadata)
         except asyncio.CancelledError:
             # Propagate cancellation so tooling can treat it as user interrupt
             log_debug(