klaude-code 2.8.1__py3-none-any.whl → 2.9.0__py3-none-any.whl

klaude_code/core/tool/__init__.py

@@ -7,7 +7,7 @@ from .file.write_tool import WriteTool
 from .report_back_tool import ReportBackTool
 from .shell.bash_tool import BashTool
 from .shell.command_safety import SafetyCheckResult, is_safe_command
-from .sub_agent_tool import SubAgentTool
+from .sub_agent import ImageGenTool, TaskTool
 from .todo.todo_write_tool import TodoWriteTool
 from .todo.update_plan_tool import UpdatePlanTool
 from .tool_abc import ToolABC
@@ -23,13 +23,14 @@ __all__ = [
     "DiffError",
     "EditTool",
     "FileTracker",
+    "ImageGenTool",
     "MermaidTool",
     "ReadTool",
     "ReportBackTool",
     "RunSubtask",
     "SafetyCheckResult",
     "SubAgentResumeClaims",
-    "SubAgentTool",
+    "TaskTool",
     "TodoContext",
     "TodoWriteTool",
     "ToolABC",

klaude_code/core/tool/file/apply_patch.py

@@ -26,33 +26,6 @@ class Commit(BaseModel):
     changes: dict[str, FileChange] = Field(default_factory=dict)
 
 
-def assemble_changes(orig: dict[str, str | None], dest: dict[str, str | None]) -> Commit:
-    commit = Commit()
-    for path in sorted(set(orig.keys()).union(dest.keys())):
-        old_content = orig.get(path)
-        new_content = dest.get(path)
-        if old_content != new_content:
-            if old_content is not None and new_content is not None:
-                commit.changes[path] = FileChange(
-                    type=ActionType.UPDATE,
-                    old_content=old_content,
-                    new_content=new_content,
-                )
-            elif new_content:
-                commit.changes[path] = FileChange(
-                    type=ActionType.ADD,
-                    new_content=new_content,
-                )
-            elif old_content:
-                commit.changes[path] = FileChange(
-                    type=ActionType.DELETE,
-                    old_content=old_content,
-                )
-            else:
-                raise AssertionError()
-    return commit
-
-
 def _new_str_list() -> list[str]:
     # Returns a new list[str] for pydantic Field default_factory
     return []

klaude_code/core/tool/file/read_tool.md

@@ -4,10 +4,11 @@ When you need to read an image, use this tool.
 
 Usage:
 - The file_path parameter must be an absolute path, not a relative path
-- By default, it reads up to 2000 lines starting from the beginning of the file
+- By default, it reads up to ${line_cap} lines starting from the beginning of the file
 - This tool allows you to read images (eg PNG, JPG, etc). When reading an image file the contents are presented visually as you are a multimodal LLM.
 - You can optionally specify a line offset and limit (especially handy for long files), but it's recommended to read the whole file by not providing these parameters
-- Any lines longer than 2000 characters will be truncated
+- Any lines longer than ${char_limit_per_line} characters will be truncated
+- Total output is capped at ${max_chars} characters
 - Results are returned using cat -n format, with line numbers starting at 1
 - This tool can only read files, not directories. To read a directory, use an ls command via the Bash tool.
 - You have the capability to call multiple tools in a single response. It is always better to speculatively read multiple files as a batch that are potentially useful. 

klaude_code/core/tool/file/read_tool.py

@@ -22,6 +22,7 @@ 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_registry import register
 from klaude_code.protocol import llm_param, message, model, tools
+from klaude_code.protocol.model import ImageUIExtra
 
 _IMAGE_MIME_TYPES: dict[str, str] = {
     ".png": "image/png",
@@ -164,7 +165,14 @@ class ReadTool(ToolABC):
         return llm_param.ToolSchema(
             name=tools.READ,
             type="function",
-            description=load_desc(Path(__file__).parent / "read_tool.md"),
+            description=load_desc(
+                Path(__file__).parent / "read_tool.md",
+                {
+                    "line_cap": str(READ_GLOBAL_LINE_CAP),
+                    "char_limit_per_line": str(READ_CHAR_LIMIT_PER_LINE),
+                    "max_chars": str(READ_MAX_CHARS),
+                },
+            ),
             parameters={
                 "type": "object",
                 "properties": {
@@ -280,7 +288,12 @@ class ReadTool(ToolABC):
             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)
-            return message.ToolResultMessage(status="success", output_text=output_text, parts=[image_part])
+            return message.ToolResultMessage(
+                status="success",
+                output_text=output_text,
+                parts=[image_part],
+                ui_extra=ImageUIExtra(file_path=file_path),
+            )
 
         offset = 1 if args.offset is None or args.offset < 1 else int(args.offset)
         limit = None if args.limit is None else int(args.limit)

klaude_code/core/tool/offload.py

@@ -68,13 +68,6 @@ class OffloadPolicy(Enum):
     ON_THRESHOLD = auto()  # Offload only when exceeding size threshold
 
 
-class TruncationStyle(Enum):
-    """How to truncate content that exceeds limits."""
-
-    HEAD_ONLY = auto()  # Keep head, discard tail (important content at top)
-    HEAD_TAIL = auto()  # Keep head and tail, discard middle (errors at end)
-
-
 @dataclass
 class OffloadResult:
     """Result of offload/truncation operation."""
@@ -94,18 +87,6 @@ class OffloadResult:
 class OffloadStrategy(ABC):
     """Base class for tool-specific offload strategies."""
 
-    @property
-    @abstractmethod
-    def offload_policy(self) -> OffloadPolicy:
-        """When to offload content to file."""
-        ...
-
-    @property
-    @abstractmethod
-    def truncation_style(self) -> TruncationStyle:
-        """How to truncate content."""
-        ...
-
     @abstractmethod
     def process(self, output: str, tool_call: ToolCallLike | None = None) -> OffloadResult:
         """Process tool output: truncate and optionally offload."""
@@ -126,14 +107,6 @@ class ReadToolStrategy(OffloadStrategy):
     This strategy is a pass-through since Read tool handles its own truncation.
     """
 
-    @property
-    def offload_policy(self) -> OffloadPolicy:
-        return OffloadPolicy.NEVER
-
-    @property
-    def truncation_style(self) -> TruncationStyle:
-        return TruncationStyle.HEAD_ONLY
-
     def process(self, output: str, tool_call: ToolCallLike | None = None) -> OffloadResult:
         return OffloadResult(output=output, was_truncated=False, original_length=len(output))
 
@@ -165,14 +138,6 @@ class HeadTailOffloadStrategy(OffloadStrategy):
         self.offload_dir = Path(offload_dir or TOOL_OUTPUT_TRUNCATION_DIR)
         self._policy = policy
 
-    @property
-    def offload_policy(self) -> OffloadPolicy:
-        return self._policy
-
-    @property
-    def truncation_style(self) -> TruncationStyle:
-        return TruncationStyle.HEAD_TAIL
-
     def _save_to_file(self, output: str, tool_call: ToolCallLike | None) -> str | None:
         """Save full output to file. Returns path or None on failure."""
         try:

klaude_code/core/tool/sub_agent/__init__.py

@@ -0,0 +1,6 @@
+"""Sub-agent tool implementations."""
+
+from .image_gen import ImageGenTool
+from .task import TaskTool
+
+__all__ = ["ImageGenTool", "TaskTool"]

klaude_code/core/tool/sub_agent/image_gen.md

@@ -0,0 +1,16 @@
+Generate one or more images from a text prompt.
+
+This tool invokes an Image Gen model to generate images. The generated image paths are automatically returned in the response.
+
+Inputs:
+- `prompt`: The main instruction describing the desired image.
+- `image_paths` (optional): Local image file paths to use as references for editing or style guidance.
+- `generation` (optional): Per-call image generation settings (aspect ratio / size).
+
+Notes:
+- Provide a short textual description of the generated image(s).
+- Do NOT include base64 image data in text output.
+- When providing multiple input images, describe each image's characteristics and purpose in the prompt, not just "image 1, image 2".
+
+Multi-turn image editing:
+- Use `resume` to continue editing a previously generated image. The agent preserves its full context including the generated image, so you don't need to pass `image_paths` again.

klaude_code/core/tool/sub_agent/image_gen.py

@@ -0,0 +1,146 @@
+"""Image generation tool implementation."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, cast
+
+from klaude_code.core.tool.context import ToolContext
+from klaude_code.core.tool.tool_abc import ToolABC, ToolConcurrencyPolicy, ToolMetadata, load_desc
+from klaude_code.core.tool.tool_registry import register
+from klaude_code.protocol import llm_param, message, model, tools
+from klaude_code.protocol.sub_agent import get_sub_agent_profile
+from klaude_code.protocol.sub_agent.image_gen import build_image_gen_prompt
+from klaude_code.session.session import Session
+
+IMAGE_GEN_PARAMETERS: dict[str, Any] = {
+    "type": "object",
+    "properties": {
+        "resume": {
+            "type": "string",
+            "description": "Optional agent ID to resume from. If provided, the agent will continue from the previous execution transcript.",
+        },
+        "description": {
+            "type": "string",
+            "description": "A short (3-5 word) description of the request.",
+        },
+        "prompt": {
+            "type": "string",
+            "description": "Text prompt for image generation.",
+        },
+        "image_paths": {
+            "type": "array",
+            "items": {"type": "string"},
+            "description": "Optional local image file paths used as references.",
+        },
+        "generation": {
+            "type": "object",
+            "description": "Optional per-call image generation settings.",
+            "properties": {
+                "aspect_ratio": {
+                    "type": "string",
+                    "description": "Aspect ratio, e.g. '16:9', '1:1', '9:16'.",
+                },
+                "image_size": {
+                    "type": "string",
+                    "enum": ["1K", "2K", "4K"],
+                    "description": "Output size for Nano Banana Pro (must use uppercase K).",
+                },
+            },
+            "additionalProperties": False,
+        },
+    },
+    "required": ["prompt"],
+    "additionalProperties": False,
+}
+
+
+@register(tools.IMAGE_GEN)
+class ImageGenTool(ToolABC):
+    """Generate or edit images using the ImageGen sub-agent."""
+
+    @classmethod
+    def metadata(cls) -> ToolMetadata:
+        return ToolMetadata(concurrency_policy=ToolConcurrencyPolicy.CONCURRENT, has_side_effects=True)
+
+    @classmethod
+    def schema(cls) -> llm_param.ToolSchema:
+        return llm_param.ToolSchema(
+            name=tools.IMAGE_GEN,
+            type="function",
+            description=load_desc(Path(__file__).parent / "image_gen.md"),
+            parameters=IMAGE_GEN_PARAMETERS,
+        )
+
+    @classmethod
+    async def call(cls, arguments: str, context: ToolContext) -> message.ToolResultMessage:
+        try:
+            args = json.loads(arguments)
+        except json.JSONDecodeError as exc:
+            return message.ToolResultMessage(status="error", output_text=f"Invalid JSON arguments: {exc}")
+
+        if not isinstance(args, dict):
+            return message.ToolResultMessage(status="error", output_text="Invalid arguments: expected object")
+
+        typed_args = cast(dict[str, Any], args)
+
+        runner = context.run_subtask
+        if runner is None:
+            return message.ToolResultMessage(status="error", output_text="No subtask runner available in this context")
+
+        resume_raw = typed_args.get("resume")
+        resume_session_id: str | None = None
+        if isinstance(resume_raw, str) and resume_raw.strip():
+            try:
+                resume_session_id = Session.resolve_sub_agent_session_id(resume_raw)
+            except ValueError as exc:
+                return message.ToolResultMessage(status="error", output_text=str(exc))
+
+            claims = context.sub_agent_resume_claims
+            if claims is not None:
+                ok = await claims.claim(resume_session_id)
+                if not ok:
+                    return message.ToolResultMessage(
+                        status="error",
+                        output_text=(
+                            "Duplicate sub-agent resume in the same response: "
+                            f"resume='{resume_raw.strip()}' (resolved='{resume_session_id[:7]}…'). "
+                            "Merge into a single call or resume in a later turn."
+                        ),
+                    )
+
+        description = str(typed_args.get("description") or "")
+        prompt = build_image_gen_prompt(typed_args)
+        generation = typed_args.get("generation")
+        generation_dict: dict[str, Any] | None = (
+            cast(dict[str, Any], generation) if isinstance(generation, dict) else None
+        )
+
+        try:
+            profile = get_sub_agent_profile(tools.IMAGE_GEN)
+        except KeyError as exc:
+            return message.ToolResultMessage(status="error", output_text=str(exc))
+
+        try:
+            result = await runner(
+                model.SubAgentState(
+                    sub_agent_type=profile.name,
+                    sub_agent_desc=description,
+                    sub_agent_prompt=prompt,
+                    resume=resume_session_id,
+                    output_schema=None,
+                    generation=generation_dict,
+                ),
+                context.record_sub_agent_session_id,
+                context.register_sub_agent_metadata_getter,
+            )
+        except Exception as exc:
+            return message.ToolResultMessage(status="error", output_text=f"Failed to run subtask: {exc}")
+
+        return message.ToolResultMessage(
+            status="success" if not result.error else "error",
+            output_text=result.task_result,
+            ui_extra=model.SessionIdUIExtra(session_id=result.session_id),
+            task_metadata=result.task_metadata,
+        )

klaude_code/core/tool/sub_agent/task.md

@@ -0,0 +1,20 @@
+Launch a new agent to handle complex, multi-step tasks autonomously.
+
+The Task tool launches specialized agents (subprocesses) that autonomously handle complex tasks. Each agent type has specific capabilities and tools available to it.
+
+When using the Task tool, you must specify a subagent_type parameter to select which agent type to use.
+
+Available agent types and the tools they have access to:
+
+${types_section}
+
+Usage notes:
+- Always include a short description (3-5 words) summarizing what the agent will do
+- Launch multiple agents concurrently whenever possible, to maximize performance; to do that, use a single message with multiple tool uses
+- Clearly tell the agent whether you expect it to write code or just to do research (search, file reads, etc.), since it is not aware of the user's intent
+- Provide clear, detailed prompts so the agent can work autonomously and return exactly the information you need.
+- 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.
+- If the agent description mentions that it should be used proactively, then you should try your best to use it without the user having to ask for it first. Use your judgement.
+- If the user specifies that they want you to run agents "in parallel", you MUST send a single message with multiple Task tool use content blocks. For example, if you need to launch both a code-reviewer agent and a test-runner agent in parallel, send a single message with both tool calls.
+- Agents can provide structured output by passing a JSON Schema in `output_schema`.

klaude_code/core/tool/sub_agent/task.py

@@ -0,0 +1,205 @@
+"""Task tool implementation for running sub-agents by type."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, cast
+
+from klaude_code.core.tool.context import ToolContext
+from klaude_code.core.tool.tool_abc import ToolABC, ToolConcurrencyPolicy, ToolMetadata, load_desc
+from klaude_code.core.tool.tool_registry import register
+from klaude_code.protocol import llm_param, message, model, tools
+from klaude_code.protocol.sub_agent import get_sub_agent_profile, iter_sub_agent_profiles
+from klaude_code.session.session import Session
+
+TASK_TYPE_TO_SUB_AGENT: dict[str, str] = {
+    "general-purpose": "Task",
+    "explore": "Explore",
+    "web": "Web",
+}
+
+
+def _task_description() -> str:
+    summaries: dict[str, str] = {}
+    for profile in iter_sub_agent_profiles():
+        if profile.invoker_type:
+            summaries[profile.invoker_type] = profile.invoker_summary.strip()
+
+    type_lines: list[str] = []
+    for invoker_type in TASK_TYPE_TO_SUB_AGENT:
+        summary = summaries.get(invoker_type, "")
+        if summary:
+            type_lines.append(f"- {invoker_type}: {summary}")
+        else:
+            type_lines.append(f"- {invoker_type}")
+
+    types_section = "\n".join(type_lines) if type_lines else "- general-purpose"
+
+    return load_desc(Path(__file__).parent / "task.md", {"types_section": types_section})
+
+
+TASK_SCHEMA = llm_param.ToolSchema(
+    name=tools.TASK,
+    type="function",
+    description=_task_description(),
+    parameters={
+        "type": "object",
+        "properties": {
+            "type": {
+                "type": "string",
+                "enum": list(TASK_TYPE_TO_SUB_AGENT.keys()),
+                "description": "Sub-agent type selector.",
+            },
+            "description": {
+                "type": "string",
+                "description": "A short (3-5 word) description of the task.",
+            },
+            "prompt": {
+                "type": "string",
+                "description": "The task for the agent to perform.",
+            },
+            "output_schema": {
+                "type": "object",
+                "description": "Optional JSON Schema for structured output.",
+            },
+            "resume": {
+                "type": "string",
+                "description": "Optional agent ID to resume from.",
+            },
+        },
+        "required": ["description", "prompt"],
+        "additionalProperties": False,
+    },
+)
+
+
+@register(tools.TASK)
+class TaskTool(ToolABC):
+    """Run a sub-agent based on the requested type."""
+
+    @classmethod
+    def metadata(cls) -> ToolMetadata:
+        return ToolMetadata(concurrency_policy=ToolConcurrencyPolicy.CONCURRENT, has_side_effects=True)
+
+    @classmethod
+    def schema(cls) -> llm_param.ToolSchema:
+        return TASK_SCHEMA
+
+    @classmethod
+    async def call(cls, arguments: str, context: ToolContext) -> message.ToolResultMessage:
+        try:
+            args = json.loads(arguments)
+        except json.JSONDecodeError as exc:
+            return message.ToolResultMessage(status="error", output_text=f"Invalid JSON arguments: {exc}")
+
+        if not isinstance(args, dict):
+            return message.ToolResultMessage(status="error", output_text="Invalid arguments: expected object")
+
+        typed_args = cast(dict[str, Any], args)
+
+        runner = context.run_subtask
+        if runner is None:
+            return message.ToolResultMessage(status="error", output_text="No subtask runner available in this context")
+
+        description = str(typed_args.get("description") or "")
+
+        resume_raw = typed_args.get("resume")
+        resume_session_id: str | None = None
+        resume_sub_agent_type: str | None = None
+        if isinstance(resume_raw, str) and resume_raw.strip():
+            try:
+                resume_session_id = Session.resolve_sub_agent_session_id(resume_raw)
+            except ValueError as exc:
+                return message.ToolResultMessage(status="error", output_text=str(exc))
+
+            try:
+                resume_session = Session.load(resume_session_id)
+            except (OSError, ValueError, json.JSONDecodeError) as exc:
+                return message.ToolResultMessage(
+                    status="error",
+                    output_text=f"Failed to resume sub-agent session '{resume_session_id}': {exc}",
+                )
+
+            if resume_session.sub_agent_state is None:
+                return message.ToolResultMessage(
+                    status="error",
+                    output_text=f"Invalid resume id '{resume_session_id}': target session is not a sub-agent session",
+                )
+
+            resume_sub_agent_type = resume_session.sub_agent_state.sub_agent_type
+            if resume_sub_agent_type == tools.IMAGE_GEN:
+                return message.ToolResultMessage(
+                    status="error",
+                    output_text="This resume id belongs to ImageGen; use the ImageGen tool to resume it.",
+                )
+
+            claims = context.sub_agent_resume_claims
+            if claims is not None:
+                ok = await claims.claim(resume_session_id)
+                if not ok:
+                    return message.ToolResultMessage(
+                        status="error",
+                        output_text=(
+                            "Duplicate sub-agent resume in the same response: "
+                            f"resume='{resume_raw.strip()}' (resolved='{resume_session_id[:7]}…'). "
+                            "Merge into a single call or resume in a later turn."
+                        ),
+                    )
+
+        type_raw = typed_args.get("type")
+        requested_type = str(type_raw).strip() if isinstance(type_raw, str) else ""
+
+        if resume_session_id and not requested_type:
+            sub_agent_type = resume_sub_agent_type or TASK_TYPE_TO_SUB_AGENT["general-purpose"]
+        else:
+            if not requested_type:
+                requested_type = "general-purpose"
+            sub_agent_type = TASK_TYPE_TO_SUB_AGENT.get(requested_type)
+            if sub_agent_type is None:
+                return message.ToolResultMessage(
+                    status="error",
+                    output_text=f"Unknown Task type '{requested_type}'.",
+                )
+
+        if resume_session_id and resume_sub_agent_type and resume_sub_agent_type != sub_agent_type:
+            return message.ToolResultMessage(
+                status="error",
+                output_text=(
+                    "Invalid resume id: sub-agent type mismatch. "
+                    f"Expected '{sub_agent_type}', got '{resume_sub_agent_type}'."
+                ),
+            )
+
+        try:
+            profile = get_sub_agent_profile(sub_agent_type)
+        except KeyError as exc:
+            return message.ToolResultMessage(status="error", output_text=str(exc))
+
+        sub_agent_prompt = profile.prompt_builder(typed_args)
+
+        output_schema_raw = typed_args.get("output_schema")
+        output_schema = cast(dict[str, Any], output_schema_raw) if isinstance(output_schema_raw, dict) else None
+
+        try:
+            result = await runner(
+                model.SubAgentState(
+                    sub_agent_type=profile.name,
+                    sub_agent_desc=description,
+                    sub_agent_prompt=sub_agent_prompt,
+                    resume=resume_session_id,
+                    output_schema=output_schema,
+                    generation=None,
+                ),
+                context.record_sub_agent_session_id,
+                context.register_sub_agent_metadata_getter,
+            )
+        except Exception as exc:
+            return message.ToolResultMessage(status="error", output_text=f"Failed to run subtask: {exc}")
+
+        return message.ToolResultMessage(
+            status="success" if not result.error else "error",
+            output_text=result.task_result,
+            ui_extra=model.SessionIdUIExtra(session_id=result.session_id),
+            task_metadata=result.task_metadata,
+        )

klaude_code/core/tool/tool_registry.py

@@ -1,10 +1,8 @@
 from collections.abc import Callable
 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
-from klaude_code.protocol.sub_agent import iter_sub_agent_profiles
 
 _REGISTRY: dict[str, type[ToolABC]] = {}
 
@@ -19,20 +17,6 @@ def register(name: str) -> Callable[[type[T]], type[T]]:
     return _decorator
 
 
-def _register_sub_agent_tools() -> None:
-    """Automatically register all sub-agent tools based on their profiles."""
-    for profile in iter_sub_agent_profiles():
-        tool_cls = SubAgentTool.for_profile(profile)
-        _REGISTRY[profile.name] = tool_cls
-
-
-_register_sub_agent_tools()
-
-
-def list_tools() -> list[str]:
-    return list(_REGISTRY.keys())
-
-
 def get_tool_schemas(tool_names: list[str]) -> list[llm_param.ToolSchema]:
     schemas: list[llm_param.ToolSchema] = []
     for tool_name in tool_names:

klaude_code/core/turn.py

@@ -243,7 +243,7 @@ class TurnExecutor:
         )
 
         # ImageGen per-call overrides (tool-level `generation` parameters)
-        if ctx.sub_agent_state is not None and ctx.sub_agent_state.sub_agent_type == "ImageGen":
+        if ctx.sub_agent_state is not None and ctx.sub_agent_state.sub_agent_type == tools.IMAGE_GEN:
             call_param.modalities = ["image", "text"]
             generation = ctx.sub_agent_state.generation or {}
             image_config = llm_param.ImageConfig()

klaude_code/llm/anthropic/input.py

@@ -18,9 +18,10 @@ from anthropic.types.beta.beta_tool_use_block_param import BetaToolUseBlockParam
 from anthropic.types.beta.beta_url_image_source_param import BetaURLImageSourceParam
 
 from klaude_code.const import EMPTY_TOOL_OUTPUT_MESSAGE
-from klaude_code.llm.image import parse_data_url
+from klaude_code.llm.image import image_file_to_data_url, parse_data_url
 from klaude_code.llm.input_common import (
     DeveloperAttachment,
+    ImagePart,
     attach_developer_messages,
     merge_reminder_text,
     split_thinking_parts,
@@ -36,8 +37,8 @@ _INLINE_IMAGE_MEDIA_TYPES: tuple[AllowedMediaType, ...] = (
 )
 
 
-def _image_part_to_block(image: message.ImageURLPart) -> BetaImageBlockParam:
-    url = image.url
+def _image_part_to_block(image: ImagePart) -> BetaImageBlockParam:
+    url = image_file_to_data_url(image) if isinstance(image, message.ImageFilePart) else image.url
     if url.startswith("data:"):
         media_type, base64_payload, _ = parse_data_url(url)
         if media_type not in _INLINE_IMAGE_MEDIA_TYPES:
@@ -64,7 +65,7 @@ def _user_message_to_message(
     for part in msg.parts:
         if isinstance(part, message.TextPart):
             blocks.append(cast(BetaTextBlockParam, {"type": "text", "text": part.text}))
-        elif isinstance(part, message.ImageURLPart):
+        elif isinstance(part, (message.ImageURLPart, message.ImageFilePart)):
             blocks.append(_image_part_to_block(part))
     if attachment.text:
         blocks.append(cast(BetaTextBlockParam, {"type": "text", "text": attachment.text}))
@@ -86,7 +87,7 @@ def _tool_message_to_block(
         attachment.text,
     )
     tool_content.append(cast(BetaTextBlockParam, {"type": "text", "text": merged_text}))
-    for image in [part for part in msg.parts if isinstance(part, message.ImageURLPart)]:
+    for image in [part for part in msg.parts if isinstance(part, (message.ImageURLPart, message.ImageFilePart))]:
         tool_content.append(_image_part_to_block(image))
     for image in attachment.images:
         tool_content.append(_image_part_to_block(image))