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

klaude_code/core/tool/file/edit_tool.py

@@ -9,10 +9,10 @@ from pathlib import Path
 from pydantic import BaseModel, Field
 
 from klaude_code.const import DIFF_DEFAULT_CONTEXT_LINES
+from klaude_code.core.tool.context import ToolContext
 from klaude_code.core.tool.file._utils import file_exists, hash_text_sha256, is_directory, read_text, write_text
 from klaude_code.core.tool.file.diff_builder import build_structured_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
 
@@ -86,7 +86,7 @@ class EditTool(ToolABC):
         return content.replace(old_string, new_string, 1)
 
     @classmethod
-    async def call(cls, arguments: str) -> message.ToolResultMessage:
+    async def call(cls, arguments: str, context: ToolContext) -> message.ToolResultMessage:
         try:
             args = EditTool.EditArguments.model_validate_json(arguments)
         except ValueError as e:  # pragma: no cover - defensive
@@ -111,7 +111,7 @@ class EditTool(ToolABC):
             )
 
         # FileTracker checks (only for editing existing files)
-        file_tracker = get_current_file_tracker()
+        file_tracker = context.file_tracker
         tracked_status: model.FileStatus | None = None
         if not file_exists(file_path):
             # We require reading before editing
@@ -119,13 +119,12 @@ class EditTool(ToolABC):
                 status="error",
                 output_text=("File has not been read yet. Read it first before writing to it."),
             )
-        if file_tracker is not None:
-            tracked_status = file_tracker.get(file_path)
-            if tracked_status is None:
-                return message.ToolResultMessage(
-                    status="error",
-                    output_text=("File has not been read yet. Read it first before writing to it."),
-                )
+        tracked_status = file_tracker.get(file_path)
+        if tracked_status is None:
+            return message.ToolResultMessage(
+                status="error",
+                output_text=("File has not been read yet. Read it first before writing to it."),
+            )
 
         # Edit existing file: validate and apply
         try:
@@ -137,29 +136,28 @@ class EditTool(ToolABC):
             )
 
         # Re-check external modifications using content hash when available.
-        if tracked_status is not None:
-            if tracked_status.content_sha256 is not None:
-                current_sha256 = hash_text_sha256(before)
-                if current_sha256 != tracked_status.content_sha256:
-                    return message.ToolResultMessage(
-                        status="error",
-                        output_text=(
-                            "File has been modified externally. Either by user or a linter. Read it first before writing to it."
-                        ),
-                    )
-            else:
-                # Backward-compat: old sessions only stored mtime.
-                try:
-                    current_mtime = Path(file_path).stat().st_mtime
-                except OSError:
-                    current_mtime = tracked_status.mtime
-                if current_mtime != tracked_status.mtime:
-                    return message.ToolResultMessage(
-                        status="error",
-                        output_text=(
-                            "File has been modified externally. Either by user or a linter. Read it first before writing to it."
-                        ),
-                    )
+        if tracked_status.content_sha256 is not None:
+            current_sha256 = hash_text_sha256(before)
+            if current_sha256 != tracked_status.content_sha256:
+                return message.ToolResultMessage(
+                    status="error",
+                    output_text=(
+                        "File has been modified externally. Either by user or a linter. Read it first before writing to it."
+                    ),
+                )
+        else:
+            # Backward-compat: old sessions only stored mtime.
+            try:
+                current_mtime = Path(file_path).stat().st_mtime
+            except OSError:
+                current_mtime = tracked_status.mtime
+            if current_mtime != tracked_status.mtime:
+                return message.ToolResultMessage(
+                    status="error",
+                    output_text=(
+                        "File has been modified externally. Either by user or a linter. Read it first before writing to it."
+                    ),
+                )
 
         err = cls.valid(
             content=before,
@@ -205,15 +203,14 @@ class EditTool(ToolABC):
         ui_extra = build_structured_diff(before, after, file_path=file_path)
 
         # Update tracker with new mtime and content hash
-        if file_tracker is not None:
-            with contextlib.suppress(Exception):
-                existing = file_tracker.get(file_path)
-                is_mem = existing.is_memory if existing else False
-                file_tracker[file_path] = model.FileStatus(
-                    mtime=Path(file_path).stat().st_mtime,
-                    content_sha256=hash_text_sha256(after),
-                    is_memory=is_mem,
-                )
+        with contextlib.suppress(Exception):
+            existing = file_tracker.get(file_path)
+            is_mem = existing.is_memory if existing else False
+            file_tracker[file_path] = model.FileStatus(
+                mtime=Path(file_path).stat().st_mtime,
+                content_sha256=hash_text_sha256(after),
+                is_memory=is_mem,
+            )
 
         # Build output message
         if args.replace_all:

klaude_code/core/tool/file/read_tool.py

@@ -17,9 +17,9 @@ from klaude_code.const import (
     READ_MAX_CHARS,
     READ_MAX_IMAGE_BYTES,
 )
+from klaude_code.core.tool.context import FileTracker, ToolContext
 from klaude_code.core.tool.file._utils import file_exists, is_directory
 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
 
@@ -121,8 +121,13 @@ def _read_segment(options: ReadOptions) -> ReadSegmentResult:
     )
 
 
