klaude-code 1.2.1__py3-none-any.whl → 1.2.3__py3-none-any.whl

klaude_code/core/executor.py

@@ -5,30 +5,90 @@ This module implements the submission_loop equivalent for klaude,
 handling operations submitted from the CLI and coordinating with agents.
 """
 
+from __future__ import annotations
+
 import asyncio
 from dataclasses import dataclass
-from uuid import uuid4
+from dataclasses import field as dataclass_field
 
-from klaude_code.command import dispatch_command
+from klaude_code.command import InputAction, InputActionType, dispatch_command
+from klaude_code.config import Config, load_config
 from klaude_code.core.agent import Agent, DefaultModelProfileProvider, ModelProfileProvider
-from klaude_code.llm import LLMClients
-from klaude_code.core.sub_agent import SubAgentResult
-from klaude_code.core.tool.tool_context import current_run_subtask_callback
-from klaude_code.protocol import events, model
-from klaude_code.protocol.op import (
-    EndOperation,
-    InitAgentOperation,
-    InterruptOperation,
-    Operation,
-    Submission,
-    UserInputOperation,
-)
+from klaude_code.core.tool import current_run_subtask_callback
+from klaude_code.llm.client import LLMClientABC
+from klaude_code.llm.registry import create_llm_client
+from klaude_code.protocol import commands, events, model, op
+from klaude_code.protocol.op_handler import OperationHandler
+from klaude_code.protocol.sub_agent import SubAgentResult, get_sub_agent_profile
+from klaude_code.protocol.tools import SubAgentType
 from klaude_code.session.session import Session
 from klaude_code.trace import DebugType, log_debug
 
 
+@dataclass
+class LLMClients:
+    """Container for LLM clients used by main agent and sub-agents."""
+
+    main: LLMClientABC
+    sub_clients: dict[SubAgentType, LLMClientABC] = dataclass_field(default_factory=lambda: {})
+
+    def get_client(self, sub_agent_type: SubAgentType | None = None) -> LLMClientABC:
+        """Get client for given sub-agent type, or main client if None."""
+        if sub_agent_type is None:
+            return self.main
+        return self.sub_clients.get(sub_agent_type) or self.main
+
+    @classmethod
+    def from_config(
+        cls,
+        config: Config,
+        model_override: str | None = None,
+        enabled_sub_agents: list[SubAgentType] | None = None,
+    ) -> LLMClients:
+        """Create LLMClients from application config.
+
+        Args:
+            config: Application configuration
+            model_override: Optional model name to override the main model
+            enabled_sub_agents: List of sub-agent types to initialize clients for
+
+        Returns:
+            LLMClients instance
+        """
+        # Resolve main agent LLM config
+        if model_override:
+            llm_config = config.get_model_config(model_override)
+        else:
+            llm_config = config.get_main_model_config()
+
+        log_debug(
+            "Main LLM config",
+            llm_config.model_dump_json(exclude_none=True),
+            style="yellow",
+            debug_type=DebugType.LLM_CONFIG,
+        )
+
+        main_client = create_llm_client(llm_config)
+        sub_clients: dict[SubAgentType, LLMClientABC] = {}
+
+        # 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):
+                continue
+            sub_llm_config = config.get_model_config(model_name)
+            sub_clients[sub_agent_type] = create_llm_client(sub_llm_config)
+
+        return cls(main=main_client, sub_clients=sub_clients)
+
+
 @dataclass
 class ActiveTask:
+    """Track an in-flight task and its owning session."""
+
     task: asyncio.Task[None]
     session_id: str
 
@@ -39,6 +99,8 @@ class ExecutorContext:
 
     This context is passed to operations when they execute, allowing them
     to access shared resources like the event queue and active sessions.
+
+    Implements the OperationHandler protocol via structural subtyping.
     """
 
     def __init__(
@@ -47,8 +109,8 @@ class ExecutorContext:
         llm_clients: LLMClients,
         model_profile_provider: ModelProfileProvider | None = None,
     ):
-        self.event_queue = event_queue
-        self.llm_clients = llm_clients
+        self.event_queue: asyncio.Queue[events.Event] = event_queue
+        self.llm_clients: LLMClients = llm_clients
         self.model_profile_provider: ModelProfileProvider = model_profile_provider or DefaultModelProfileProvider()
 
         # Track active agents by session ID
@@ -60,56 +122,68 @@ class ExecutorContext:
         """Emit an event to the UI display system."""
         await self.event_queue.put(event)
 
