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

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)

klaude_code/core/tool/file/edit_tool.py

@@ -8,10 +8,11 @@ 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
 
@@ -85,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
@@ -110,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
@@ -118,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:
@@ -136,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,
@@ -191,28 +190,27 @@ class EditTool(ToolABC):
         except (OSError, UnicodeError) as e:  # pragma: no cover
             return message.ToolResultMessage(status="error", output_text=f"<tool_use_error>{e}</tool_use_error>")
 
-        # Prepare UI extra: unified diff with 3 context lines
+        # Prepare UI extra: unified diff with default context lines
         diff_lines = list(
             difflib.unified_diff(
                 before.splitlines(),
                 after.splitlines(),
                 fromfile=file_path,
                 tofile=file_path,
-                n=3,
+                n=DIFF_DEFAULT_CONTEXT_LINES,
             )
         )
         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()