-def _track_file_access(file_path: str, *, content_sha256: str | None = None, is_memory: bool = False) -> None:
-    file_tracker = get_current_file_tracker()
+def _track_file_access(
+    file_tracker: FileTracker | None,
+    file_path: str,
+    *,
+    content_sha256: str | None = None,
+    is_memory: bool = False,
+) -> None:
     if file_tracker is None or not file_exists(file_path) or is_directory(file_path):
         return
     with contextlib.suppress(Exception):
@@ -182,12 +187,12 @@ class ReadTool(ToolABC):
         )
 
     @classmethod
-    async def call(cls, arguments: str) -> message.ToolResultMessage:
+    async def call(cls, arguments: str, context: ToolContext) -> message.ToolResultMessage:
         try:
             args = ReadTool.ReadArguments.model_validate_json(arguments)
         except Exception as e:  # pragma: no cover - defensive
             return message.ToolResultMessage(status="error", output_text=f"Invalid arguments: {e}")
-        return await cls.call_with_args(args)
+        return await cls.call_with_args(args, context)
 
     @classmethod
     def _effective_limits(cls) -> tuple[int | None, int | None, int | None]:
@@ -198,7 +203,7 @@ class ReadTool(ToolABC):
         )
 
     @classmethod
-    async def call_with_args(cls, args: ReadTool.ReadArguments) -> message.ToolResultMessage:
+    async def call_with_args(cls, args: ReadTool.ReadArguments, context: ToolContext) -> message.ToolResultMessage:
         file_path = os.path.abspath(args.file_path)
         char_per_line, line_cap, max_chars = cls._effective_limits()
 
@@ -271,7 +276,7 @@ class ReadTool(ToolABC):
                     output_text=f"<tool_use_error>Failed to read image file: {exc}</tool_use_error>",
                 )
 
-            _track_file_access(file_path, content_sha256=hashlib.sha256(image_bytes).hexdigest())
+            _track_file_access(context.file_tracker, file_path, content_sha256=hashlib.sha256(image_bytes).hexdigest())
             size_kb = size_bytes / 1024.0 if size_bytes else 0.0
             output_text = f"[image] {Path(file_path).name} ({size_kb:.1f}KB)"
             image_part = message.ImageURLPart(url=data_url, id=None)
@@ -308,7 +313,7 @@ class ReadTool(ToolABC):
 
         if offset > max(read_result.total_lines, 0):
             warn = f"<system-reminder>Warning: the file exists but is shorter than the provided offset ({offset}). The file has {read_result.total_lines} lines.</system-reminder>"
