klaude-code 2.0.2__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+))')
 
@@ -131,57 +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)]
-
-            tool_args = args.model_dump_json(exclude_none=True)
-            formatted_blocks.append(
-                f"""Called the {tools.READ} tool with the following input: {tool_args}
+    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_args = args.model_dump_json(exclude_none=True)
-            formatted_blocks.append(
-                f"""Called the {tools.BASH} tool with the following input: {tool_args}
+        )
+        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))
-    finally:
-        reset_tool_context(context_token)
+        )
+        at_ops.append(model.AtFileOp(operation="List", path=path_str + "/", mentioned_in=mentioned_in))
 
 
 async def at_file_reader_reminder(
@@ -329,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,
@@ -427,9 +428,7 @@ 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>"
-        ),
+        parts=[],
         ui_extra=model.DeveloperUIExtra(items=[model.UserImagesUIItem(count=image_count)]),
     )
 
@@ -603,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,3 +1,4 @@
+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
@@ -11,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
@@ -37,30 +28,27 @@ __all__ = [
     "MermaidTool",
     "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",
 ]

klaude_code/core/tool/context.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable, MutableMapping
+from dataclasses import dataclass, replace
+
+from klaude_code.protocol import model
+from klaude_code.protocol.sub_agent import SubAgentResult
+from klaude_code.session.session import Session
+
+type FileTracker = MutableMapping[str, model.FileStatus]
+
+RunSubtask = Callable[[model.SubAgentState, Callable[[str], None] | None], Awaitable[SubAgentResult]]
+
+
+@dataclass
+class TodoContext:
+    """Todo access interface exposed to tools.
+
+    Tools can only read the current todo list and replace it with
+    a new list; they cannot access the full Session object.
+    """
+
+    get_todos: Callable[[], list[model.TodoItem]]
+    set_todos: Callable[[list[model.TodoItem]], None]
+
+
+@dataclass
+class SessionTodoStore:
+    """Adapter exposing session todos through an explicit interface."""
+
+    session: Session
+
+    def get(self) -> list[model.TodoItem]:
+        return self.session.todos
+
+    def set(self, todos: list[model.TodoItem]) -> None:
+        self.session.todos = todos
+
+
+def build_todo_context(session: Session) -> TodoContext:
+    """Create a TodoContext backed by the given session."""
+
+    store = SessionTodoStore(session)
+    return TodoContext(get_todos=store.get, set_todos=store.set)
+
+
+class SubAgentResumeClaims:
+    """Track sub-agent resume claims for a single turn.
+
+    Multiple concurrent sub-agent tool calls can attempt to resume the same
+    session id in a single model response. This class provides an atomic
+    `claim()` operation to reject duplicates.
+    """
+
+    def __init__(self) -> None:
+        self._claims: set[str] = set()
+        self._lock = asyncio.Lock()
+
+    async def claim(self, session_id: str) -> bool:
+        async with self._lock:
+            if session_id in self._claims:
+                return False
+            self._claims.add(session_id)
+            return True
+
+
+@dataclass(frozen=True)
+class ToolContext:
+    """Tool execution context.
+
+    This object is shallow-immutable: fields cannot be reassigned, but fields
+    may reference mutable objects (e.g., FileTracker).
+    """
+
+    file_tracker: FileTracker
+    todo_context: TodoContext
+    session_id: str
+    run_subtask: RunSubtask | None = None
+    sub_agent_resume_claims: SubAgentResumeClaims | None = None
+    record_sub_agent_session_id: Callable[[str], None] | None = None
+
+    def with_record_sub_agent_session_id(self, callback: Callable[[str], None] | None) -> ToolContext:
+        return replace(self, record_sub_agent_session_id=callback)

klaude_code/core/tool/file/apply_patch_tool.py

@@ -7,20 +7,20 @@ from pathlib import Path
 
 from pydantic import BaseModel
 
+from klaude_code.core.tool.context import FileTracker, ToolContext
 from klaude_code.core.tool.file import apply_patch as apply_patch_module
 from klaude_code.core.tool.file._utils import hash_text_sha256
 from klaude_code.core.tool.file.diff_builder import build_structured_file_diff
 from klaude_code.core.tool.tool_abc import ToolABC, load_desc
