klaude-code 2.0.1__py3-none-any.whl → 2.1.0__py3-none-any.whl

klaude_code/core/manager/llm_clients_builder.py

@@ -6,9 +6,9 @@ from klaude_code.config import Config
 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
-from klaude_code.trace import DebugType, log_debug
 
 
 def build_llm_clients(

klaude_code/core/manager/sub_agent_manager.py

@@ -9,15 +9,15 @@ from __future__ import annotations
 
 import asyncio
 import json
+from collections.abc import Callable
 
-from klaude_code.core.agent import Agent, AgentProfile, ModelProfileProvider
+from klaude_code.core.agent import Agent
+from klaude_code.core.agent_profile import ModelProfileProvider
 from klaude_code.core.manager.llm_clients import LLMClients
-from klaude_code.core.tool import ReportBackTool
-from klaude_code.core.tool.tool_context import record_sub_agent_session_id
+from klaude_code.log import DebugType, log_debug
 from klaude_code.protocol import events, message, model
 from klaude_code.protocol.sub_agent import SubAgentResult
 from klaude_code.session.session import Session
-from klaude_code.trace import DebugType, log_debug
 
 
 class SubAgentManager:
@@ -38,7 +38,13 @@ class SubAgentManager:
 
         await self._event_queue.put(event)
 
-    async def run_sub_agent(self, parent_agent: Agent, state: model.SubAgentState) -> SubAgentResult:
+    async def run_sub_agent(
+        self,
+        parent_agent: Agent,
+        state: model.SubAgentState,
+        *,
+        record_session_id: Callable[[str], None] | None = None,
+    ) -> SubAgentResult:
         """Run a nested sub-agent task and return its result."""
 
         parent_session = parent_agent.session
@@ -77,9 +83,8 @@ class SubAgentManager:
                     error=True,
                 )
 
-            # Expose the session id immediately so ToolExecutor.cancel() can attach
-            # it to the synthesized cancellation ToolResult.
-            record_sub_agent_session_id(child_session.id)
+            if record_session_id is not None:
+                record_session_id(child_session.id)
 
             # Update persisted sub-agent state to reflect the current invocation.
             child_session.sub_agent_state.sub_agent_desc = state.sub_agent_desc
@@ -92,33 +97,15 @@ class SubAgentManager:
             child_session = Session(work_dir=parent_session.work_dir)
             child_session.sub_agent_state = state
 
-            # Expose the new session id immediately so ToolExecutor.cancel() can attach
-            # it to the synthesized cancellation ToolResult.
-            record_sub_agent_session_id(child_session.id)
+            if record_session_id is not None:
+                record_session_id(child_session.id)
 
         child_profile = self._model_profile_provider.build_profile(
             self._llm_clients.get_client(state.sub_agent_type),
             state.sub_agent_type,
+            output_schema=state.output_schema,
         )
 