-            _track_file_access(file_path, content_sha256=read_result.content_sha256)
+            _track_file_access(context.file_tracker, file_path, content_sha256=read_result.content_sha256)
             return message.ToolResultMessage(status="success", output_text=warn)
 
         lines_out: list[str] = [_format_numbered_line(no, content) for no, content in read_result.selected_lines]
@@ -326,6 +331,6 @@ class ReadTool(ToolABC):
             )
 
         read_result_str = "\n".join(lines_out)
-        _track_file_access(file_path, content_sha256=read_result.content_sha256)
+        _track_file_access(context.file_tracker, file_path, content_sha256=read_result.content_sha256)
 
         return message.ToolResultMessage(status="success", output_text=read_result_str)

klaude_code/core/tool/file/write_tool.py

@@ -7,10 +7,10 @@ from pathlib import Path
 
 from pydantic import BaseModel
 
+from klaude_code.core.tool.context import ToolContext
 from klaude_code.core.tool.file._utils import file_exists, hash_text_sha256, is_directory, read_text, write_text
 from klaude_code.core.tool.file.diff_builder import build_structured_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
 
@@ -46,7 +46,7 @@ class WriteTool(ToolABC):
         )
 
     @classmethod
-    async def call(cls, arguments: str) -> message.ToolResultMessage:
+    async def call(cls, arguments: str, context: ToolContext) -> message.ToolResultMessage:
         try:
             args = WriteArguments.model_validate_json(arguments)
         except ValueError as e:  # pragma: no cover - defensive
@@ -60,12 +60,12 @@ class WriteTool(ToolABC):
                 output_text="<tool_use_error>Illegal operation on a directory. write</tool_use_error>",
             )
 
-        file_tracker = get_current_file_tracker()
+        file_tracker = context.file_tracker
         exists = file_exists(file_path)
         tracked_status: model.FileStatus | None = None
 
         if exists:
-            tracked_status = file_tracker.get(file_path) if file_tracker is not None else None
+            tracked_status = file_tracker.get(file_path)
             if tracked_status is None:
                 return message.ToolResultMessage(
                     status="error",
@@ -114,15 +114,14 @@ class WriteTool(ToolABC):
         except (OSError, UnicodeError) as e:  # pragma: no cover
             return message.ToolResultMessage(status="error", output_text=f"<tool_use_error>{e}</tool_use_error>")
 
-        if file_tracker is not None:
-            with contextlib.suppress(Exception):
-                existing = file_tracker.get(file_path)
-                is_mem = existing.is_memory if existing else False
-                file_tracker[file_path] = model.FileStatus(
-                    mtime=Path(file_path).stat().st_mtime,
-                    content_sha256=hash_text_sha256(args.content),
-                    is_memory=is_mem,
-                )
+        with contextlib.suppress(Exception):
+            existing = file_tracker.get(file_path)
+            is_mem = existing.is_memory if existing else False
+            file_tracker[file_path] = model.FileStatus(
+                mtime=Path(file_path).stat().st_mtime,
+                content_sha256=hash_text_sha256(args.content),
+                is_memory=is_mem,
+            )
 
         # For markdown files, use MarkdownDocUIExtra to render content as markdown
         # Otherwise, build diff between previous and new content

klaude_code/core/tool/report_back_tool.py

@@ -2,6 +2,7 @@
 
 from typing import Any, ClassVar, cast
 
+from klaude_code.core.tool.context import ToolContext
 from klaude_code.protocol import llm_param, message, tools
 
 
@@ -72,7 +73,9 @@ class ReportBackTool:
         )
 
     @classmethod
