klaude-code 2.4.1__py3-none-any.whl → 2.5.0__py3-none-any.whl

klaude_code/app/runtime.py

@@ -11,7 +11,6 @@ from klaude_code.config import Config, load_config
 from klaude_code.core.agent import Agent
 from klaude_code.core.agent_profile import (
     DefaultModelProfileProvider,
-    NanoBananaModelProfileProvider,
     VanillaModelProfileProvider,
 )
 from klaude_code.core.executor import Executor
@@ -28,7 +27,6 @@ class AppInitConfig:
     model: str | None
     debug: bool
     vanilla: bool
-    banana: bool
     debug_filters: set[DebugType] | None = None
 
 
@@ -59,7 +57,7 @@ async def initialize_app_components(
         llm_clients = build_llm_clients(
             config,
             model_override=init_config.model,
-            skip_sub_agents=init_config.vanilla or init_config.banana,
+            skip_sub_agents=init_config.vanilla,
         )
     except ValueError as exc:
         if init_config.model:
@@ -74,9 +72,7 @@ async def initialize_app_components(
             log((f"Error: failed to load the default model configuration: {exc}", "red"))
         raise typer.Exit(2) from None
 
-    if init_config.banana:
-        model_profile_provider = NanoBananaModelProfileProvider()
-    elif init_config.vanilla:
+    if init_config.vanilla:
         model_profile_provider = VanillaModelProfileProvider()
     else:
         model_profile_provider = DefaultModelProfileProvider(config=config)

klaude_code/cli/main.py

@@ -222,7 +222,6 @@ def main_callback(
             model=chosen_model,
             debug=debug_enabled,
             vanilla=vanilla,
-            banana=banana,
             debug_filters=debug_filters,
         )
 

klaude_code/config/assets/builtin_config.yaml

@@ -152,6 +152,13 @@ provider_list:
     modalities:
     - image
     - text
+  - model_name: flux
+    model_id: black-forest-labs/flux.2-max
+    context_limit: 47000
+    cost: {input: 7.32, output: 7.32, image: 7.32}
+    modalities:
+    - image
+    - text
 - provider_name: google
   protocol: google
   api_key: ${GOOGLE_API_KEY}

klaude_code/const.py

@@ -47,7 +47,6 @@ THROUGHPUT_MIN_DURATION_SEC = 0.15  # Minimum duration (seconds) for throughput
 INITIAL_RETRY_DELAY_S = 1.0  # Initial delay before retrying a failed turn (seconds)
 MAX_RETRY_DELAY_S = 30.0  # Maximum delay between retries (seconds)
 CANCEL_OUTPUT = "[Request interrupted by user for tool use]"  # Message shown when tool call is cancelled
-INTERRUPT_MARKER = " <system>interrupted by user</system>"  # Marker appended when assistant is interrupted
 EMPTY_TOOL_OUTPUT_MESSAGE = (
     "<system-reminder>Tool ran without output or errors</system-reminder>"  # Tool output placeholder
 )
@@ -112,7 +111,10 @@ DIFF_DEFAULT_CONTEXT_LINES = 3  # Default number of context lines in diff output
 TOOL_OUTPUT_MAX_LENGTH = 40000  # Maximum length for tool output before truncation
 TOOL_OUTPUT_DISPLAY_HEAD = 10000  # Characters to show from the beginning of truncated output
 TOOL_OUTPUT_DISPLAY_TAIL = 10000  # Characters to show from the end of truncated output
-TOOL_OUTPUT_TRUNCATION_DIR = "/tmp/klaude"  # Directory for saving full truncated output
+TOOL_OUTPUT_MAX_LINES = 2000  # Maximum lines for tool output before truncation
+TOOL_OUTPUT_DISPLAY_HEAD_LINES = 1000  # Lines to show from the beginning of truncated output
+TOOL_OUTPUT_DISPLAY_TAIL_LINES = 1000  # Lines to show from the end of truncated output
+TOOL_OUTPUT_TRUNCATION_DIR = "/tmp"  # Directory for saving full truncated output
 
 
 # =============================================================================
@@ -156,8 +158,8 @@ STATUS_HINT_TEXT = " (esc to interrupt)"  # Status hint text shown after spinner
 
 # Spinner status texts
 STATUS_WAITING_TEXT = "Connecting …"
-STATUS_THINKING_TEXT = "Reasoning …"
-STATUS_COMPOSING_TEXT = "Generating"
+STATUS_THINKING_TEXT = "Thinking …"
+STATUS_COMPOSING_TEXT = "Composing …"
 
 # Backwards-compatible alias for the default spinner status text.
 STATUS_DEFAULT_TEXT = STATUS_WAITING_TEXT
@@ -166,6 +168,7 @@ SPINNER_BREATH_PERIOD_SECONDS: float = 2.0  # Spinner breathing animation period
 STATUS_SHIMMER_PADDING = 10  # Horizontal padding for shimmer band position
 STATUS_SHIMMER_BAND_HALF_WIDTH = 5.0  # Half-width of shimmer band in characters
 STATUS_SHIMMER_ALPHA_SCALE = 0.7  # Scale factor for shimmer intensity
+STATUS_SHOW_BUFFER_LENGTH = False  # Show character count (e.g., "(213)") during text generation
 
 
 # =============================================================================

klaude_code/core/agent.py

@@ -8,7 +8,7 @@ from klaude_code.core.tool import build_todo_context, get_registry
 from klaude_code.core.tool.context import RunSubtask
 from klaude_code.llm import LLMClientABC
 from klaude_code.log import DebugType, log_debug
-from klaude_code.protocol import events
+from klaude_code.protocol import events, model
 from klaude_code.protocol.message import UserInputPayload
 from klaude_code.session import Session
 
@@ -93,3 +93,12 @@ class Agent:
 
     def get_llm_client(self) -> LLMClientABC:
         return self.profile.llm_client
+
+    def get_partial_metadata(self) -> model.TaskMetadata | None:
+        """Get partial metadata from the currently running task.
+
+        Returns None if no task is running or no usage data has been accumulated.
+        """
+        if self._current_task is None:
+            return None
+        return self._current_task.get_partial_metadata()

klaude_code/core/agent_profile.py

@@ -62,10 +62,7 @@ PROMPT_FILES: dict[str, str] = {
 }
 
 
-NANO_BANANA_SYSTEM_PROMPT_PATH = "prompts/prompt-nano-banana.md"
-
-
-STRUCTURED_OUTPUT_PROMPT = """\
+STRUCTURED_OUTPUT_PROMPT_FOR_SUB_AGENT = """\
 
 # Structured Output
 You have a `report_back` tool available. When you complete the task,\
@@ -74,6 +71,20 @@ Only the content passed to `report_back` will be returned to user.\
 """
 
 
+SUB_AGENT_COMMON_PROMPT_FOR_MAIN_AGENT = """\
+
+# Sub-agent capabilities
+You have sub-agents (e.g. Task, Explore, WebAgent, ImageGen) with structured output and resume capabilites:
+- Agents can be provided with an `output_format` (JSON Schema) parameter for structured output
+-   Example: `output_format={"type": "object", "properties": {"files": {"type": "array", "items": {"type": "string"}, "description": "List of file paths that match the search criteria, e.g. ['src/main.py', 'src/utils/helper.py']"}}, "required": ["files"]}`
+- Agents can be resumed using the `resume` parameter by passing the agent ID from a previous invocation. \
+When resumed, the agent continues with its full previous context preserved. \
+When NOT resuming, each invocation starts fresh and you should provide a detailed task description with all necessary context.
+- When the agent is done, it will return a single message back to you along with its agent ID. \
+You can use this ID to resume the agent later if needed for follow-up work.
+"""
+
+
 @cache
 def _load_prompt_by_path(prompt_path: str) -> str:
     """Load and cache prompt content from a file path relative to core package."""
@@ -142,10 +153,24 @@ def _build_env_info(model_name: str) -> str:
     return "\n".join(env_lines)
 
 
+def _has_sub_agents(config: Config | None) -> bool:
+    """Check if there are any sub-agent tools available for the main agent."""
+    if config is not None:
+        from klaude_code.config.sub_agent_model_helper import SubAgentModelHelper
+
+        helper = SubAgentModelHelper(config)
+        return bool(helper.get_enabled_sub_agent_tool_names())
+
+    from klaude_code.protocol.sub_agent import sub_agent_tool_names
+
+    return bool(sub_agent_tool_names(enabled_only=True))
+
+
 def load_system_prompt(
     model_name: str,
     protocol: llm_param.LLMClientProtocol,
     sub_agent_type: str | None = None,
+    config: Config | None = None,
 ) -> str:
     """Get system prompt content for the given model and sub-agent type."""
 
@@ -161,13 +186,18 @@ def load_system_prompt(
         return base_prompt
 
     skills_prompt = ""
+    sub_agent_prompt = ""
     if sub_agent_type is None:
         # Skills are progressive-disclosure: keep only metadata in the system prompt.
         from klaude_code.skill.manager import format_available_skills_for_system_prompt
 
         skills_prompt = format_available_skills_for_system_prompt()
 
-    return base_prompt + _build_env_info(model_name) + skills_prompt
+        # Add sub-agent resume instructions if there are sub-agent tools available.
+        if _has_sub_agents(config):
+            sub_agent_prompt = "\n" + SUB_AGENT_COMMON_PROMPT_FOR_MAIN_AGENT
+
+    return base_prompt + _build_env_info(model_name) + skills_prompt + sub_agent_prompt
 
 
 def load_agent_tools(
@@ -245,7 +275,7 @@ def with_structured_output(profile: AgentProfile, output_schema: dict[str, Any])
     base_prompt = profile.system_prompt or ""
     return AgentProfile(
         llm_client=profile.llm_client,
-        system_prompt=base_prompt + STRUCTURED_OUTPUT_PROMPT,
+        system_prompt=base_prompt + STRUCTURED_OUTPUT_PROMPT_FOR_SUB_AGENT,
         tools=[*profile.tools, report_back_tool_class.schema()],
         reminders=profile.reminders,
     )
@@ -279,17 +309,24 @@ class DefaultModelProfileProvider(ModelProfileProvider):
         model_name = llm_client.model_name
         llm_config = llm_client.get_llm_config()
 
-        # Image generation models should not have tools
-        if llm_config.modalities and "image" in llm_config.modalities:
+        # Image generation models should not have system prompt, tools, or reminders
+        is_image_model = llm_config.modalities and "image" in llm_config.modalities
+        if is_image_model:
+            agent_system_prompt: str | None = None
             agent_tools: list[llm_param.ToolSchema] = []
+            agent_reminders: list[Reminder] = []
         else:
+            agent_system_prompt = load_system_prompt(
+                model_name, llm_client.protocol, sub_agent_type, config=self._config
+            )
             agent_tools = load_agent_tools(model_name, sub_agent_type, config=self._config)
+            agent_reminders = load_agent_reminders(model_name, sub_agent_type)
 
         profile = AgentProfile(
             llm_client=llm_client,
-            system_prompt=load_system_prompt(model_name, llm_client.protocol, sub_agent_type),
+            system_prompt=agent_system_prompt,
             tools=agent_tools,
-            reminders=load_agent_reminders(model_name, sub_agent_type),
+            reminders=agent_reminders,
         )
         if output_schema:
             return with_structured_output(profile, output_schema)
@@ -316,28 +353,3 @@ class VanillaModelProfileProvider(ModelProfileProvider):
         if output_schema:
             return with_structured_output(profile, output_schema)
         return profile
-
-
-class NanoBananaModelProfileProvider(ModelProfileProvider):
-    """Provider for the Nano Banana image generation model.
-
-    This mode uses a dedicated system prompt and strips all tools/reminders.
-    """
-
-    def build_profile(
-        self,
-        llm_client: LLMClientABC,
-        sub_agent_type: tools.SubAgentType | None = None,
-        *,
-        output_schema: dict[str, Any] | None = None,
-    ) -> AgentProfile:
-        del sub_agent_type
-        profile = AgentProfile(
-            llm_client=llm_client,
-            system_prompt=_load_prompt_by_path(NANO_BANANA_SYSTEM_PROMPT_PATH),
-            tools=[],
-            reminders=[],
-        )
-        if output_schema:
-            return with_structured_output(profile, output_schema)
-        return profile

klaude_code/core/executor.py

@@ -17,14 +17,14 @@ from pathlib import Path
 from klaude_code.config import load_config
 from klaude_code.config.sub_agent_model_helper import SubAgentModelHelper
 from klaude_code.core.agent import Agent
-from klaude_code.core.agent_profile import AgentProfile, DefaultModelProfileProvider, ModelProfileProvider
+from klaude_code.core.agent_profile import DefaultModelProfileProvider, ModelProfileProvider
 from klaude_code.core.manager import LLMClients, SubAgentManager
 from klaude_code.llm.registry import create_llm_client
 from klaude_code.log import DebugType, log_debug
 from klaude_code.protocol import commands, events, message, model, op
 from klaude_code.protocol.llm_param import LLMConfigParameter, Thinking
 from klaude_code.protocol.op_handler import OperationHandler
-from klaude_code.protocol.sub_agent import SubAgentResult, get_sub_agent_profile_by_tool
+from klaude_code.protocol.sub_agent import SubAgentResult
 from klaude_code.session.export import build_export_html, get_default_export_path
 from klaude_code.session.session import Session
 
@@ -110,19 +110,6 @@ class AgentRuntime:
     def current_agent(self) -> Agent | None:
         return self._agent
 
-    def _get_sub_agent_models(self, profile: AgentProfile) -> dict[str, LLMConfigParameter]:
-        """Build a dict of sub-agent type to LLMConfigParameter based on profile tools."""
-        enabled_types: set[str] = set()
-        for tool in profile.tools:
-            sub_profile = get_sub_agent_profile_by_tool(tool.name)
-            if sub_profile is not None:
-                enabled_types.add(sub_profile.name)
-        return {
-            sub_agent_type: client.get_llm_config()
-            for sub_agent_type, client in self._llm_clients.sub_clients.items()
-            if sub_agent_type in enabled_types
-        }
-
     async def ensure_agent(self, session_id: str | None = None) -> Agent:
         """Return the active agent, creating or loading a session as needed."""
 
@@ -149,7 +136,6 @@ class AgentRuntime:
                 session_id=session.id,
                 work_dir=str(session.work_dir),
                 llm_config=self._llm_clients.main.get_llm_config(),
-                sub_agent_models=self._get_sub_agent_models(profile),
             )
         )
 
@@ -206,7 +192,6 @@ class AgentRuntime:
                 session_id=agent.session.id,
                 work_dir=str(agent.session.work_dir),
                 llm_config=self._llm_clients.main.get_llm_config(),
-                sub_agent_models=self._get_sub_agent_models(agent.profile),
             )
         )
 
@@ -230,7 +215,6 @@ class AgentRuntime:
                 session_id=target_session.id,
                 work_dir=str(target_session.work_dir),
                 llm_config=self._llm_clients.main.get_llm_config(),
-                sub_agent_models=self._get_sub_agent_models(profile),
             )
         )
 
@@ -291,8 +275,11 @@ class AgentRuntime:
             async def _runner(
                 state: model.SubAgentState,
                 record_session_id: Callable[[str], None] | None,
+                register_metadata_getter: Callable[[Callable[[], model.TaskMetadata | None]], None] | None,
             ) -> SubAgentResult:
-                return await self._sub_agent_manager.run_sub_agent(agent, state, record_session_id=record_session_id)
+                return await self._sub_agent_manager.run_sub_agent(
+                    agent, state, record_session_id=record_session_id, register_metadata_getter=register_metadata_getter
+                )
 
             async for event in agent.run_task(user_input, run_subtask=_runner):
                 await self._emit_event(event)
@@ -464,7 +451,6 @@ class ExecutorContext:
                     llm_config=llm_config,
                     work_dir=str(agent.session.work_dir),
                     show_klaude_code_info=False,
-                    show_sub_agent_models=False,
                 )
             )
 
@@ -512,7 +498,6 @@ class ExecutorContext:
                     work_dir=str(agent.session.work_dir),
                     llm_config=agent.profile.llm_client.get_llm_config(),
                     show_klaude_code_info=False,
-                    show_sub_agent_models=False,
                 )
             )
 

klaude_code/core/manager/sub_agent_manager.py

@@ -44,6 +44,7 @@ class SubAgentManager:
         state: model.SubAgentState,
         *,
         record_session_id: Callable[[str], None] | None = None,
+        register_metadata_getter: Callable[[Callable[[], model.TaskMetadata | None]], None] | None = None,
     ) -> SubAgentResult:
         """Run a nested sub-agent task and return its result."""
 
@@ -114,6 +115,16 @@ class SubAgentManager:
             debug_type=DebugType.EXECUTION,
         )
 
+        # Register metadata getter so parent can retrieve partial metadata on cancel
+        def _get_partial_metadata() -> model.TaskMetadata | None:
+            metadata = child_agent.get_partial_metadata()
+            if metadata is not None:
+                metadata.description = state.sub_agent_desc or None
+            return metadata
+
+        if register_metadata_getter is not None:
+            register_metadata_getter(_get_partial_metadata)
+
         try:
             # Not emit the subtask's user input since task tool call is already rendered
             result: str = ""
@@ -138,6 +149,7 @@ class SubAgentManager:
                 # Capture TaskMetadataEvent for metadata propagation
                 elif isinstance(event, events.TaskMetadataEvent):
                     task_metadata = event.metadata.main_agent
+                    task_metadata.description = state.sub_agent_desc or None
                 await self.emit_event(event)
 
             # Ensure the sub-agent session is persisted before returning its id for resume.
@@ -148,7 +160,11 @@ class SubAgentManager:
                 task_metadata=task_metadata,
             )
         except asyncio.CancelledError:
-            # Propagate cancellation so tooling can treat it as user interrupt
+            # Call cancel() on child agent to emit cleanup events
+            # Note: Parent retrieves partial metadata via registered getter before this runs
+            for evt in child_agent.cancel():
+                await self.emit_event(evt)
+
             log_debug(
                 f"Sub-agent task for {state.sub_agent_type} was cancelled",
                 style="yellow",

klaude_code/core/prompts/prompt-sub-agent-web.md

@@ -17,7 +17,7 @@ You are a web research subagent that searches and fetches web content to provide
 - HTML pages are automatically converted to Markdown
 - JSON responses are auto-formatted with indentation
 - Other text content returned as-is
-- **Content is always saved to a local file** - check `<file_saved>` tag for the path
+- **Content is always saved to a local file** - path shown in `[Web content saved to ...]` at output start
 
 ## Tool Usage Strategy
 
@@ -54,9 +54,9 @@ Balance efficiency with thoroughness. For open-ended questions (e.g., "recommend
 ## Response Guidelines
 
 - Only your last message is returned to the main agent
+- Include the file path from `[Web content saved to ...]` so the main agent can access full content
 - **DO NOT copy full web page content** - the main agent can read the saved files directly
 - Provide a concise summary/analysis of key findings
-- Include the file path from `<file_saved>` so the main agent can access full content if needed
 - Lead with the most recent info for evolving topics
 - Favor original sources (company blogs, papers, gov sites) over aggregators
 - When sources conflict, explain the discrepancy and which source is more authoritative
@@ -73,5 +73,5 @@ Stop only when all are true:
 You MUST end every response with a "Sources:" section listing all URLs with their saved file paths:
 
 Sources:
-- [Source Title](https://example.com) -> /tmp/klaude/web/example_com-123456.md
-- [Another Source](https://example.com/page) -> /tmp/klaude/web/example_com_page-123456.md
+- [Source Title](https://example.com) -> /tmp/klaude-webfetch-example_com.txt
+- [Another Source](https://example.com/page) -> /tmp/klaude-webfetch-example_com_page.txt

klaude_code/core/task.py

@@ -66,6 +66,49 @@ class MetadataAccumulator:
         """Add sub-agent task metadata to the accumulated state."""
         self._sub_agent_metadata.append(sub_agent_metadata)
 
+    def get_partial(self, task_duration_s: float) -> model.TaskMetadata | None:
+        """Return a snapshot of main agent metadata without modifying accumulator state.
+
+        Returns None if no usage data has been accumulated yet.
+        """
+        if self._main_agent.usage is None:
+            return None
+
+        # Create a copy to avoid modifying the original
+        usage_copy = self._main_agent.usage.model_copy(deep=True)
+
+        if self._throughput_tracked_tokens > 0:
+            usage_copy.throughput_tps = self._throughput_weighted_sum / self._throughput_tracked_tokens
+        else:
+            usage_copy.throughput_tps = None
+
+        if self._first_token_latency_count > 0:
+            usage_copy.first_token_latency_ms = self._first_token_latency_sum / self._first_token_latency_count
+        else:
+            usage_copy.first_token_latency_ms = None
+
+        return model.TaskMetadata(
+            model_name=self._main_agent.model_name,
+            provider=self._main_agent.provider,
+            usage=usage_copy,
+            task_duration_s=task_duration_s,
+            turn_count=self._turn_count,
+        )
+
+    def get_partial_item(self, task_duration_s: float) -> model.TaskMetadataItem | None:
+        """Return a snapshot of full metadata (main + sub-agents) without modifying state.
+
+        Returns None if no usage data has been accumulated yet.
+        """
+        main_agent = self.get_partial(task_duration_s)
+        if main_agent is None:
+            return None
+
+        return model.TaskMetadataItem(
+            main_agent=main_agent,
+            sub_agent_task_metadata=list(self._sub_agent_metadata),
+        )
+
     def finalize(self, task_duration_s: float) -> model.TaskMetadataItem:
         """Return the final accumulated metadata with computed throughput and duration."""
         if self._main_agent.usage is not None:
@@ -129,20 +172,38 @@ class TaskExecutor:
     def current_turn(self) -> TurnExecutor | None:
         return self._current_turn
 
+    def get_partial_metadata(self) -> model.TaskMetadata | None:
+        """Get the currently accumulated metadata without finalizing.
+
+        Returns partial metadata that can be used if the task is interrupted.
+        """
+        if self._metadata_accumulator is None or self._started_at <= 0:
+            return None
+        task_duration_s = time.perf_counter() - self._started_at
+        return self._metadata_accumulator.get_partial(task_duration_s)
+
     def cancel(self) -> list[events.Event]:
         """Cancel the current turn and return any resulting events including metadata."""
         ui_events: list[events.Event] = []
         if self._current_turn is not None:
-            ui_events.extend(self._current_turn.cancel())
+            for evt in self._current_turn.cancel():
+                # Collect sub-agent task metadata from cancelled tool results
+                if (
+                    isinstance(evt, events.ToolResultEvent)
+                    and evt.task_metadata is not None
+                    and self._metadata_accumulator is not None
+                ):
+                    self._metadata_accumulator.add_sub_agent_metadata(evt.task_metadata)
+                ui_events.append(evt)
             self._current_turn = None
 
         # Emit partial metadata on cancellation
         if self._metadata_accumulator is not None and self._started_at > 0:
             task_duration_s = time.perf_counter() - self._started_at
-            accumulated = self._metadata_accumulator.finalize(task_duration_s)
-            if accumulated.main_agent.usage is not None:
+            accumulated = self._metadata_accumulator.get_partial_item(task_duration_s)
+            if accumulated is not None:
                 session_id = self._context.session_ctx.session_id
-                ui_events.append(events.TaskMetadataEvent(metadata=accumulated, session_id=session_id))
+                ui_events.append(events.TaskMetadataEvent(metadata=accumulated, session_id=session_id, cancelled=True))
                 self._context.session_ctx.append_history([accumulated])
 
         return ui_events

klaude_code/core/tool/__init__.py

@@ -13,7 +13,6 @@ from .todo.update_plan_tool import UpdatePlanTool
 from .tool_abc import ToolABC
 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
 from .web.web_fetch_tool import WebFetchTool
 from .web.web_search_tool import WebSearchTool
@@ -29,14 +28,12 @@ __all__ = [
     "ReportBackTool",
     "RunSubtask",
     "SafetyCheckResult",
-    "SimpleTruncationStrategy",
     "SubAgentResumeClaims",
     "SubAgentTool",
     "TodoContext",
     "TodoWriteTool",
     "ToolABC",
     "ToolContext",
-    "TruncationStrategy",
     "UpdatePlanTool",
     "WebFetchTool",
     "WebSearchTool",
@@ -44,9 +41,7 @@ __all__ = [
     "build_todo_context",
     "get_registry",
     "get_tool_schemas",
-    "get_truncation_strategy",
     "is_safe_command",
     "process_patch",
     "run_tool",
-    "set_truncation_strategy",
 ]

klaude_code/core/tool/context.py

@@ -10,7 +10,12 @@ from klaude_code.session.session import Session
 
 type FileTracker = MutableMapping[str, model.FileStatus]
 
-RunSubtask = Callable[[model.SubAgentState, Callable[[str], None] | None], Awaitable[SubAgentResult]]
+GetMetadataFn = Callable[[], model.TaskMetadata | None]
+
+RunSubtask = Callable[
+    [model.SubAgentState, Callable[[str], None] | None, Callable[[GetMetadataFn], None] | None],
+    Awaitable[SubAgentResult],
+]
 
 
 @dataclass
@@ -79,6 +84,12 @@ class ToolContext:
     run_subtask: RunSubtask | None = None
     sub_agent_resume_claims: SubAgentResumeClaims | None = None
     record_sub_agent_session_id: Callable[[str], None] | None = None
+    register_sub_agent_metadata_getter: Callable[[GetMetadataFn], 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)
+
+    def with_register_sub_agent_metadata_getter(
+        self, callback: Callable[[GetMetadataFn], None] | None
+    ) -> ToolContext:
+        return replace(self, register_sub_agent_metadata_getter=callback)