-from klaude_code.core.tool.tool_context import get_current_file_tracker
 from klaude_code.core.tool.tool_registry import register
 from klaude_code.protocol import llm_param, message, model, tools
 
 
 class ApplyPatchHandler:
     @classmethod
-    async def handle_apply_patch(cls, patch_text: str) -> message.ToolResultMessage:
+    async def handle_apply_patch(cls, patch_text: str, context: ToolContext) -> message.ToolResultMessage:
         try:
-            output, ui_extra = await asyncio.to_thread(cls._apply_patch_in_thread, patch_text)
+            output, ui_extra = await asyncio.to_thread(cls._apply_patch_in_thread, patch_text, context.file_tracker)
         except apply_patch_module.DiffError as error:
             return message.ToolResultMessage(status="error", output_text=str(error))
         except Exception as error:  # pragma: no cover  # unexpected errors bubbled to tool result
@@ -32,14 +32,13 @@ class ApplyPatchHandler:
         )
 
     @staticmethod
-    def _apply_patch_in_thread(patch_text: str) -> tuple[str, model.ToolResultUIExtra]:
+    def _apply_patch_in_thread(patch_text: str, file_tracker: FileTracker) -> tuple[str, model.ToolResultUIExtra]:
         ap = apply_patch_module
         normalized_start = patch_text.lstrip()
         if not normalized_start.startswith("*** Begin Patch"):
             raise ap.DiffError("apply_patch content must start with *** Begin Patch")
 
         workspace_root = os.path.realpath(os.getcwd())
-        file_tracker = get_current_file_tracker()
 
         def resolve_path(path: str) -> str:
             candidate = os.path.realpath(path if os.path.isabs(path) else os.path.join(workspace_root, path))
@@ -89,15 +88,14 @@ class ApplyPatchHandler:
             with open(resolved, "w", encoding="utf-8") as handle:
                 handle.write(content)
 
-            if file_tracker is not None:
-                with contextlib.suppress(Exception):  # pragma: no cover - file tracker best-effort
-                    existing = file_tracker.get(resolved)
-                    is_mem = existing.is_memory if existing else False
-                    file_tracker[resolved] = model.FileStatus(
-                        mtime=Path(resolved).stat().st_mtime,
-                        content_sha256=hash_text_sha256(content),
-                        is_memory=is_mem,
-                    )
+            with contextlib.suppress(Exception):  # pragma: no cover - file tracker best-effort
+                existing = file_tracker.get(resolved)
+                is_mem = existing.is_memory if existing else False
+                file_tracker[resolved] = model.FileStatus(
+                    mtime=Path(resolved).stat().st_mtime,
+                    content_sha256=hash_text_sha256(content),
+                    is_memory=is_mem,
+                )
 
         def remove_fn(path: str) -> None:
             resolved = resolve_path(path)
@@ -107,9 +105,8 @@ class ApplyPatchHandler:
                 raise ap.DiffError(f"Cannot delete directory: {path}")
             os.remove(resolved)
 
-            if file_tracker is not None:
-                with contextlib.suppress(Exception):  # pragma: no cover - file tracker best-effort
-                    file_tracker.pop(resolved, None)
+            with contextlib.suppress(Exception):  # pragma: no cover - file tracker best-effort
+                file_tracker.pop(resolved, None)
 
         ap.apply_commit(commit, write_fn, remove_fn)
 
@@ -172,13 +169,13 @@ class ApplyPatchTool(ToolABC):
         )
 
     @classmethod
-    async def call(cls, arguments: str) -> message.ToolResultMessage:
+    async def call(cls, arguments: str, context: ToolContext) -> message.ToolResultMessage:
         try:
             args = cls.ApplyPatchArguments.model_validate_json(arguments)
         except ValueError as exc:
             return message.ToolResultMessage(status="error", output_text=f"Invalid arguments: {exc}")
-        return await cls.call_with_args(args)
+        return await cls.call_with_args(args, context)
 
     @classmethod
-    async def call_with_args(cls, args: ApplyPatchArguments) -> message.ToolResultMessage:
-        return await ApplyPatchHandler.handle_apply_patch(args.patch)
+    async def call_with_args(cls, args: ApplyPatchArguments, context: ToolContext) -> message.ToolResultMessage:
+        return await ApplyPatchHandler.handle_apply_patch(args.patch, context)