-    async def call(cls, arguments: str) -> message.ToolResultMessage:
+    async def call(cls, arguments: str, context: ToolContext) -> message.ToolResultMessage:
+        del arguments
+        del context
         """Execute the report_back tool.
 
         The actual handling of report_back results is done by TurnExecutor.

klaude_code/core/tool/shell/bash_tool.py

@@ -11,9 +11,9 @@ from typing import Any
 from pydantic import BaseModel
 
 from klaude_code.const import BASH_DEFAULT_TIMEOUT_MS, BASH_TERMINATE_TIMEOUT_SEC
+from klaude_code.core.tool.context import ToolContext
 from klaude_code.core.tool.shell.command_safety import is_safe_command
 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
 
@@ -71,7 +71,7 @@ class BashTool(ToolABC):
         timeout_ms: int = BASH_DEFAULT_TIMEOUT_MS
 
     @classmethod
-    async def call(cls, arguments: str) -> message.ToolResultMessage:
+    async def call(cls, arguments: str, context: ToolContext) -> message.ToolResultMessage:
         try:
             args = BashTool.BashArguments.model_validate_json(arguments)
         except ValueError as e:
@@ -79,10 +79,10 @@ class BashTool(ToolABC):
                 status="error",
                 output_text=f"Invalid arguments: {e}",
             )
-        return await cls.call_with_args(args)
+        return await cls.call_with_args(args, context)
 
     @classmethod
-    async def call_with_args(cls, args: BashArguments) -> message.ToolResultMessage:
+    async def call_with_args(cls, args: BashArguments, context: ToolContext) -> message.ToolResultMessage:
         # Safety check: only execute commands proven as "known safe"
         result = is_safe_command(args.command)
         if not result.is_safe:
@@ -119,6 +119,8 @@ class BashTool(ToolABC):
             }
         )
 
+        file_tracker = context.file_tracker
+
         def _hash_file_content_sha256(file_path: str) -> str | None:
             try:
                 suffix = Path(file_path).suffix.lower()
@@ -144,9 +146,6 @@ class BashTool(ToolABC):
             return os.path.abspath(os.path.join(base_dir, path))
 
         def _track_files_read(file_paths: list[str], *, base_dir: str) -> None:
-            file_tracker = get_current_file_tracker()
-            if file_tracker is None:
-                return
             for p in file_paths:
                 abs_path = _resolve_in_dir(base_dir, p)
                 if not os.path.exists(abs_path) or os.path.isdir(abs_path):
@@ -168,10 +167,6 @@ class BashTool(ToolABC):
             _track_files_read(file_paths, base_dir=base_dir)
 
         def _track_mv(src_paths: list[str], dest_path: str, *, base_dir: str) -> None:
-            file_tracker = get_current_file_tracker()
-            if file_tracker is None:
-                return
-
             abs_dest = _resolve_in_dir(base_dir, dest_path)
             dest_is_dir = os.path.isdir(abs_dest)
 

klaude_code/core/tool/skill/skill_tool.py

@@ -4,6 +4,7 @@ from pathlib import Path
 
 from pydantic import BaseModel
 
+from klaude_code.core.tool.context import ToolContext
 from klaude_code.core.tool.tool_abc import ToolABC, load_desc
 from klaude_code.core.tool.tool_registry import register
 from klaude_code.protocol import llm_param, message, tools
@@ -55,7 +56,8 @@ class SkillTool(ToolABC):
         command: str
 
     @classmethod
-    async def call(cls, arguments: str) -> message.ToolResultMessage:
+    async def call(cls, arguments: str, context: ToolContext) -> message.ToolResultMessage:
+        del context
         """Load and return full skill content."""
         try:
             args = cls.SkillArguments.model_validate_json(arguments)

klaude_code/core/tool/sub_agent_tool.py

@@ -10,8 +10,8 @@ import asyncio
 import json
 from typing import TYPE_CHECKING, Any, ClassVar, cast
 
+from klaude_code.core.tool.context import ToolContext
 from klaude_code.core.tool.tool_abc import ToolABC, ToolConcurrencyPolicy, ToolMetadata
-from klaude_code.core.tool.tool_context import current_run_subtask_callback, current_sub_agent_resume_claims
 from klaude_code.protocol import llm_param, message, model
 from klaude_code.session.session import Session
 
@@ -52,7 +52,7 @@ class SubAgentTool(ToolABC):
         )
 
     @classmethod
-    async def call(cls, arguments: str) -> message.ToolResultMessage:
+    async def call(cls, arguments: str, context: ToolContext) -> message.ToolResultMessage:
         profile = cls._profile
 
         try:
@@ -60,7 +60,7 @@ class SubAgentTool(ToolABC):
         except json.JSONDecodeError as e:
             return message.ToolResultMessage(status="error", output_text=f"Invalid JSON arguments: {e}")
 
-        runner = current_run_subtask_callback.get()
+        runner = context.run_subtask
         if runner is None:
             return message.ToolResultMessage(status="error", output_text="No subtask runner available in this context")
 
@@ -76,9 +76,10 @@ class SubAgentTool(ToolABC):
             except ValueError as exc:
                 return message.ToolResultMessage(status="error", output_text=str(exc))
 
-            claims = current_sub_agent_resume_claims.get()
+            claims = context.sub_agent_resume_claims
             if claims is not None:
-                if resume_session_id in claims:
+                ok = await claims.claim(resume_session_id)
+                if not ok:
                     return message.ToolResultMessage(
                         status="error",
                         output_text=(
@@ -87,7 +88,6 @@ class SubAgentTool(ToolABC):
                             "Merge into a single call or resume in a later turn."
                         ),
                     )
-                claims.add(resume_session_id)
 
         generation = args.get("generation")
         generation_dict: dict[str, Any] | None = (
@@ -108,7 +108,8 @@ class SubAgentTool(ToolABC):
                     resume=resume_session_id,
                     output_schema=output_schema,
                     generation=generation_dict,
-                )
+                ),
+                context.record_sub_agent_session_id,
             )
         except asyncio.CancelledError:
             raise

klaude_code/core/tool/todo/todo_write_tool.py

@@ -2,8 +2,8 @@ from pathlib import Path
 
 from pydantic import BaseModel
 
+from klaude_code.core.tool.context import ToolContext
 from klaude_code.core.tool.tool_abc import ToolABC, load_desc
-from klaude_code.core.tool.tool_context import get_current_todo_context
 from klaude_code.core.tool.tool_registry import register
 from klaude_code.protocol import llm_param, message, model, tools
 
@@ -76,7 +76,7 @@ class TodoWriteTool(ToolABC):
         )
 
     @classmethod
-    async def call(cls, arguments: str) -> message.ToolResultMessage:
+    async def call(cls, arguments: str, context: ToolContext) -> message.ToolResultMessage:
         try:
             args = TodoWriteArguments.model_validate_json(arguments)
         except ValueError as e:
@@ -85,13 +85,7 @@ class TodoWriteTool(ToolABC):
                 output_text=f"Invalid arguments: {e}",
             )
 
-        # Get current todo context to store todos
-        todo_context = get_current_todo_context()
-        if todo_context is None:
-            return message.ToolResultMessage(
-                status="error",
-                output_text="No active session found",
-            )
+        todo_context = context.todo_context
 
         # Get current todos before updating (for comparison)
         old_todos = todo_context.get_todos()

klaude_code/core/tool/todo/update_plan_tool.py

@@ -6,8 +6,8 @@ from pathlib import Path
 
 from pydantic import BaseModel, field_validator
 
+from klaude_code.core.tool.context import ToolContext
 from klaude_code.core.tool.tool_abc import ToolABC, load_desc
-from klaude_code.core.tool.tool_context import get_current_todo_context
 from klaude_code.core.tool.tool_registry import register
 from klaude_code.protocol import llm_param, message, model, tools
 
@@ -79,15 +79,13 @@ class UpdatePlanTool(ToolABC):
         )
 
     @classmethod
-    async def call(cls, arguments: str) -> message.ToolResultMessage:
+    async def call(cls, arguments: str, context: ToolContext) -> message.ToolResultMessage:
         try:
             args = UpdatePlanArguments.model_validate_json(arguments)
         except ValueError as exc:
             return message.ToolResultMessage(status="error", output_text=f"Invalid arguments: {exc}")
 
-        todo_context = get_current_todo_context()
-        if todo_context is None:
-            return message.ToolResultMessage(status="error", output_text="No active session found")
+        todo_context = context.todo_context
 
         new_todos = [model.TodoItem(content=item.step, status=item.status) for item in args.plan]
         old_todos = todo_context.get_todos()

klaude_code/core/tool/tool_abc.py

@@ -4,6 +4,7 @@ from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
 
+from klaude_code.core.tool.context import ToolContext
 from klaude_code.protocol import llm_param, message
 
 
@@ -27,7 +28,7 @@ class ToolABC(ABC):
 
     @classmethod
     @abstractmethod
-    async def call(cls, arguments: str) -> message.ToolResultMessage:
+    async def call(cls, arguments: str, context: ToolContext) -> message.ToolResultMessage:
         raise NotImplementedError
 
 

klaude_code/core/tool/tool_registry.py

@@ -3,8 +3,8 @@ from typing import TypeVar
 
 from klaude_code.core.tool.sub_agent_tool import SubAgentTool
 from klaude_code.core.tool.tool_abc import ToolABC
-from klaude_code.protocol import llm_param, tools
-from klaude_code.protocol.sub_agent import get_sub_agent_profile, iter_sub_agent_profiles, sub_agent_tool_names
+from klaude_code.protocol import llm_param
+from klaude_code.protocol.sub_agent import iter_sub_agent_profiles
 
 _REGISTRY: dict[str, type[ToolABC]] = {}
 
@@ -45,34 +45,3 @@ def get_tool_schemas(tool_names: list[str]) -> list[llm_param.ToolSchema]:
 def get_registry() -> dict[str, type[ToolABC]]:
     """Get the global tool registry."""
     return _REGISTRY
-
-
-def load_agent_tools(
-    model_name: str, sub_agent_type: tools.SubAgentType | None = None, *, vanilla: bool = False
-) -> list[llm_param.ToolSchema]:
-    """Get tools for an agent based on model and agent type.
-
-    Args:
-        model_name: The model name.
-        sub_agent_type: If None, returns main agent tools. Otherwise returns sub-agent tools.
-        vanilla: If True, returns minimal vanilla tools (ignores sub_agent_type).
-    """
-    if vanilla:
-        return get_tool_schemas([tools.BASH, tools.EDIT, tools.WRITE, tools.READ])
-
-    if sub_agent_type is not None:
-        profile = get_sub_agent_profile(sub_agent_type)
-        return get_tool_schemas(list(profile.tool_set))
-
-    # Main agent tools
-    if "gpt-5" in model_name:
-        tool_names = [tools.BASH, tools.READ, tools.APPLY_PATCH, tools.UPDATE_PLAN]
-    elif "gemini-3" in model_name:
-        tool_names = [tools.BASH, tools.READ, tools.EDIT, tools.WRITE]
-    else:
-        tool_names = [tools.BASH, tools.READ, tools.EDIT, tools.WRITE, tools.TODO_WRITE]
-
-    tool_names.extend(sub_agent_tool_names(enabled_only=True, model_name=model_name))
-    tool_names.extend([tools.SKILL, tools.MERMAID])
-    # tool_names.extend([tools.MEMORY])
-    return get_tool_schemas(tool_names)

klaude_code/core/tool/tool_runner.py

@@ -3,9 +3,9 @@ from collections.abc import AsyncGenerator, Callable, Iterable, Sequence
 from dataclasses import dataclass
 
 from klaude_code.const import CANCEL_OUTPUT
+from klaude_code.core.tool.context import ToolContext
 from klaude_code.core.tool.report_back_tool import ReportBackTool
 from klaude_code.core.tool.tool_abc import ToolABC, ToolConcurrencyPolicy
-from klaude_code.core.tool.tool_context import current_sub_agent_session_id_recorder
 from klaude_code.core.tool.truncation import truncate_tool_output
 from klaude_code.protocol import message, model, tools
 
@@ -18,19 +18,24 @@ class ToolCallRequest:
     arguments_json: str
 
 
-async def run_tool(tool_call: ToolCallRequest, registry: dict[str, type[ToolABC]]) -> message.ToolResultMessage:
+async def run_tool(
+    tool_call: ToolCallRequest,
+    registry: dict[str, type[ToolABC]],
+    context: ToolContext,
+) -> message.ToolResultMessage:
     """Execute a tool call and return the result.
 
     Args:
         tool_call: The tool call to execute.
         registry: The tool registry mapping tool names to tool classes.