-        # Inject report_back tool if output_schema is provided
-        if state.output_schema:
-            report_back_tool_class = ReportBackTool.for_schema(state.output_schema)
-            report_back_prompt = """\
-
-# Structured Output
-You have a `report_back` tool available. When you complete the task,\
-you MUST call `report_back` with the structured result matching the required schema.\
-Only the content passed to `report_back` will be returned to user.\
-"""
-            base_prompt = child_profile.system_prompt or ""
-            child_profile = AgentProfile(
-                llm_client=child_profile.llm_client,
-                system_prompt=base_prompt + report_back_prompt,
-                tools=[*child_profile.tools, report_back_tool_class.schema()],
-                reminders=child_profile.reminders,
-            )
-
         child_agent = Agent(session=child_session, profile=child_profile)
 
         log_debug(

klaude_code/core/reminders.py

@@ -1,22 +1,19 @@
 import hashlib
 import re
 import shlex
-from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from pathlib import Path
 
 from pydantic import BaseModel
 
 from klaude_code.const import MEMORY_FILE_NAMES, REMINDER_COOLDOWN_TURNS, TODO_REMINDER_TOOL_CALL_THRESHOLD
-from klaude_code.core.tool import BashTool, ReadTool, reset_tool_context, set_tool_context_from_session
+from klaude_code.core.tool import BashTool, ReadTool, build_todo_context
+from klaude_code.core.tool.context import ToolContext
 from klaude_code.core.tool.file._utils import hash_text_sha256
 from klaude_code.protocol import message, model, tools
 from klaude_code.session import Session
 from klaude_code.skill import get_skill
 
-type Reminder = Callable[[Session], Awaitable[message.DeveloperMessage | None]]
-
-
 # Match @ preceded by whitespace, start of line, or → (ReadTool line number arrow)
 AT_FILE_PATTERN = re.compile(r'(?:(?<!\S)|(?<=\u2192))@("(?P<quoted>[^\"]+)"|(?P<plain>\S+))')
 
@@ -71,10 +68,13 @@ def get_at_patterns_with_source(session: Session) -> list[AtPatternSource]:
                 patterns.append(AtPatternSource(pattern=path_str, mentioned_in=None))
             break
 
-        if isinstance(item, message.DeveloperMessage) and item.memory_mentioned:
-            for memory_path, mentioned_patterns in item.memory_mentioned.items():
-                for pattern in mentioned_patterns:
-                    patterns.append(AtPatternSource(pattern=pattern, mentioned_in=memory_path))
+        if isinstance(item, message.DeveloperMessage) and item.ui_extra:
+            for ui_item in item.ui_extra.items:
+                if not isinstance(ui_item, model.MemoryLoadedUIItem):
+                    continue
+                for mem in ui_item.files:
+                    for pattern in mem.mentioned_patterns:
+                        patterns.append(AtPatternSource(pattern=pattern, mentioned_in=mem.path))
     return patterns
 
 
@@ -113,7 +113,8 @@ def _is_tracked_file_unchanged(session: Session, path: str) -> bool:
 async def _load_at_file_recursive(
     session: Session,
     pattern: str,
-    at_files: dict[str, model.AtPatternParseResult],
+    at_ops: list[model.AtFileOp],
+    formatted_blocks: list[str],
     collected_images: list[message.ImageURLPart],
     visited: set[str],
     base_dir: Path | None = None,
@@ -127,53 +128,59 @@ async def _load_at_file_recursive(
         return
     visited.add(path_str)
 
-    context_token = set_tool_context_from_session(session)
-    try:
-        if path.exists() and path.is_file():
-            if _is_tracked_file_unchanged(session, path_str):
-                return
-            args = ReadTool.ReadArguments(file_path=path_str)
-            tool_result = await ReadTool.call_with_args(args)
-            images = [part for part in tool_result.parts if isinstance(part, message.ImageURLPart)]
-            at_files[path_str] = model.AtPatternParseResult(
-                path=path_str,
-                tool_name=tools.READ,
-                result=tool_result.output_text,
-                tool_args=args.model_dump_json(exclude_none=True),
-                operation="Read",
-                mentioned_in=mentioned_in,
-            )
-            if images:
-                collected_images.extend(images)
-
-            # Recursively parse @ references from ReadTool output
-            output = tool_result.output_text
-            if "@" in output:
-                for match in AT_FILE_PATTERN.finditer(output):
-                    nested = match.group("quoted") or match.group("plain")
-                    if nested:
-                        await _load_at_file_recursive(
-                            session,
-                            nested,
-                            at_files,
-                            collected_images,
-                            visited,
-                            base_dir=path.parent,
-                            mentioned_in=path_str,
-                        )
-        elif path.exists() and path.is_dir():
-            quoted_path = shlex.quote(path_str)
-            args = BashTool.BashArguments(command=f"ls {quoted_path}")
-            tool_result = await BashTool.call_with_args(args)
-            at_files[path_str] = model.AtPatternParseResult(
-                path=path_str + "/",
-                tool_name=tools.BASH,
-                result=tool_result.output_text,
-                tool_args=args.model_dump_json(exclude_none=True),
-                operation="List",
-            )
-    finally:
-        reset_tool_context(context_token)
+    tool_context = ToolContext(
+        file_tracker=session.file_tracker,
+        todo_context=build_todo_context(session),
+        session_id=session.id,
+    )
+
+    if path.exists() and path.is_file():
+        if _is_tracked_file_unchanged(session, path_str):
+            return
+        args = ReadTool.ReadArguments(file_path=path_str)
+        tool_result = await ReadTool.call_with_args(args, tool_context)
+        images = [part for part in tool_result.parts if isinstance(part, message.ImageURLPart)]
+
+        tool_args = args.model_dump_json(exclude_none=True)
+        formatted_blocks.append(
+            f"""Called the {tools.READ} tool with the following input: {tool_args}
+Result of calling the {tools.READ} tool:
+{tool_result.output_text}
+"""
+        )
+        at_ops.append(model.AtFileOp(operation="Read", path=path_str, mentioned_in=mentioned_in))
+        if images:
+            collected_images.extend(images)
+
+        # Recursively parse @ references from ReadTool output
+        output = tool_result.output_text
+        if "@" in output:
+            for match in AT_FILE_PATTERN.finditer(output):
+                nested = match.group("quoted") or match.group("plain")
+                if nested:
+                    await _load_at_file_recursive(
+                        session,
+                        nested,
+                        at_ops,
+                        formatted_blocks,
+                        collected_images,
+                        visited,
+                        base_dir=path.parent,
+                        mentioned_in=path_str,
+                    )
+    elif path.exists() and path.is_dir():
+        quoted_path = shlex.quote(path_str)
+        args = BashTool.BashArguments(command=f"ls {quoted_path}")
+        tool_result = await BashTool.call_with_args(args, tool_context)
+
+        tool_args = args.model_dump_json(exclude_none=True)
+        formatted_blocks.append(
+            f"""Called the {tools.BASH} tool with the following input: {tool_args}
+Result of calling the {tools.BASH} tool:
+{tool_result.output_text}
+"""
+        )
+        at_ops.append(model.AtFileOp(operation="List", path=path_str + "/", mentioned_in=mentioned_in))
 
 
 async def at_file_reader_reminder(
@@ -184,7 +191,8 @@ async def at_file_reader_reminder(
     if not at_pattern_sources:
         return None
 
-    at_files: dict[str, model.AtPatternParseResult] = {}  # path -> content
+    at_ops: list[model.AtFileOp] = []
+    formatted_blocks: list[str] = []
     collected_images: list[message.ImageURLPart] = []
     visited: set[str] = set()
 
@@ -192,30 +200,23 @@ async def at_file_reader_reminder(
         await _load_at_file_recursive(
             session,
             source.pattern,
-            at_files,
+            at_ops,
+            formatted_blocks,
             collected_images,
             visited,
             mentioned_in=source.mentioned_in,
         )
 
-    if len(at_files) == 0:
+    if len(formatted_blocks) == 0:
         return None
 
-    at_files_str = "\n\n".join(
-        [
-            f"""Called the {result.tool_name} tool with the following input: {result.tool_args}
-Result of calling the {result.tool_name} tool:
-{result.result}
-"""
-            for result in at_files.values()
-        ]
-    )
+    at_files_str = "\n\n".join(formatted_blocks)
     return message.DeveloperMessage(
         parts=message.parts_from_text_and_images(
             f"""<system-reminder>{at_files_str}\n</system-reminder>""",
             collected_images or None,
         ),
-        at_files=list(at_files.values()),
+        ui_extra=model.DeveloperUIExtra(items=[model.AtFileOpsUIItem(ops=at_ops)]),
     )
 
 
@@ -299,7 +300,7 @@ Here are the existing contents of your todo list:
 
 {model.todo_list_str(session.todos)}</system-reminder>"""
             ),
-            todo_use=True,
+            ui_extra=model.DeveloperUIExtra(items=[model.TodoReminderUIItem(reason="not_used_recently")]),
         )
 
     if session.need_todo_not_used_cooldown_counter > 0:
@@ -327,18 +328,20 @@ async def file_changed_externally_reminder(
                     changed = current_mtime != status.mtime
 
                 if changed:
-                    context_token = set_tool_context_from_session(session)
-                    try:
-                        tool_result = await ReadTool.call_with_args(
-                            ReadTool.ReadArguments(file_path=path)
-                        )  # This tool will update file tracker
-                        if tool_result.status == "success":
-                            images = [part for part in tool_result.parts if isinstance(part, message.ImageURLPart)]
-                            changed_files.append((path, tool_result.output_text, images or None))
-                            if images:
-                                collected_images.extend(images)
-                    finally:
-                        reset_tool_context(context_token)
+                    tool_context = ToolContext(
+                        file_tracker=session.file_tracker,
+                        todo_context=build_todo_context(session),
+                        session_id=session.id,
+                    )
+                    tool_result = await ReadTool.call_with_args(
+                        ReadTool.ReadArguments(file_path=path),
+                        tool_context,
+                    )  # This tool will update file tracker
+                    if tool_result.status == "success":
+                        images = [part for part in tool_result.parts if isinstance(part, message.ImageURLPart)]
+                        changed_files.append((path, tool_result.output_text, images or None))
+                        if images:
+                            collected_images.extend(images)
             except (
                 FileNotFoundError,
                 IsADirectoryError,
@@ -357,10 +360,12 @@ async def file_changed_externally_reminder(
         )
         return message.DeveloperMessage(
             parts=message.parts_from_text_and_images(
-                f"""<system-reminder>{changed_files_str}""",
+                f"""<system-reminder>{changed_files_str}</system-reminder>""",
                 collected_images or None,
             ),
-            external_file_changes=[file_path for file_path, _, _ in changed_files],
+            ui_extra=model.DeveloperUIExtra(
+                items=[model.ExternalFileChangesUIItem(paths=[file_path for file_path, _, _ in changed_files])]
+            ),
         )
 
     return None
@@ -423,10 +428,8 @@ async def image_reminder(session: Session) -> message.DeveloperMessage | None:
         return None
 
     return message.DeveloperMessage(
-        parts=message.text_parts_from_str(
-            f"<system-reminder>User attached {image_count} image{'s' if image_count > 1 else ''} in their message. Make sure to analyze and reference these images as needed.</system-reminder>"
-        ),
-        user_image_count=image_count,
+        parts=[],
+        ui_extra=model.DeveloperUIExtra(items=[model.UserImagesUIItem(count=image_count)]),
     )
 
 
@@ -456,7 +459,7 @@ async def skill_reminder(session: Session) -> message.DeveloperMessage | None:
 
     return message.DeveloperMessage(
         parts=message.text_parts_from_str(content),
-        skill_name=skill.name,
+        ui_extra=model.DeveloperUIExtra(items=[model.SkillActivatedUIItem(name=skill.name)]),
     )
 
 
@@ -487,7 +490,7 @@ async def memory_reminder(session: Session) -> message.DeveloperMessage | None:
         path_str = str(memory_path)
         if memory_path.exists() and memory_path.is_file() and not _is_memory_loaded(session, path_str):
             try:
-                text = memory_path.read_text()
+                text = memory_path.read_text(encoding="utf-8", errors="replace")
                 _mark_memory_loaded(session, path_str)
                 memories.append(Memory(path=path_str, instruction=instruction, content=text))
             except (PermissionError, UnicodeDecodeError, OSError):
@@ -496,32 +499,19 @@ async def memory_reminder(session: Session) -> message.DeveloperMessage | None:
         memories_str = "\n\n".join(
             [f"Contents of {memory.path} ({memory.instruction}):\n\n{memory.content}" for memory in memories]
         )
-        # Build memory_mentioned: extract @ patterns from each memory's content
-        memory_mentioned: dict[str, list[str]] = {}
-        for memory in memories:
-            patterns = _extract_at_patterns(memory.content)
-            if patterns:
-                memory_mentioned[memory.path] = patterns
-
+        loaded_files = [
+            model.MemoryFileLoaded(path=memory.path, mentioned_patterns=_extract_at_patterns(memory.content))
+            for memory in memories
+        ]
         return message.DeveloperMessage(
             parts=message.text_parts_from_str(
-                f"""<system-reminder>As you answer the user's questions, you can use the following context:
+                f"""<system-reminder>
+Loaded memory files. Follow these instructions. Do not mention them to the user unless explicitly asked.
 
-# claudeMd
-Codebase and user instructions are shown below. Be sure to adhere to these instructions. IMPORTANT: These instructions OVERRIDE any default behavior and you MUST follow them exactly as written.
 {memories_str}
-
-#important-instruction-reminders
-Do what has been asked; nothing more, nothing less.
-NEVER create files unless they're absolutely necessary for achieving your goal.
-ALWAYS prefer editing an existing file to creating a new one.
-NEVER proactively create documentation files (*.md) or README files. Only create documentation files if explicitly requested by the User.
-
-IMPORTANT: this context may or may not be relevant to your tasks. You should not respond to this context unless it is highly relevant to your task.
 </system-reminder>"""
             ),
-            memory_paths=[memory.path for memory in memories],
-            memory_mentioned=memory_mentioned or None,
+            ui_extra=model.DeveloperUIExtra(items=[model.MemoryLoadedUIItem(files=loaded_files)]),
         )
     return None
 
@@ -572,7 +562,7 @@ async def last_path_memory_reminder(
                     continue
                 if mem_path.exists() and mem_path.is_file():
                     try:
-                        text = mem_path.read_text()
+                        text = mem_path.read_text(encoding="utf-8", errors="replace")
                     except (PermissionError, UnicodeDecodeError, OSError):
                         continue
                     _mark_memory_loaded(session, mem_path_str)
@@ -589,20 +579,16 @@ async def last_path_memory_reminder(
         memories_str = "\n\n".join(
             [f"Contents of {memory.path} ({memory.instruction}):\n\n{memory.content}" for memory in memories]
         )
-        # Build memory_mentioned: extract @ patterns from each memory's content
-        memory_mentioned: dict[str, list[str]] = {}
-        for memory in memories:
-            patterns = _extract_at_patterns(memory.content)
-            if patterns:
-                memory_mentioned[memory.path] = patterns
-
+        loaded_files = [
+            model.MemoryFileLoaded(path=memory.path, mentioned_patterns=_extract_at_patterns(memory.content))
+            for memory in memories
+        ]
         return message.DeveloperMessage(
             parts=message.text_parts_from_str(
                 f"""<system-reminder>{memories_str}
 </system-reminder>"""
             ),
-            memory_paths=[memory.path for memory in memories],
-            memory_mentioned=memory_mentioned or None,
+            ui_extra=model.DeveloperUIExtra(items=[model.MemoryLoadedUIItem(files=loaded_files)]),
         )
 
 
@@ -616,37 +602,3 @@ ALL_REMINDERS = [
     image_reminder,
     skill_reminder,
 ]
-
-
-def load_agent_reminders(
-    model_name: str, sub_agent_type: str | None = None, *, vanilla: bool = False
-) -> list[Reminder]:
-    """Get reminders for an agent based on model and agent type.
-
-    Args:
-        model_name: The model name.
-        sub_agent_type: If None, returns main agent reminders. Otherwise returns sub-agent reminders.
-        vanilla: If True, returns minimal vanilla reminders (ignores sub_agent_type).
-    """
-    if vanilla:
-        return [at_file_reader_reminder]
-
-    reminders: list[Reminder] = []
-
-    # Only main agent (not sub-agent) gets todo reminders, and not for GPT-5
-    if sub_agent_type is None and "gpt-5" not in model_name:
-        reminders.append(empty_todo_reminder)
-        reminders.append(todo_not_used_recently_reminder)
-
-    reminders.extend(
-        [
-            memory_reminder,
-            at_file_reader_reminder,
-            last_path_memory_reminder,
-            file_changed_externally_reminder,
-            image_reminder,
-            skill_reminder,
-        ]
-    )
-
-    return reminders

klaude_code/core/task.py

@@ -4,17 +4,14 @@ import asyncio
 import time
 from collections.abc import AsyncGenerator, Callable, Sequence
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
 
 from klaude_code.const import INITIAL_RETRY_DELAY_S, MAX_FAILED_TURN_RETRIES, MAX_RETRY_DELAY_S
-from klaude_code.core.reminders import Reminder
+from klaude_code.core.agent_profile import AgentProfile, Reminder
 from klaude_code.core.tool import FileTracker, TodoContext, ToolABC
+from klaude_code.core.tool.context import RunSubtask
 from klaude_code.core.turn import TurnError, TurnExecutionContext, TurnExecutor
+from klaude_code.log import DebugType, log_debug
 from klaude_code.protocol import events, message, model
-from klaude_code.trace import DebugType, log_debug
-
-if TYPE_CHECKING:
-    from klaude_code.core.agent import AgentProfile
 
 
 class MetadataAccumulator:
@@ -101,6 +98,7 @@ class SessionContext:
     append_history: Callable[[Sequence[message.HistoryEvent]], None]
     file_tracker: FileTracker
     todo_context: TodoContext
+    run_subtask: RunSubtask | None
 
 
 @dataclass
@@ -189,26 +187,20 @@ class TaskExecutor:
                 self._current_turn = turn
 
                 try:
-                    async for turn_event in turn.run():
-                        match turn_event:
-                            case events.AssistantMessageEvent() as am:
+                    async for e in turn.run():
+                        match e:
+                            case events.ResponseCompleteEvent() as am:
                                 yield am
-                            case events.ResponseMetadataEvent() as e:
-                                metadata_accumulator.add(e.metadata)
-                                # Emit context usage event if available
-                                context_percent = e.metadata.context_usage_percent
-                                if context_percent is not None:
-                                    yield events.ContextUsageEvent(
-                                        session_id=session_ctx.session_id,
-                                        context_percent=context_percent,
-                                    )
+                            case events.UsageEvent() as e:
+                                metadata_accumulator.add(e.usage)
+                                yield e
                             case events.ToolResultEvent() as e:
                                 # Collect sub-agent task metadata from tool results
                                 if e.task_metadata is not None:
                                     metadata_accumulator.add_sub_agent_metadata(e.task_metadata)
-                                yield turn_event
+                                yield e
                             case _:
-                                yield turn_event
+                                yield e
 
                     turn_succeeded = True
                     break

klaude_code/core/tool/__init__.py

@@ -1,7 +1,7 @@
+from .context import FileTracker, RunSubtask, SubAgentResumeClaims, TodoContext, ToolContext, build_todo_context
 from .file.apply_patch import DiffError, process_patch
 from .file.apply_patch_tool import ApplyPatchTool
 from .file.edit_tool import EditTool
-from .file.move_tool import MoveTool
 from .file.read_tool import ReadTool
 from .file.write_tool import WriteTool
 from .report_back_tool import ReportBackTool
@@ -12,17 +12,7 @@ from .sub_agent_tool import SubAgentTool
 from .todo.todo_write_tool import TodoWriteTool
 from .todo.update_plan_tool import UpdatePlanTool
 from .tool_abc import ToolABC
-from .tool_context import (
-    FileTracker,
-    TodoContext,
-    ToolContextToken,
-    build_todo_context,
-    current_run_subtask_callback,
-    reset_tool_context,
-    set_tool_context_from_session,
-    tool_context,
-)
-from .tool_registry import get_registry, get_tool_schemas, load_agent_tools
+from .tool_registry import get_registry, get_tool_schemas
 from .tool_runner import run_tool
 from .truncation import SimpleTruncationStrategy, TruncationStrategy, get_truncation_strategy, set_truncation_strategy
 from .web.mermaid_tool import MermaidTool
@@ -36,33 +26,29 @@ __all__ = [
     "EditTool",
     "FileTracker",
     "MermaidTool",
-    "MoveTool",
     "ReadTool",
     "ReportBackTool",
+    "RunSubtask",
     "SafetyCheckResult",
     "SimpleTruncationStrategy",
     "SkillTool",
+    "SubAgentResumeClaims",
     "SubAgentTool",
     "TodoContext",
     "TodoWriteTool",
     "ToolABC",
-    "ToolContextToken",
+    "ToolContext",
     "TruncationStrategy",
     "UpdatePlanTool",
     "WebFetchTool",
     "WebSearchTool",
     "WriteTool",
     "build_todo_context",
-    "current_run_subtask_callback",
     "get_registry",
     "get_tool_schemas",
     "get_truncation_strategy",
     "is_safe_command",
-    "load_agent_tools",
     "process_patch",
-    "reset_tool_context",
     "run_tool",
-    "set_tool_context_from_session",
     "set_truncation_strategy",
-    "tool_context",
 ]