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

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))

klaude_code/llm/antigravity/input.py

@@ -6,9 +6,10 @@ from binascii import Error as BinasciiError
 from typing import Any, TypedDict
 
 from klaude_code.const import EMPTY_TOOL_OUTPUT_MESSAGE
-from klaude_code.llm.image import assistant_image_to_data_url, parse_data_url
+from klaude_code.llm.image import assistant_image_to_data_url, 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,
@@ -66,9 +67,9 @@ def _data_url_to_inline_data(url: str) -> InlineData:
     return InlineData(mimeType=media_type, data=base64.b64encode(decoded).decode("ascii"))
 
 
-def _image_part_to_part(image: message.ImageURLPart) -> Part:
-    """Convert ImageURLPart to Part dict."""
-    url = image.url
+def _image_part_to_part(image: ImagePart) -> Part:
+    """Convert ImageURLPart or ImageFilePart to Part dict."""
+    url = image_file_to_data_url(image) if isinstance(image, message.ImageFilePart) else image.url
     if url.startswith("data:"):
         return Part(inlineData=_data_url_to_inline_data(url))
     # For non-data URLs, best-effort using inline_data format
@@ -81,7 +82,7 @@ def _user_message_to_content(msg: message.UserMessage, attachment: DeveloperAtta
     for part in msg.parts:
         if isinstance(part, message.TextPart):
             parts.append(Part(text=part.text))
-        elif isinstance(part, message.ImageURLPart):
+        elif isinstance(part, (message.ImageURLPart, message.ImageFilePart)):
             parts.append(_image_part_to_part(part))
     if attachment.text:
         parts.append(Part(text=attachment.text))
@@ -108,14 +109,20 @@ def _tool_messages_to_contents(
         )
         has_text = merged_text.strip() != ""
 
-        images = [part for part in msg.parts if isinstance(part, message.ImageURLPart)] + attachment.images
+        images: list[ImagePart] = [
+            part for part in msg.parts if isinstance(part, (message.ImageURLPart, message.ImageFilePart))
+        ]
+        images.extend(attachment.images)
         image_parts: list[Part] = []
         function_response_parts: list[dict[str, Any]] = []
 
         for image in images:
             try:
                 image_parts.append(_image_part_to_part(image))
-                if image.url.startswith("data:"):
+                if isinstance(image, message.ImageFilePart):
+                    inline_data = _data_url_to_inline_data(image_file_to_data_url(image))
+                    function_response_parts.append({"inlineData": inline_data})
+                elif image.url.startswith("data:"):
                     inline_data = _data_url_to_inline_data(image.url)
                     function_response_parts.append({"inlineData": inline_data})
             except ValueError:

klaude_code/llm/codex/client.py

@@ -146,6 +146,28 @@ class CodexClient(LLMClientABC):
             )
         except (openai.OpenAIError, httpx.HTTPError) as e:
             error_message = f"{e.__class__.__name__} {e!s}"
+
+            # Check for invalid instruction error and invalidate prompt cache
+            if _is_invalid_instruction_error(e) and param.model_id:
+                _invalidate_prompt_cache_for_model(param.model_id)
+
             return error_llm_stream(metadata_tracker, error=error_message)
 
         return ResponsesLLMStream(stream, param=param, metadata_tracker=metadata_tracker)
+
+
+def _is_invalid_instruction_error(e: Exception) -> bool:
+    """Check if the error is related to invalid instructions."""
+    error_str = str(e).lower()
+    return "invalid instruction" in error_str or "invalid_instruction" in error_str
+
+
+def _invalidate_prompt_cache_for_model(model_id: str) -> None:
+    """Invalidate the cached prompt for a model to force refresh."""
+    from klaude_code.llm.codex.prompt_sync import invalidate_cache
+
+    log_debug(
+        f"Invalidating prompt cache for model {model_id} due to invalid instruction error",
+        debug_type=DebugType.GENERAL,
+    )
+    invalidate_cache(model_id)

klaude_code/llm/codex/prompt_sync.py