-    async def handle_init_agent(self, operation: InitAgentOperation) -> None:
+    async def _ensure_agent(self, session_id: str) -> 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,
+        )
+
+        async for evt in agent.replay_history():
+            await self.emit_event(evt)
+
+        await self.emit_event(
+            events.WelcomeEvent(
+                work_dir=str(session.work_dir),
+                llm_config=self.llm_clients.main.get_llm_config(),
+            )
+        )
+
+        self.active_agents[session_id] = agent
+        log_debug(
+            f"Initialized agent for session: {session_id}",
+            style="cyan",
+            debug_type=DebugType.EXECUTION,
+        )
+        return agent
+
+    async def handle_init_agent(self, operation: op.InitAgentOperation) -> None:
         """Initialize an agent for a session and replay history to UI."""
         if operation.session_id is None:
             raise ValueError("session_id cannot be None")
 
-        # Load or create session first
-        session = Session.load(operation.session_id)
-
-        # Create agent if not exists
-        if operation.session_id not in self.active_agents:
-            profile = self.model_profile_provider.build_profile(self.llm_clients.main)
-            agent = Agent(
-                session=session,
-                profile=profile,
-                model_profile_provider=self.model_profile_provider,
-            )
-            async for evt in agent.replay_history():
-                await self.emit_event(evt)
-            await self.emit_event(
-                events.WelcomeEvent(
-                    work_dir=str(session.work_dir),
-                    llm_config=self.llm_clients.main.get_llm_config(),
-                )
-            )
-            self.active_agents[operation.session_id] = agent
-            log_debug(
-                f"Initialized agent for session: {operation.session_id}",
-                style="cyan",
-                debug_type=DebugType.EXECUTION,
-            )
+        await self._ensure_agent(operation.session_id)
 
-    async def handle_user_input(self, operation: UserInputOperation) -> None:
+    async def handle_user_input(self, operation: op.UserInputOperation) -> None:
         """Handle a user input operation by running it through an agent."""
 
         if operation.session_id is None:
             raise ValueError("session_id cannot be None")
 
-        # Ensure initialized via init_agent
-        if operation.session_id not in self.active_agents:
-            await self.handle_init_agent(InitAgentOperation(id=str(uuid4()), session_id=operation.session_id))
-
-        agent = self.active_agents[operation.session_id]
+        session_id = operation.session_id
+        agent = await self._ensure_agent(session_id)
+        user_input = operation.input
 
         # emit user input event
-        await self.emit_event(events.UserMessageEvent(content=operation.content, session_id=operation.session_id))
+        await self.emit_event(
+            events.UserMessageEvent(content=user_input.text, session_id=session_id, images=user_input.images)
+        )
 
-        result = await dispatch_command(operation.content, agent)
-        if not result.agent_input:
-            # If this command do not need run agent, we should append user message to session history here
-            agent.session.append_history([model.UserMessageItem(content=operation.content)])
+        result = await dispatch_command(user_input.text, agent)
+
+        actions: list[InputAction] = list(result.actions or [])
+
+        has_run_agent_action = any(action.type is InputActionType.RUN_AGENT for action in actions)
+        if not has_run_agent_action:
+            # No async agent task will run, append user message directly
+            agent.session.append_history([model.UserMessageItem(content=user_input.text, images=user_input.images)])
 
         if result.events:
             agent.session.append_history(
@@ -118,15 +192,82 @@ class ExecutorContext:
             for evt in result.events:
                 await self.emit_event(evt)
 
-        if result.agent_input:
-            # Start task to process user input (do NOT await here so the executor loop stays responsive)
+        for action in actions:
+            await self._run_input_action(action, operation, agent)
+
+    async def _run_input_action(self, action: InputAction, operation: op.UserInputOperation, agent: Agent) -> None:
+        if operation.session_id is None:
+            raise ValueError("session_id cannot be None for input actions")
+
+        session_id = operation.session_id
+
+        if action.type == InputActionType.RUN_AGENT:
+            task_input = model.UserInputPayload(text=action.text, images=operation.input.images)
+
+            existing_active = self.active_tasks.get(operation.id)
+            if existing_active is not None and not existing_active.task.done():
+                raise RuntimeError(f"Active task already registered for operation {operation.id}")
+
             task: asyncio.Task[None] = asyncio.create_task(
-                self._run_agent_task(agent, result.agent_input, operation.id, operation.session_id)
+                self._run_agent_task(agent, task_input, operation.id, session_id)
             )
-            self.active_tasks[operation.id] = ActiveTask(task=task, session_id=operation.session_id)
-            # Do not await task here; completion will be tracked by the executor
+            self.active_tasks[operation.id] = ActiveTask(task=task, session_id=session_id)
+            return
+
+        if action.type == InputActionType.CHANGE_MODEL:
+            if not action.model_name:
+                raise ValueError("ChangeModel action requires model_name")
+
+            await self._apply_model_change(agent, action.model_name)
+            return
+
+        if action.type == InputActionType.CLEAR:
+            await self._apply_clear(agent)
+            return
+
+        raise ValueError(f"Unsupported input action type: {action.type}")
+
+    async def _apply_model_change(self, agent: Agent, model_name: str) -> None:
+        config = load_config()
+        if config is None:
+            raise ValueError("Configuration must be initialized before changing model")
+
+        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))
+
+        developer_item = model.DeveloperMessageItem(
+            content=f"switched to model: {model_name}",
+            command_output=model.CommandOutput(command_name=commands.CommandName.MODEL),
+        )
+        agent.session.append_history([developer_item])
+
+        await self.emit_event(events.DeveloperMessageEvent(session_id=agent.session.id, item=developer_item))
+        await self.emit_event(events.WelcomeEvent(llm_config=llm_config, work_dir=str(agent.session.work_dir)))
+
+    async def _apply_clear(self, agent: Agent) -> None:
+        old_session_id = agent.session.id
+
+        # Create a new session instance to replace the current one
+        new_session = Session(work_dir=agent.session.work_dir)
+        new_session.model_name = agent.session.model_name
 
