klaude-code 1.9.0__py3-none-any.whl → 2.0.0__py3-none-any.whl

klaude_code/core/tool/web/web_fetch_tool.py

@@ -10,14 +10,17 @@ from urllib.parse import quote, urlparse, urlunparse
 
 from pydantic import BaseModel
 
-from klaude_code import const
+from klaude_code.const import (
+    TOOL_OUTPUT_TRUNCATION_DIR,
+    URL_FILENAME_MAX_LENGTH,
+    WEB_FETCH_DEFAULT_TIMEOUT_SEC,
+    WEB_FETCH_USER_AGENT,
+)
 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, model, tools
+from klaude_code.protocol import llm_param, message, tools
 
-DEFAULT_TIMEOUT_SEC = 30
-DEFAULT_USER_AGENT = "Mozilla/5.0 (compatible; KlaudeCode/1.0)"
-WEB_FETCH_SAVE_DIR = Path(const.TOOL_OUTPUT_TRUNCATION_DIR) / "web"
+WEB_FETCH_SAVE_DIR = Path(TOOL_OUTPUT_TRUNCATION_DIR) / "web"
 
 
 def _encode_url(url: str) -> str:
@@ -110,16 +113,16 @@ def _extract_url_filename(url: str) -> str:
     path = parsed.path.strip("/").replace("/", "_")
     name = f"{host}_{path}" if path else host
     name = re.sub(r"[^a-zA-Z0-9_\-]", "_", name)
-    return name[:80] if len(name) > 80 else name
+    return name[:URL_FILENAME_MAX_LENGTH] if len(name) > URL_FILENAME_MAX_LENGTH else name
 
 
-def _save_web_content(url: str, content: str) -> str | None:
+def _save_web_content(url: str, content: str, extension: str = ".md") -> str | None:
     """Save web content to file. Returns file path or None on failure."""
     try:
         WEB_FETCH_SAVE_DIR.mkdir(parents=True, exist_ok=True)
         timestamp = int(time.time())
         identifier = _extract_url_filename(url)
-        filename = f"{identifier}-{timestamp}.md"
+        filename = f"{identifier}-{timestamp}{extension}"
         file_path = WEB_FETCH_SAVE_DIR / filename
         file_path.write_text(content, encoding="utf-8")
         return str(file_path)
@@ -127,6 +130,26 @@ def _save_web_content(url: str, content: str) -> str | None:
         return None
 
 
+def _save_binary_content(url: str, data: bytes, extension: str = ".bin") -> str | None:
+    """Save binary content to file. Returns file path or None on failure."""
+    try:
+        WEB_FETCH_SAVE_DIR.mkdir(parents=True, exist_ok=True)
+        timestamp = int(time.time())
+        identifier = _extract_url_filename(url)
+        filename = f"{identifier}-{timestamp}{extension}"
+        file_path = WEB_FETCH_SAVE_DIR / filename
+        file_path.write_bytes(data)
+        return str(file_path)
+    except OSError:
+        return None
+
+
+def _is_pdf_url(url: str) -> bool:
+    """Check if URL points to a PDF file."""
+    parsed = urlparse(url)
+    return parsed.path.lower().endswith(".pdf") or "/pdf/" in parsed.path.lower()
+
+
 def _process_content(content_type: str, text: str) -> str:
     """Process content based on Content-Type header."""
     if content_type == "text/html":
@@ -139,19 +162,19 @@ def _process_content(content_type: str, text: str) -> str:
         return text
 
 
-def _fetch_url(url: str, timeout: int = DEFAULT_TIMEOUT_SEC) -> tuple[str, str]:
+def _fetch_url(url: str, timeout: int = WEB_FETCH_DEFAULT_TIMEOUT_SEC) -> tuple[str, bytes, str | None]:
     """
     Fetch URL content synchronously.
 
     Returns:
-        Tuple of (content_type, response_text)
+        Tuple of (content_type, raw_data, charset)
 
     Raises:
         Various exceptions on failure
     """
     headers = {
         "Accept": "text/markdown, */*",
-        "User-Agent": DEFAULT_USER_AGENT,
+        "User-Agent": WEB_FETCH_USER_AGENT,
     }
     encoded_url = _encode_url(url)
     request = urllib.request.Request(encoded_url, headers=headers)