@@ -0,0 +1,237 @@
+"""Dynamic prompt synchronization from OpenAI Codex GitHub repository."""
+
+import json
+import time
+from functools import cache
+from importlib.resources import files
+from pathlib import Path
+from typing import Any, Literal
+
+import httpx
+
+from klaude_code.log import DebugType, log_debug
+
+GITHUB_API_RELEASES = "https://api.github.com/repos/openai/codex/releases/latest"
+GITHUB_HTML_RELEASES = "https://github.com/openai/codex/releases/latest"
+GITHUB_RAW_BASE = "https://raw.githubusercontent.com/openai/codex"
+
+CACHE_DIR = Path.home() / ".klaude" / "codex-prompts"
+CACHE_TTL_SECONDS = 24 * 60 * 60  # 24 hours
+
+type ModelFamily = Literal["gpt-5.2-codex", "codex-max", "codex", "gpt-5.2", "gpt-5.1"]
+
+PROMPT_FILES: dict[ModelFamily, str] = {
+    "gpt-5.2-codex": "gpt-5.2-codex_prompt.md",
+    "codex-max": "gpt-5.1-codex-max_prompt.md",
+    "codex": "gpt_5_codex_prompt.md",
+    "gpt-5.2": "gpt_5_2_prompt.md",
+    "gpt-5.1": "gpt_5_1_prompt.md",
+}
+
+CACHE_FILES: dict[ModelFamily, str] = {
+    "gpt-5.2-codex": "gpt-5.2-codex-instructions.md",
+    "codex-max": "codex-max-instructions.md",
+    "codex": "codex-instructions.md",
+    "gpt-5.2": "gpt-5.2-instructions.md",
+    "gpt-5.1": "gpt-5.1-instructions.md",
+}
+
+
+@cache
+def _load_bundled_prompt(prompt_path: str) -> str:
+    """Load bundled prompt from package resources."""
+    return files("klaude_code.core").joinpath(prompt_path).read_text(encoding="utf-8").strip()
+
+
+class CacheMetadata:
+    def __init__(self, etag: str | None, tag: str, last_checked: int, url: str):
+        self.etag = etag
+        self.tag = tag
+        self.last_checked = last_checked
+        self.url = url
+
+    def to_dict(self) -> dict[str, str | int | None]:
+        return {
+            "etag": self.etag,
+            "tag": self.tag,
+            "last_checked": self.last_checked,
+            "url": self.url,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, object]) -> "CacheMetadata":
+        etag = data.get("etag")
+        last_checked = data.get("last_checked")
+        return cls(
+            etag=etag if isinstance(etag, str) else None,
+            tag=str(data.get("tag", "")),
+            last_checked=int(last_checked) if isinstance(last_checked, int | float) else 0,
+            url=str(data.get("url", "")),
+        )
+
+
+def get_model_family(model: str) -> ModelFamily:
+    """Determine model family from model name."""
+    if "gpt-5.2-codex" in model or "gpt 5.2 codex" in model:
+        return "gpt-5.2-codex"
+    if "codex-max" in model:
+        return "codex-max"
+    if "codex" in model or model.startswith("codex-"):
+        return "codex"
+    if "gpt-5.2" in model:
+        return "gpt-5.2"
+    return "gpt-5.1"
+
+
+def _get_latest_release_tag(client: httpx.Client) -> str:
+    """Get latest release tag from GitHub."""
+    try:
+        response = client.get(GITHUB_API_RELEASES)
+        if response.status_code == 200:
+            data: dict[str, Any] = response.json()
+            tag_name: Any = data.get("tag_name")
+            if isinstance(tag_name, str):
+                return tag_name
+    except httpx.HTTPError:
+        pass
+
+    # Fallback: follow redirect from releases/latest
+    response = client.get(GITHUB_HTML_RELEASES, follow_redirects=True)
+    if response.status_code == 200:
+        final_url = str(response.url)
+        if "/tag/" in final_url:
+            parts = final_url.split("/tag/")
+            if len(parts) > 1 and "/" not in parts[-1]:
+                return parts[-1]
+
+    raise RuntimeError("Failed to determine latest release tag from GitHub")
+
+
+def _load_cache_metadata(meta_file: Path) -> CacheMetadata | None:
+    if not meta_file.exists():
+        return None
+    try:
+        data = json.loads(meta_file.read_text())
+        return CacheMetadata.from_dict(data)
+    except (json.JSONDecodeError, ValueError):
+        return None
+
+
+def _save_cache_metadata(meta_file: Path, metadata: CacheMetadata) -> None:
+    meta_file.parent.mkdir(parents=True, exist_ok=True)
+    meta_file.write_text(json.dumps(metadata.to_dict(), indent=2))
+
+
+def get_codex_instructions(model: str = "gpt-5.1-codex", force_refresh: bool = False) -> str:
+    """Get Codex instructions for the given model.
+
+    Args:
+        model: Model name to get instructions for.
+        force_refresh: If True, bypass cache TTL and fetch fresh instructions.
+
+    Returns:
+        The Codex system prompt instructions.
+    """
+    model_family = get_model_family(model)
+    prompt_file = PROMPT_FILES[model_family]
+    cache_file = CACHE_DIR / CACHE_FILES[model_family]
+    meta_file = CACHE_DIR / f"{CACHE_FILES[model_family].replace('.md', '-meta.json')}"
+
+    # Check cache unless force refresh
+    if not force_refresh:
+        metadata = _load_cache_metadata(meta_file)
+        if metadata and cache_file.exists():
+            age = int(time.time()) - metadata.last_checked
+            if age < CACHE_TTL_SECONDS:
+                log_debug(f"Using cached {model_family} instructions (age: {age}s)", debug_type=DebugType.GENERAL)
+                return cache_file.read_text()
+
+    try:
+        with httpx.Client(timeout=30.0) as client:
+            latest_tag = _get_latest_release_tag(client)
+            instructions_url = f"{GITHUB_RAW_BASE}/{latest_tag}/codex-rs/core/{prompt_file}"
+
+            # Load existing metadata for conditional request
+            metadata = _load_cache_metadata(meta_file)
+            headers: dict[str, str] = {}
+
+            # Only use ETag if tag matches (different release = different content)
+            if metadata and metadata.tag == latest_tag and metadata.etag:
+                headers["If-None-Match"] = metadata.etag
+
+            response = client.get(instructions_url, headers=headers)
+
+            if response.status_code == 304 and cache_file.exists():
+                # Not modified, update last_checked and return cached
+                if metadata:
+                    metadata.last_checked = int(time.time())
+                    _save_cache_metadata(meta_file, metadata)
+                log_debug(f"Codex {model_family} instructions not modified", debug_type=DebugType.GENERAL)
+                return cache_file.read_text()
+
+            if response.status_code == 200:
+                instructions = response.text
+                new_etag = response.headers.get("etag")
+
+                # Save to cache
+                cache_file.parent.mkdir(parents=True, exist_ok=True)
+                cache_file.write_text(instructions)
+                _save_cache_metadata(
+                    meta_file,
+                    CacheMetadata(
+                        etag=new_etag,
+                        tag=latest_tag,
+                        last_checked=int(time.time()),
+                        url=instructions_url,
+                    ),
+                )
+
+                log_debug(f"Updated {model_family} instructions from GitHub", debug_type=DebugType.GENERAL)
+                return instructions
+
+            raise RuntimeError(f"HTTP {response.status_code}")
+
+    except Exception as e:
+        log_debug(f"Failed to fetch {model_family} instructions: {e}", debug_type=DebugType.GENERAL)
+
+        # Fallback to cached version
+        if cache_file.exists():
+            log_debug(f"Using cached {model_family} instructions (fallback)", debug_type=DebugType.GENERAL)
+            return cache_file.read_text()
+
+        # Last resort: use bundled prompt
+        bundled_path = _get_bundled_prompt_path(model_family)
+        if bundled_path:
+            log_debug(f"Using bundled {model_family} instructions (fallback)", debug_type=DebugType.GENERAL)
+            return _load_bundled_prompt(bundled_path)
+
+        raise RuntimeError(f"No Codex instructions available for {model_family}") from e
+
+
+def _get_bundled_prompt_path(model_family: ModelFamily) -> str | None:
+    """Get bundled prompt path for model family."""
+    if model_family == "gpt-5.2-codex":
+        return "prompts/prompt-codex-gpt-5-2-codex.md"
+    if model_family == "gpt-5.2":
+        return "prompts/prompt-codex-gpt-5-2.md"
+    if model_family in ("codex", "codex-max", "gpt-5.1"):
+        return "prompts/prompt-codex.md"
+    return None
+
+
+def invalidate_cache(model: str | None = None) -> None:
+    """Invalidate cached instructions to force refresh on next access.
+
+    Args:
+        model: If provided, only invalidate cache for this model's family.
+               If None, invalidate all cached instructions.
+    """
+    if model:
+        model_family = get_model_family(model)
+        meta_file = CACHE_DIR / f"{CACHE_FILES[model_family].replace('.md', '-meta.json')}"
+        if meta_file.exists():
+            meta_file.unlink()
+    else:
+        if CACHE_DIR.exists():
+            for meta_file in CACHE_DIR.glob("*-meta.json"):
+                meta_file.unlink()