+        context: The explicit tool execution context.
 
     Returns:
         The result of the tool execution.
     """
     # Special handling for report_back tool (not registered in global registry)
     if tool_call.tool_name == tools.REPORT_BACK:
-        tool_result = await ReportBackTool.call(tool_call.arguments_json)
+        tool_result = await ReportBackTool.call(tool_call.arguments_json, context)
         tool_result.call_id = tool_call.call_id
         tool_result.tool_name = tool_call.tool_name
         return tool_result
@@ -43,7 +48,7 @@ async def run_tool(tool_call: ToolCallRequest, registry: dict[str, type[ToolABC]
             tool_name=tool_call.tool_name,
         )
     try:
-        tool_result = await registry[tool_call.tool_name].call(tool_call.arguments_json)
+        tool_result = await registry[tool_call.tool_name].call(tool_call.arguments_json, context)
         tool_result.call_id = tool_call.call_id
         tool_result.tool_name = tool_call.tool_name
         if tool_result.output_text:
@@ -109,9 +114,11 @@ class ToolExecutor:
     def __init__(
         self,
         *,
+        context: ToolContext,
         registry: dict[str, type[ToolABC]],
         append_history: Callable[[Sequence[message.HistoryEvent]], None],
     ) -> None:
+        self._context = context
         self._registry = registry
         self._append_history = append_history
 
@@ -268,15 +275,11 @@ class ToolExecutor:
 
     async def _run_single_tool_call(self, tool_call: ToolCallRequest) -> list[ToolExecutorEvent]:
         def _record_sub_agent_session_id(session_id: str) -> None:
-            # Keep the first recorded id if multiple writes happen.
             if tool_call.call_id not in self._sub_agent_session_ids:
                 self._sub_agent_session_ids[tool_call.call_id] = session_id
 
-        recorder_token = current_sub_agent_session_id_recorder.set(_record_sub_agent_session_id)
-        try:
-            tool_result: message.ToolResultMessage = await run_tool(tool_call, self._registry)
-        finally:
-            current_sub_agent_session_id_recorder.reset(recorder_token)
+        call_context = self._context.with_record_sub_agent_session_id(_record_sub_agent_session_id)
+        tool_result: message.ToolResultMessage = await run_tool(tool_call, self._registry, call_context)
 
         self._append_history([tool_result])
 

klaude_code/core/tool/web/mermaid_tool.py

@@ -8,6 +8,7 @@ from pathlib import Path
 from pydantic import BaseModel, Field
 
 from klaude_code.const import MERMAID_LIVE_PREFIX
+from klaude_code.core.tool.context import ToolContext
 from klaude_code.core.tool.tool_abc import ToolABC, load_desc
 from klaude_code.core.tool.tool_registry import register
 from klaude_code.protocol import llm_param, message, model, tools
@@ -40,7 +41,8 @@ class MermaidTool(ToolABC):
         )
 
     @classmethod
-    async def call(cls, arguments: str) -> message.ToolResultMessage:
+    async def call(cls, arguments: str, context: ToolContext) -> message.ToolResultMessage:
+        del context
         try:
             args = cls.MermaidArguments.model_validate_json(arguments)
         except Exception as exc:  # pragma: no cover - defensive