-    async def handle_interrupt(self, operation: InterruptOperation) -> None:
+        # Replace the agent's session with the new one
+        agent.session = new_session
+        agent.session.save()
+
+        # Update the active_agents mapping
+        self.active_agents.pop(old_session_id, None)
+        self.active_agents[new_session.id] = agent
+
+        developer_item = model.DeveloperMessageItem(
+            content="started new conversation",
+            command_output=model.CommandOutput(command_name=commands.CommandName.CLEAR),
+        )
+
+        await self.emit_event(events.DeveloperMessageEvent(session_id=agent.session.id, item=developer_item))
+
+    async def handle_interrupt(self, operation: op.InterruptOperation) -> None:
         """Handle an interrupt by invoking agent.cancel() and cancelling tasks."""
 
         # Determine affected sessions
@@ -170,7 +311,9 @@ class ExecutorContext:
             # Remove from active tasks immediately
             self.active_tasks.pop(task_id, None)
 
-    async def _run_agent_task(self, agent: Agent, user_input: str, task_id: str, session_id: str) -> None:
+    async def _run_agent_task(
+        self, agent: Agent, user_input: model.UserInputPayload, task_id: str, session_id: str
+    ) -> None:
         """
         Run an agent task and forward all events to the UI.
 
@@ -225,7 +368,11 @@ class ExecutorContext:
         finally:
             # Clean up the task from active tasks
             self.active_tasks.pop(task_id, None)
-            log_debug(f"Cleaned up agent task {task_id}", style="cyan", debug_type=DebugType.EXECUTION)
+            log_debug(
+                f"Cleaned up agent task {task_id}",
+                style="cyan",
+                debug_type=DebugType.EXECUTION,
+            )
 
     async def _run_subagent_task(self, parent_agent: Agent, state: model.SubAgentState) -> SubAgentResult:
         """Run a nested sub-agent task and return the final task_result text.
@@ -246,7 +393,6 @@ class ExecutorContext:
         child_agent = Agent(
             session=child_session,
             profile=child_profile,
-            model_profile_provider=self.model_profile_provider,
         )
 
         log_debug(
@@ -258,7 +404,8 @@ class ExecutorContext:
         try:
             # Not emit the subtask's user input since task tool call is already rendered
             result: str = ""
-            async for event in child_agent.run_task(state.sub_agent_prompt):
+            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
@@ -279,7 +426,9 @@ class ExecutorContext:
                 debug_type=DebugType.EXECUTION,
             )
             return SubAgentResult(
-                task_result=f"Subagent task failed: [{e.__class__.__name__}] {str(e)}", session_id="", error=True
+                task_result=f"Subagent task failed: [{e.__class__.__name__}] {str(e)}",
+                session_id="",
+                error=True,
             )
 
 
@@ -298,11 +447,11 @@ class Executor:
         model_profile_provider: ModelProfileProvider | None = None,
     ):
         self.context = ExecutorContext(event_queue, llm_clients, model_profile_provider)