klaude_code/llm/google/client.py

@@ -3,7 +3,6 @@
 # pyright: reportUnknownArgumentType=false
 # pyright: reportAttributeAccessIssue=false
 
-import json
 from base64 import b64encode
 from collections.abc import AsyncGenerator, AsyncIterator
 from typing import Any, cast, override
@@ -33,6 +32,7 @@ from klaude_code.llm.client import LLMClientABC, LLMStreamABC
 from klaude_code.llm.google.input import convert_history_to_contents, convert_tool_schema
 from klaude_code.llm.image import save_assistant_image
 from klaude_code.llm.input_common import apply_config_defaults
+from klaude_code.llm.json_stable import dumps_canonical_json
 from klaude_code.llm.registry import register
 from klaude_code.llm.stream_parts import (
     append_text_part,
@@ -122,6 +122,8 @@ def _usage_from_metadata(
     if usage is None:
         return None
 
+    # In Gemini usage metadata, prompt_token_count represents the full prompt tokens
+    # (including cached tokens). cached_content_token_count is a subset of prompt tokens.
     cached = usage.cached_content_token_count or 0
     prompt = usage.prompt_token_count or 0
     response = usage.candidates_token_count or 0
@@ -136,10 +138,10 @@ def _usage_from_metadata(
 
     total = usage.total_token_count
     if total is None:
-        total = prompt + cached + response + thoughts
+        total = prompt + response + thoughts
 
     return model.Usage(
-        input_tokens=prompt + cached,
+        input_tokens=prompt,
         cached_tokens=cached,
         output_tokens=response + thoughts,
         reasoning_tokens=thoughts,
@@ -385,7 +387,7 @@ async def parse_google_stream(
             args_obj = function_call.args
             if args_obj is not None:
                 # Add ToolCallPart, then ThinkingSignaturePart after it
-                state.append_tool_call(call_id, name, json.dumps(args_obj, ensure_ascii=False))
+                state.append_tool_call(call_id, name, dumps_canonical_json(args_obj))
                 encoded_sig = _encode_thought_signature(thought_signature)
                 if encoded_sig:
                     state.append_thinking_signature(encoded_sig)
@@ -400,7 +402,7 @@ async def parse_google_stream(
             will_continue = function_call.will_continue
             if will_continue is False and call_id in partial_args_by_call and call_id not in completed_tool_items:
                 # Add ToolCallPart, then ThinkingSignaturePart after it
-                state.append_tool_call(call_id, name, json.dumps(partial_args_by_call[call_id], ensure_ascii=False))
+                state.append_tool_call(call_id, name, dumps_canonical_json(partial_args_by_call[call_id]))
                 stored_sig = started_tool_calls.get(call_id, (name, None))[1]
                 encoded_stored_sig = _encode_thought_signature(stored_sig)
                 if encoded_stored_sig:
@@ -412,7 +414,7 @@ async def parse_google_stream(
         if call_id in completed_tool_items:
             continue
         args = partial_args_by_call.get(call_id, {})
-        state.append_tool_call(call_id, name, json.dumps(args, ensure_ascii=False))
+        state.append_tool_call(call_id, name, dumps_canonical_json(args))
         encoded_stored_sig = _encode_thought_signature(stored_sig)
         if encoded_stored_sig:
             state.append_thinking_signature(encoded_stored_sig)