@@ -159,8 +182,7 @@ def _fetch_url(url: str, timeout: int = DEFAULT_TIMEOUT_SEC) -> tuple[str, str]:
     with urllib.request.urlopen(request, timeout=timeout) as response:
         content_type, charset = _extract_content_type_and_charset(response)
         data = response.read()
-        text = _decode_content(data, charset)
-        return content_type, text
+        return content_type, data, charset
 
 
 @register(tools.WEB_FETCH)
@@ -191,29 +213,45 @@ class WebFetchTool(ToolABC):
         url: str
 
     @classmethod
-    async def call(cls, arguments: str) -> model.ToolResultItem:
+    async def call(cls, arguments: str) -> message.ToolResultMessage:
         try:
             args = WebFetchTool.WebFetchArguments.model_validate_json(arguments)
         except ValueError as e:
-            return model.ToolResultItem(
+            return message.ToolResultMessage(
                 status="error",
-                output=f"Invalid arguments: {e}",
+                output_text=f"Invalid arguments: {e}",
             )
         return await cls.call_with_args(args)
 
     @classmethod
-    async def call_with_args(cls, args: WebFetchArguments) -> model.ToolResultItem:
+    async def call_with_args(cls, args: WebFetchArguments) -> message.ToolResultMessage:
         url = args.url
 
         # Basic URL validation
         if not url.startswith(("http://", "https://")):
-            return model.ToolResultItem(
+            return message.ToolResultMessage(
                 status="error",
-                output=f"Invalid URL: must start with http:// or https:// (url={url})",
+                output_text=f"Invalid URL: must start with http:// or https:// (url={url})",
             )
 
         try:
-            content_type, text = await asyncio.to_thread(_fetch_url, url)
+            content_type, data, charset = await asyncio.to_thread(_fetch_url, url)
+
+            # Handle PDF files
+            if content_type == "application/pdf" or _is_pdf_url(url):
+                saved_path = _save_binary_content(url, data, ".pdf")
+                if saved_path:
+                    return message.ToolResultMessage(
+                        status="success",
+                        output_text=f"PDF file saved to: {saved_path}\n\nTo read the PDF content, use the Read tool on this file path.",
+                    )
+                return message.ToolResultMessage(
+                    status="error",
+                    output_text=f"Failed to save PDF file (url={url})",
+                )
+
+            # Handle text content
+            text = _decode_content(data, charset)
             processed = _process_content(content_type, text)
 
             # Always save content to file
@@ -222,28 +260,28 @@ class WebFetchTool(ToolABC):
             # Build output with file path info
             output = f"<file_saved>{saved_path}</file_saved>\n\n{processed}" if saved_path else processed
 
-            return model.ToolResultItem(
+            return message.ToolResultMessage(
                 status="success",
-                output=output,
+                output_text=output,
             )
 
         except urllib.error.HTTPError as e:
-            return model.ToolResultItem(
+            return message.ToolResultMessage(
                 status="error",
-                output=f"HTTP error {e.code}: {e.reason} (url={url})",
+                output_text=f"HTTP error {e.code}: {e.reason} (url={url})",
             )
         except urllib.error.URLError as e:
-            return model.ToolResultItem(
+            return message.ToolResultMessage(
                 status="error",
-                output=f"URL error: {e.reason} (url={url})",
+                output_text=f"URL error: {e.reason} (url={url})",
             )
         except TimeoutError:
-            return model.ToolResultItem(
+            return message.ToolResultMessage(
                 status="error",
-                output=f"Request timed out after {DEFAULT_TIMEOUT_SEC} seconds (url={url})",
+                output_text=f"Request timed out after {WEB_FETCH_DEFAULT_TIMEOUT_SEC} seconds (url={url})",
             )
         except Exception as e:
-            return model.ToolResultItem(
+            return message.ToolResultMessage(
                 status="error",
-                output=f"Failed to fetch URL: {e} (url={url})",
+                output_text=f"Failed to fetch URL: {e} (url={url})",
             )

klaude_code/core/tool/web/web_search_tool.py

@@ -4,12 +4,10 @@ from pathlib import Path
 
 from pydantic import BaseModel
 
+from klaude_code.const import WEB_SEARCH_DEFAULT_MAX_RESULTS, WEB_SEARCH_MAX_RESULTS_LIMIT
 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, model, tools
-
-DEFAULT_MAX_RESULTS = 10
-MAX_RESULTS_LIMIT = 20
+from klaude_code.protocol import llm_param, message, tools
 
 
 @dataclass
@@ -81,7 +79,7 @@ class WebSearchTool(ToolABC):
                     },
                     "max_results": {
                         "type": "integer",
-                        "description": f"Maximum number of results to return (default: {DEFAULT_MAX_RESULTS}, max: {MAX_RESULTS_LIMIT})",
+                        "description": f"Maximum number of results to return (default: {WEB_SEARCH_DEFAULT_MAX_RESULTS}, max: {WEB_SEARCH_MAX_RESULTS_LIMIT})",
                     },
                 },
                 "required": ["query"],
@@ -90,41 +88,41 @@ class WebSearchTool(ToolABC):
 
     class WebSearchArguments(BaseModel):
         query: str
-        max_results: int = DEFAULT_MAX_RESULTS
+        max_results: int = WEB_SEARCH_DEFAULT_MAX_RESULTS
 
     @classmethod
-    async def call(cls, arguments: str) -> model.ToolResultItem:
+    async def call(cls, arguments: str) -> message.ToolResultMessage:
         try:
             args = WebSearchTool.WebSearchArguments.model_validate_json(arguments)
         except ValueError as e:
-            return model.ToolResultItem(
+            return message.ToolResultMessage(
                 status="error",
-                output=f"Invalid arguments: {e}",
+                output_text=f"Invalid arguments: {e}",
             )
         return await cls.call_with_args(args)
 
     @classmethod
-    async def call_with_args(cls, args: WebSearchArguments) -> model.ToolResultItem:
+    async def call_with_args(cls, args: WebSearchArguments) -> message.ToolResultMessage:
         query = args.query.strip()
         if not query:
-            return model.ToolResultItem(
+            return message.ToolResultMessage(
                 status="error",
-                output="Query cannot be empty",
+                output_text="Query cannot be empty",
             )
 
-        max_results = min(max(args.max_results, 1), MAX_RESULTS_LIMIT)
+        max_results = min(max(args.max_results, 1), WEB_SEARCH_MAX_RESULTS_LIMIT)
 
         try:
             results = await asyncio.to_thread(_search_duckduckgo, query, max_results)
             formatted = _format_results(results)
 
-            return model.ToolResultItem(
+            return message.ToolResultMessage(
                 status="success",
-                output=formatted,
+                output_text=formatted,
             )
 
         except Exception as e:
-            return model.ToolResultItem(
+            return message.ToolResultMessage(
                 status="error",
-                output=f"Search failed: {e}",
+                output_text=f"Search failed: {e}",
             )

klaude_code/core/turn.py

@@ -4,12 +4,15 @@ from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
+from klaude_code.const import INTERRUPT_MARKER, SUPPORTED_IMAGE_SIZES
 from klaude_code.core.tool import ToolABC, tool_context
+from klaude_code.core.tool.tool_context import current_sub_agent_resume_claims
 
 if TYPE_CHECKING:
     from klaude_code.core.task import SessionContext
 
 from klaude_code.core.tool.tool_runner import (
+    ToolCallRequest,
     ToolExecutionCallStarted,
     ToolExecutionResult,
     ToolExecutionTodoChange,
@@ -17,7 +20,7 @@ from klaude_code.core.tool.tool_runner import (
     ToolExecutorEvent,
 )
 from klaude_code.llm import LLMClientABC
-from klaude_code.protocol import events, llm_param, model, tools
+from klaude_code.protocol import events, llm_param, message, model, tools
 from klaude_code.trace import DebugType, log_debug
 
 
@@ -36,16 +39,16 @@ class TurnExecutionContext:
     system_prompt: str | None
     tools: list[llm_param.ToolSchema]
     tool_registry: dict[str, type[ToolABC]]
+    sub_agent_state: model.SubAgentState | None = None
 
 
 @dataclass
 class TurnResult:
     """Aggregated state produced while executing a turn."""
 
-    reasoning_items: list[model.ReasoningTextItem | model.ReasoningEncryptedItem]
-    assistant_message: model.AssistantMessageItem | None
-    tool_calls: list[model.ToolCallItem]
-    stream_error: model.StreamErrorItem | None
+    assistant_message: message.AssistantMessage | None
+    tool_calls: list[ToolCallRequest]
+    stream_error: message.StreamErrorItem | None
     report_back_result: str | None = field(default=None)
 
 
@@ -61,23 +64,27 @@ def build_events_from_tool_executor_event(session_id: str, event: ToolExecutorEv
                     session_id=session_id,
                     response_id=tool_call.response_id,
                     tool_call_id=tool_call.call_id,
-                    tool_name=tool_call.name,
-                    arguments=tool_call.arguments,
+                    tool_name=tool_call.tool_name,
+                    arguments=tool_call.arguments_json,
                 )
             )
-        case ToolExecutionResult(tool_call=tool_call, tool_result=tool_result):
+        case ToolExecutionResult(tool_call=tool_call, tool_result=tool_result, is_last_in_turn=is_last_in_turn):
+            status = "success" if tool_result.status == "success" else "error"
             ui_events.append(
                 events.ToolResultEvent(
                     session_id=session_id,
                     response_id=tool_call.response_id,
                     tool_call_id=tool_call.call_id,
-                    tool_name=tool_call.name,
-                    result=tool_result.output or "",
+                    tool_name=tool_call.tool_name,
+                    result=tool_result.output_text,
                     ui_extra=tool_result.ui_extra,
-                    status=tool_result.status,
+                    status=status,
                     task_metadata=tool_result.task_metadata,
+                    is_last_in_turn=is_last_in_turn,
                 )
             )
+            if tool_result.status == "aborted":
+                ui_events.append(events.InterruptEvent(session_id=session_id))
         case ToolExecutionTodoChange(todos=todos):
             ui_events.append(
                 events.TodoChangeEvent(
@@ -129,7 +136,10 @@ class TurnExecutor:
         if self._turn_result is not None and self._turn_result.report_back_result is not None:
             return self._turn_result.report_back_result
         if self._turn_result is not None and self._turn_result.assistant_message is not None:
-            return self._turn_result.assistant_message.content or ""
+            assistant_message = self._turn_result.assistant_message
+            text = message.join_text_parts(assistant_message.parts)
+            images = [part for part in assistant_message.parts if isinstance(part, message.ImageFilePart)]
+            return message.format_saved_images(images, text)
         return ""
 
     @property
@@ -160,7 +170,6 @@ class TurnExecutor:
         yield events.TurnStartEvent(session_id=session_ctx.session_id)
 
         self._turn_result = TurnResult(
-            reasoning_items=[],
             assistant_message=None,
             tool_calls=[],
             stream_error=None,
@@ -188,8 +197,8 @@ class TurnExecutor:
     def _detect_report_back(self, turn_result: TurnResult) -> None:
         """Detect report_back tool call and store its arguments as JSON string."""
         for tool_call in turn_result.tool_calls:
-            if tool_call.name == tools.REPORT_BACK:
-                turn_result.report_back_result = tool_call.arguments
+            if tool_call.tool_name == tools.REPORT_BACK:
+                turn_result.report_back_result = tool_call.arguments_json
                 break
 
     async def _consume_llm_stream(self, turn_result: TurnResult) -> AsyncGenerator[events.Event]:
@@ -197,98 +206,134 @@ class TurnExecutor:
 
         ctx = self._context
         session_ctx = ctx.session_ctx
-        async for response_item in ctx.llm_client.call(
-            llm_param.LLMCallParameter(
-                input=session_ctx.get_conversation_history(),
-                system=ctx.system_prompt,
-                tools=ctx.tools,
-                session_id=session_ctx.session_id,
-            )
-        ):
+        message_types = (
+            message.SystemMessage,
+            message.DeveloperMessage,
+            message.UserMessage,
+            message.AssistantMessage,
+            message.ToolResultMessage,
+        )
+        messages = [item for item in session_ctx.get_conversation_history() if isinstance(item, message_types)]
+        call_param = llm_param.LLMCallParameter(
+            input=messages,
+            system=ctx.system_prompt,
+            tools=ctx.tools,
+            session_id=session_ctx.session_id,
+        )
+
+        # 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":
+            call_param.modalities = ["image", "text"]
+            generation = ctx.sub_agent_state.generation or {}
+            image_config = llm_param.ImageConfig()
+            aspect_ratio = generation.get("aspect_ratio")
+            if isinstance(aspect_ratio, str) and aspect_ratio.strip():
+                image_config.aspect_ratio = aspect_ratio.strip()
+            image_size = generation.get("image_size")
+            if image_size in SUPPORTED_IMAGE_SIZES:
+                image_config.image_size = image_size
+            extra = generation.get("extra")
+            if isinstance(extra, dict) and extra:
+                image_config.extra = extra
+            if image_config.model_dump(exclude_none=True):
+                call_param.image_config = image_config
+
+        async for delta in ctx.llm_client.call(call_param):
             log_debug(
-                f"[{response_item.__class__.__name__}]",
-                response_item.model_dump_json(exclude_none=True),
+                f"[{delta.__class__.__name__}]",
+                delta.model_dump_json(exclude_none=True),
                 style="green",
                 debug_type=DebugType.RESPONSE,
             )
-            match response_item:
-                case model.StartItem():
-                    continue
-                case model.ReasoningTextItem() as item:
-                    turn_result.reasoning_items.append(item)
-                    yield events.ThinkingEvent(
-                        content=item.content,
-                        response_id=item.response_id,
-                        session_id=session_ctx.session_id,
-                    )
-                case model.ReasoningEncryptedItem() as item:
-                    turn_result.reasoning_items.append(item)
-                case model.ReasoningTextDelta() as item:
+            match delta:
+                case message.ThinkingTextDelta() as delta:
                     yield events.ThinkingDeltaEvent(
-                        content=item.content,
-                        response_id=item.response_id,
+                        content=delta.content,
+                        response_id=delta.response_id,
                         session_id=session_ctx.session_id,
                     )
-                case model.AssistantMessageDelta() as item:
-                    if item.response_id:
-                        self._assistant_response_id = item.response_id
-                    self._assistant_delta_buffer.append(item.content)
-                    yield events.AssistantMessageDeltaEvent(
-                        content=item.content,
-                        response_id=item.response_id,
+                case message.AssistantTextDelta() as delta:
+                    if delta.response_id:
+                        self._assistant_response_id = delta.response_id
+                    self._assistant_delta_buffer.append(delta.content)
+                    yield events.AssistantTextDeltaEvent(
+                        content=delta.content,
+                        response_id=delta.response_id,
                         session_id=session_ctx.session_id,
                     )
-                case model.AssistantMessageItem() as item:
-                    turn_result.assistant_message = item
-                    yield events.AssistantMessageEvent(
-                        content=item.content or "",
-                        response_id=item.response_id,
+                case message.AssistantImageDelta() as delta:
+                    yield events.AssistantImageDeltaEvent(
+                        file_path=delta.file_path,
+                        response_id=delta.response_id,
                         session_id=session_ctx.session_id,
                     )
-                case model.ResponseMetadataItem() as item:
-                    yield events.ResponseMetadataEvent(
+                case message.AssistantMessage() as msg:
+                    if msg.response_id is None and self._assistant_response_id:
+                        msg.response_id = self._assistant_response_id
+                    turn_result.assistant_message = msg
+                    for part in msg.parts:
+                        if isinstance(part, message.ToolCallPart):
+                            turn_result.tool_calls.append(
+                                ToolCallRequest(
+                                    response_id=msg.response_id,
+                                    call_id=part.call_id,
+                                    tool_name=part.tool_name,
+                                    arguments_json=part.arguments_json,
+                                )
+                            )
+                    yield events.AssistantMessageEvent(
+                        content=message.join_text_parts(msg.parts),
+                        response_id=msg.response_id,
                         session_id=session_ctx.session_id,
-                        metadata=item,
                     )
-                case model.StreamErrorItem() as item:
-                    turn_result.stream_error = item
+                    if msg.stop_reason == "aborted":
+                        yield events.InterruptEvent(session_id=session_ctx.session_id)
+                    if msg.usage:
+                        metadata = msg.usage
+                        if metadata.response_id is None:
+                            metadata.response_id = msg.response_id
+                        if not metadata.model_name:
+                            metadata.model_name = ctx.llm_client.model_name
+                        if metadata.provider is None:
+                            metadata.provider = ctx.llm_client.get_llm_config().provider_name or None
+                        yield events.ResponseMetadataEvent(
+                            session_id=session_ctx.session_id,
+                            metadata=metadata,
+                        )
+                case message.StreamErrorItem() as msg:
+                    turn_result.stream_error = msg
                     log_debug(
                         "[StreamError]",
-                        item.error,
+                        msg.error,
                         style="red",
                         debug_type=DebugType.RESPONSE,
                     )
-                case model.ToolCallStartItem() as item:
+                case message.ToolCallStartItem() as msg:
                     yield events.TurnToolCallStartEvent(
                         session_id=session_ctx.session_id,
-                        response_id=item.response_id,
-                        tool_call_id=item.call_id,
-                        tool_name=item.name,
+                        response_id=msg.response_id,
+                        tool_call_id=msg.call_id,
+                        tool_name=msg.name,
                         arguments="",
                     )
-                case model.ToolCallItem() as item:
-                    turn_result.tool_calls.append(item)
                 case _:
                     continue
 
     def _append_success_history(self, turn_result: TurnResult) -> None:
         """Persist successful turn artifacts to conversation history."""
         session_ctx = self._context.session_ctx
-        if turn_result.reasoning_items:
-            session_ctx.append_history(turn_result.reasoning_items)
         if turn_result.assistant_message:
             session_ctx.append_history([turn_result.assistant_message])
-        if turn_result.tool_calls:
-            session_ctx.append_history(turn_result.tool_calls)
         self._assistant_delta_buffer.clear()
         self._assistant_response_id = None
 
-    async def _run_tool_executor(self, tool_calls: list[model.ToolCallItem]) -> AsyncGenerator[events.Event]:
+    async def _run_tool_executor(self, tool_calls: list[ToolCallRequest]) -> AsyncGenerator[events.Event]:
         """Run tools for the turn and translate executor events to UI events."""
 
         ctx = self._context
         session_ctx = ctx.session_ctx
         with tool_context(session_ctx.file_tracker, session_ctx.todo_context):
+            resume_claims_token = current_sub_agent_resume_claims.set(set())
             executor = ToolExecutor(
                 registry=ctx.tool_registry,
                 append_history=session_ctx.append_history,
@@ -300,6 +345,7 @@ class TurnExecutor:
                         yield ui_event
             finally:
                 self._tool_executor = None
+                current_sub_agent_resume_claims.reset(resume_claims_token)
 
     def _persist_partial_assistant_on_cancel(self) -> None:
         """Persist streamed assistant text when a turn is interrupted.
@@ -311,12 +357,13 @@ class TurnExecutor:
 
         if not self._assistant_delta_buffer:
             return
-        partial_text = "".join(self._assistant_delta_buffer) + "<system interrupted by user>"
+        partial_text = "".join(self._assistant_delta_buffer) + INTERRUPT_MARKER
         if not partial_text:
             return
-        message_item = model.AssistantMessageItem(
-            content=partial_text,
+        partial_message = message.AssistantMessage(
+            parts=message.text_parts_from_str(partial_text),
             response_id=self._assistant_response_id,
+            stop_reason="aborted",
         )
-        self._context.session_ctx.append_history([message_item])
+        self._context.session_ctx.append_history([partial_message])
         self._assistant_delta_buffer.clear()