-        self.submission_queue: asyncio.Queue[Submission] = asyncio.Queue()
-        self.running = False
-        self.task_completion_events: dict[str, asyncio.Event] = {}
+        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] = {}
 
-    async def submit(self, operation: Operation) -> str:
+    async def submit(self, operation: op.Operation) -> str:
         """
         Submit an operation to the executor for processing.
 
@@ -313,12 +462,11 @@ class Executor:
             Unique submission ID for tracking
         """
 
-        submission = Submission(id=operation.id, operation=operation)
+        submission = op.Submission(id=operation.id, operation=operation)
         await self.submission_queue.put(submission)
 
         # Create completion event for tracking
-        completion_event = asyncio.Event()
-        self.task_completion_events[operation.id] = completion_event
+        self._completion_events[operation.id] = asyncio.Event()
 
         log_debug(
             f"Submitted operation {operation.type} with ID {operation.id}",
@@ -328,12 +476,17 @@ class Executor:
 
         return operation.id
 
-    async def wait_for_completion(self, submission_id: str) -> None:
+    async def wait_for(self, submission_id: str) -> None:
         """Wait for a specific submission to complete."""
-        if submission_id in self.task_completion_events:
-            await self.task_completion_events[submission_id].wait()
-            # Clean up the completion event
-            self.task_completion_events.pop(submission_id, None)
+        event = self._completion_events.get(submission_id)
+        if event is not None:
+            await event.wait()
+            self._completion_events.pop(submission_id, None)
+
+    async def submit_and_wait(self, operation: op.Operation) -> None:
+        """Submit an operation and wait for it to complete."""
+        submission_id = await self.submit(operation)
+        await self.wait_for(submission_id)
 
     async def start(self) -> None:
         """
@@ -342,17 +495,15 @@ class Executor:
         This method runs continuously, processing submissions from the queue
         until the executor is stopped.
         """
-        self.running = True
-
         log_debug("Executor started", style="green", debug_type=DebugType.EXECUTION)
 
-        while self.running:
+        while True:
             try:
                 # Wait for next submission
                 submission = await self.submission_queue.get()
 
                 # Check for end operation to gracefully exit
-                if isinstance(submission.operation, EndOperation):
+                if isinstance(submission.operation, op.EndOperation):
                     log_debug(
                         "Received EndOperation, stopping executor",
                         style="yellow",
@@ -369,15 +520,17 @@ class Executor:
 
             except Exception as e:
                 # Handle unexpected errors
-                log_debug(f"Executor error: {str(e)}", style="red", debug_type=DebugType.EXECUTION)
+                log_debug(
+                    f"Executor error: {str(e)}",
+                    style="red",
+                    debug_type=DebugType.EXECUTION,
+                )
                 await self.context.emit_event(
                     events.ErrorEvent(error_message=f"Executor error: {str(e)}", can_retry=False)
                 )
 
     async def stop(self) -> None:
         """Stop the executor and clean up resources."""
-        self.running = False
-
         # Cancel all active tasks and collect them for awaiting
         tasks_to_await: list[asyncio.Task[None]] = []
         for active in self.context.active_tasks.values():
@@ -395,15 +548,19 @@ class Executor:
 
         # Send EndOperation to wake up the start() loop
         try:
-            end_operation = EndOperation()
-            submission = Submission(id=end_operation.id, operation=end_operation)
+            end_operation = op.EndOperation()
+            submission = op.Submission(id=end_operation.id, operation=end_operation)
             await self.submission_queue.put(submission)
         except Exception as e:
-            log_debug(f"Failed to send EndOperation: {str(e)}", style="red", debug_type=DebugType.EXECUTION)
+            log_debug(
+                f"Failed to send EndOperation: {str(e)}",
+                style="red",
+                debug_type=DebugType.EXECUTION,
+            )
 
         log_debug("Executor stopped", style="yellow", debug_type=DebugType.EXECUTION)
 
-    async def _handle_submission(self, submission: Submission) -> None:
+    async def _handle_submission(self, submission: op.Submission) -> None:
         """
         Handle a single submission by executing its operation.
 
@@ -418,23 +575,28 @@ class Executor:
             )
 
             # Execute to spawn the agent task in context
-            await submission.operation.execute(self.context)
+            await submission.operation.execute(handler=self.context)
 
             async def _await_agent_and_complete() -> None:
-                try:
-                    # Wait for the agent task tied to this submission id
-                    active = self.context.active_tasks.get(submission.id)
-                    if active is not None:
+                # Wait for the agent task tied to this submission id
+                active = self.context.active_tasks.get(submission.id)
+                if active is not None:
+                    try:
                         await active.task
-                finally:
-                    # Signal completion of this submission when agent task completes
-                    if submission.id in self.task_completion_events:
-                        self.task_completion_events[submission.id].set()
-
-                # Run in background so the submission loop can continue (e.g., to handle interrupts)
+                    finally:
+                        event = self._completion_events.get(submission.id)
+                        if event is not None:
+                            event.set()
 
+            # Run in background so the submission loop can continue (e.g., to handle interrupts)
             asyncio.create_task(_await_agent_and_complete())
 
+            # For operations without ActiveTask (e.g., InitAgentOperation), signal completion immediately
+            if submission.id not in self.context.active_tasks:
+                event = self._completion_events.get(submission.id)
+                if event is not None:
+                    event.set()
+
         except Exception as e:
             log_debug(
                 f"Failed to handle submission {submission.id}: {str(e)}",
@@ -445,6 +607,11 @@ class Executor:
                 events.ErrorEvent(error_message=f"Operation failed: {str(e)}", can_retry=False)
             )
             # Set completion event even on error to prevent wait_for_completion from hanging
-            completion_event = self.task_completion_events.get(submission.id)
-            if completion_event is not None:
-                completion_event.set()
+            event = self._completion_events.get(submission.id)
+            if event is not None:
+                event.set()
+
+
+# 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]

klaude_code/core/prompt.py

@@ -13,20 +13,20 @@ COMMAND_DESCRIPTIONS: dict[str, str] = {
 
 # Mapping from logical prompt keys to resource file paths under the core/prompt directory.
 PROMPT_FILES: dict[str, str] = {
-    "main_codex": "prompt/prompt-codex.md",
-    "main_claude": "prompt/prompt-claude-code.md",
-    "main_gemini": "prompt/prompt-gemini.md",  # https://ai.google.dev/gemini-api/docs/prompting-strategies?hl=zh-cn#agentic-si-template
+    "main_codex": "prompts/prompt-codex.md",
+    "main_claude": "prompts/prompt-claude-code.md",
+    "main_gemini": "prompts/prompt-gemini.md",  # https://ai.google.dev/gemini-api/docs/prompting-strategies?hl=zh-cn#agentic-si-template
     # Sub-agent prompts keyed by their name
-    "Task": "prompt/prompt-subagent.md",
-    "Oracle": "prompt/prompt-subagent-oracle.md",
-    "Explore": "prompt/prompt-subagent-explore.md",
-    "WebFetchAgent": "prompt/prompt-subagent-webfetch.md",
+    "Task": "prompts/prompt-subagent.md",
+    "Oracle": "prompts/prompt-subagent-oracle.md",
+    "Explore": "prompts/prompt-subagent-explore.md",
+    "WebFetchAgent": "prompts/prompt-subagent-webfetch.md",
 }
 
 
 @lru_cache(maxsize=None)
-def get_system_prompt(model_name: str, key: str = "main") -> str:
-    """Get system prompt content for the given model and prompt key."""
+def get_system_prompt(model_name: str, sub_agent_type: str | None = None) -> str:
+    """Get system prompt content for the given model and sub-agent type."""
 
     cwd = Path.cwd()
     today = datetime.datetime.now().strftime("%Y-%m-%d")
@@ -37,7 +37,7 @@ def get_system_prompt(model_name: str, key: str = "main") -> str:
         if shutil.which(command) is not None:
             available_tools.append(f"{command}: {desc}")
 
-    if key == "main":
+    if sub_agent_type is None:
         match model_name:
             case name if "gpt-5" in name:
                 file_key = "main_codex"
@@ -46,12 +46,12 @@ def get_system_prompt(model_name: str, key: str = "main") -> str:
             case _:
                 file_key = "main_claude"
     else:
-        file_key = key
+        file_key = sub_agent_type
 
     try:
         prompt_path = PROMPT_FILES[file_key]
     except KeyError as exc:
-        raise ValueError(f"Unknown prompt key: {key}") from exc
+        raise ValueError(f"Unknown prompt key: {file_key}") from exc
 
     base_prompt = (
         files(__package__)

klaude_code/core/{prompt → prompts}/prompt-gemini.md

@@ -1,4 +1,4 @@
-You are a very strong reasoner and planner. Use these critical instructions to structure your plans, thoughts, and responses.
+You are an interactive CLI tool that helps users with software engineering tasks. Use the instructions below and the tools available to you to assist the user.
 
 Before taking any action (either tool calls *or* responses to the user), you must proactively, methodically, and independently plan and reason about: