klaude-code 1.2.6__py3-none-any.whl → 1.8.0__py3-none-any.whl

klaude_code/protocol/sub_agent/__init__.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from klaude_code.protocol import tools
+
+if TYPE_CHECKING:
+    from klaude_code.protocol import model
+
+AvailabilityPredicate = Callable[[str], bool]
+PromptBuilder = Callable[[dict[str, Any]], str]
+
+
+@dataclass
+class SubAgentResult:
+    task_result: str
+    session_id: str
+    error: bool = False
+    task_metadata: model.TaskMetadata | None = None
+
+
+def _default_prompt_builder(args: dict[str, Any]) -> str:
+    """Default prompt builder that just returns the 'prompt' field."""
+    return args.get("prompt", "")
+
+
+@dataclass(frozen=True)
+class SubAgentProfile:
+    """Metadata describing a sub agent and how it integrates with the system.
+
+    This dataclass contains all the information needed to:
+    1. Register the sub agent with the system
+    2. Generate the tool schema for the main agent
+    3. Build the prompt for the sub agent
+    """
+
+    # Identity - single name used for type, tool_name, config_key, and prompt_key
+    name: str  # e.g., "Task", "Oracle", "Explore"
+
+    # Tool schema
+    description: str  # Tool description shown to the main agent
+    parameters: dict[str, Any] = field(
+        default_factory=lambda: dict[str, Any](), hash=False
+    )  # JSON Schema for tool parameters
+
+    # System prompt
+    prompt_file: str = ""  # Resource file path relative to core package (e.g., "prompts/prompt-sub-agent.md")
+
+    # Sub agent configuration
+    tool_set: tuple[str, ...] = ()  # Tools available to this sub agent
+    prompt_builder: PromptBuilder = _default_prompt_builder  # Builds the sub agent prompt from tool arguments
+
+    # UI display
+    active_form: str = ""  # Active form for spinner status (e.g., "Tasking", "Exploring")
+
+    # Availability
+    enabled_by_default: bool = True
+    show_in_main_agent: bool = True
+    target_model_filter: AvailabilityPredicate | None = None
+
+    # Structured output support: specifies which parameter in the tool schema contains the output schema
+    output_schema_arg: str | None = None
+
+    def enabled_for_model(self, model_name: str | None) -> bool:
+        if not self.enabled_by_default:
+            return False
+        if model_name is None or self.target_model_filter is None:
+            return True
+        return self.target_model_filter(model_name)
+
+
+_PROFILES: dict[str, SubAgentProfile] = {}
+
+
+def register_sub_agent(profile: SubAgentProfile) -> None:
+    if profile.name in _PROFILES:
+        raise ValueError(f"Duplicate sub agent profile: {profile.name}")
+    _PROFILES[profile.name] = profile
+
+
+def get_sub_agent_profile(sub_agent_type: tools.SubAgentType) -> SubAgentProfile:
+    try:
+        return _PROFILES[sub_agent_type]
+    except KeyError as exc:
+        raise KeyError(f"Unknown sub agent type: {sub_agent_type}") from exc
+
+
+def iter_sub_agent_profiles(enabled_only: bool = False, model_name: str | None = None) -> list[SubAgentProfile]:
+    profiles = list(_PROFILES.values())
+    if not enabled_only:
+        return profiles
+    return [p for p in profiles if p.enabled_for_model(model_name)]
+
+
+def get_sub_agent_profile_by_tool(tool_name: str) -> SubAgentProfile | None:
+    return _PROFILES.get(tool_name)
+
+
+def is_sub_agent_tool(tool_name: str) -> bool:
+    return tool_name in _PROFILES
+
+
+def sub_agent_tool_names(enabled_only: bool = False, model_name: str | None = None) -> list[str]:
+    return [
+        profile.name
+        for profile in iter_sub_agent_profiles(enabled_only=enabled_only, model_name=model_name)
+        if profile.show_in_main_agent
+    ]
+
+
+# Import sub-agent modules to trigger registration
+from klaude_code.protocol.sub_agent import explore as explore  # noqa: E402
+from klaude_code.protocol.sub_agent import oracle as oracle  # noqa: E402
+from klaude_code.protocol.sub_agent import task as task  # noqa: E402
+from klaude_code.protocol.sub_agent import web as web  # noqa: E402

klaude_code/protocol/sub_agent/explore.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from typing import Any
+
+from klaude_code.protocol import tools
+from klaude_code.protocol.sub_agent import SubAgentProfile, register_sub_agent
+
+EXPLORE_DESCRIPTION = """\
+Spin up a fast agent specialized for exploring codebases. Use this when you need to quickly find files by patterns (eg. "src/components/**/*.tsx"), \
+search code for keywords (eg. "API endpoints"), or answer questions about the codebase (eg. "how do API endpoints work?")\
+When calling this agent, specify the desired thoroughness level: "quick" for basic searches, "medium" for moderate exploration, or "very thorough" for comprehensive analysis across multiple locations and naming conventions.
+Always spawn multiple search agents in parallel to maximise speed.
+
+Structured output:
+- Provide an `output_format` (JSON Schema) parameter for structured data back from the sub-agent
+- 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"]}`\
+"""
+
+EXPLORE_PARAMETERS = {
+    "type": "object",
+    "properties": {
+        "description": {
+            "type": "string",
+            "description": "Short (3-5 words) label for the exploration goal",
+        },
+        "prompt": {
+            "type": "string",
+            "description": "The task for the agent to perform",
+        },
+        "thoroughness": {
+            "type": "string",
+            "enum": ["quick", "medium", "very thorough"],
+            "description": "Controls how deep the sub-agent should search the repo",
+        },
+        "output_format": {
+            "type": "object",
+            "description": "Optional JSON Schema for sub-agent structured output",
+        },
+    },
+    "required": ["description", "prompt"],
+    "additionalProperties": False,
+}
+
+
+def _explore_prompt_builder(args: dict[str, Any]) -> str:
+    """Build the Explore prompt from tool arguments."""
+    prompt = args.get("prompt", "").strip()
+    thoroughness = args.get("thoroughness", "medium")
+    return f"{prompt}\nthoroughness: {thoroughness}"
+
+
+register_sub_agent(
+    SubAgentProfile(
+        name="Explore",
+        description=EXPLORE_DESCRIPTION,
+        parameters=EXPLORE_PARAMETERS,
+        prompt_file="prompts/prompt-sub-agent-explore.md",
+        tool_set=(tools.BASH, tools.READ),
+        prompt_builder=_explore_prompt_builder,
+        active_form="Exploring",
+        output_schema_arg="output_format",
+    )
+)

klaude_code/protocol/sub_agent/oracle.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from typing import Any
+
+from klaude_code.protocol import tools
+from klaude_code.protocol.sub_agent import SubAgentProfile, register_sub_agent
+
+ORACLE_DESCRIPTION = """\
+Consult the Oracle - an AI advisor powered by OpenAI's premium reasoning model that can plan, review, and provide expert guidance.
+
+The Oracle has access to the following tools: Read, Bash.
+
+The Oracle acts as your senior engineering advisor and can help with:
+
+WHEN TO USE THE ORACLE:
+- Code reviews and architecture feedback
+- Finding a bug in multiple files
+- Planning complex implementations or refactoring
+- Analyzing code quality and suggesting improvements
+- Answering complex technical questions that require deep reasoning
+
+WHEN NOT TO USE THE ORACLE:
+- Simple file reading or searching tasks (use Read or Grep directly)
+- Codebase searches (use Task)
+- Basic code modifications and when you need to execute code changes (do it yourself or use Task)
+
+USAGE GUIDELINES:
+1. Be specific about what you want the Oracle to review, plan, or debug
+2. Provide relevant context about what you're trying to achieve. If you know that any files are involved, list them and they will be attached.
+
+
+EXAMPLES:
+- "Review the authentication system architecture and suggest improvements"
+- "Plan the implementation of real-time collaboration features"
+- "Analyze the performance bottlenecks in the data processing pipeline"
+- "Review this API design and suggest better patterns"\
+"""
+
+ORACLE_PARAMETERS = {
+    "properties": {
+        "context": {
+            "description": "Optional context about the current situation, what you've tried, or background information that would help the Oracle provide better guidance.",
+            "type": "string",
+        },
+        "files": {
+            "description": "Optional list of specific file paths (text files, images) that the Oracle should examine as part of its analysis. These files will be attached to the Oracle input.",
+            "items": {"type": "string"},
+            "type": "array",
+        },
+        "task": {
+            "description": "The task or question you want the Oracle to help with. Be specific about what kind of guidance, review, or planning you need.",
+            "type": "string",
+        },
+        "description": {
+            "description": "A short (3-5 word) description of the task",
+            "type": "string",
+        },
+    },
+    "required": ["task", "description"],
+    "type": "object",
+}
+
+
+def _oracle_prompt_builder(args: dict[str, Any]) -> str:
+    """Build the Oracle prompt from tool arguments."""
+    context = args.get("context", "")
+    task = args.get("task", "")
+    files = args.get("files", [])
+
+    prompt = f"""\
+Context: {context}
+Task: {task}\
+"""
+    if files:
+        files_str = "\n".join(f"@{file}" for file in files)
+        prompt += f"\nRelated files to review:\n{files_str}"
+    return prompt
+
+
+register_sub_agent(
+    SubAgentProfile(
+        name="Oracle",
+        description=ORACLE_DESCRIPTION,
+        parameters=ORACLE_PARAMETERS,
+        prompt_file="prompts/prompt-sub-agent-oracle.md",
+        tool_set=(tools.READ, tools.BASH),
+        prompt_builder=_oracle_prompt_builder,
+        active_form="Consulting Oracle",
+        target_model_filter=lambda model: ("gpt-5" not in model) and ("gemini-3" not in model),
+    )
+)

klaude_code/protocol/sub_agent/task.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from klaude_code.protocol import tools
+from klaude_code.protocol.sub_agent import SubAgentProfile, register_sub_agent
+
+TASK_DESCRIPTION = """\
+Launch a new agent to handle complex, multi-step tasks autonomously. \
+
+When NOT to use the Task tool:
+- If you want to read a specific file path, use the Read or Bash tool for `rg` instead of the Task tool, to find the match more quickly
+- If you are searching for a specific class definition like "class Foo", use the Bash tool for `rg` instead, to find the match more quickly
+- If you are searching for code within a specific file or set of 2-3 files, use the Read tool instead of the Task tool, to find the match more quickly
+- Other tasks that are not related to the agent descriptions above
+
+Usage notes:
+- Launch multiple agents concurrently whenever possible, to maximize performance; to do that, use a single message with multiple tool uses
+- When the agent is done, it will return a single message back to you. The result returned by the agent is not visible to the user. To show the user the result, you should send a text message back to the user with a concise summary of the result.
+- Each agent invocation is stateless. You will not be able to send additional messages to the agent, nor will the agent be able to communicate with you outside of its final report. Therefore, your prompt should contain a highly detailed task description for the agent to perform autonomously and you should specify exactly what information the agent should return back to you in its final and only message to you.
+- The agent's outputs should generally be trusted
+- 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
+- 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.
+
+Structured output:
+- Provide an `output_format` (JSON Schema) parameter for structured data back from the agent
+- 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"]}`\
+"""
+
+TASK_PARAMETERS = {
+    "type": "object",
+    "properties": {
+        "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_format": {
+            "type": "object",
+            "description": (
+                "Optional JSON Schema for structured output, better with examples in argument descriptions."
+            ),
+        },
+    },
+    "required": ["description", "prompt"],
+    "additionalProperties": False,
+}
+
+register_sub_agent(
+    SubAgentProfile(
+        name="Task",
+        description=TASK_DESCRIPTION,
+        parameters=TASK_PARAMETERS,
+        prompt_file="prompts/prompt-sub-agent.md",
+        tool_set=(tools.BASH, tools.READ, tools.EDIT, tools.WRITE, tools.MOVE),
+        active_form="Tasking",
+        output_schema_arg="output_format",
+    )
+)

klaude_code/protocol/sub_agent/web.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from typing import Any
+
+from klaude_code.protocol import tools
+from klaude_code.protocol.sub_agent import SubAgentProfile, register_sub_agent
+
+WEB_AGENT_DESCRIPTION = """\
+Launch a sub-agent to search the web, fetch pages, and analyze content. Use this for:
+- Accessing up-to-date information beyond your knowledge cutoff (current events, recent releases, latest docs)
+- Researching topics, news, APIs, or technical references
+- Fetching and analyzing specific URLs
+- Gathering comprehensive information from multiple web sources
+
+Capabilities:
+- Search the web to find relevant pages (no URL required)
+- Fetch and parse web pages (HTML-to-Markdown conversion)
+- Follow links across multiple pages autonomously
+- Aggregate findings from multiple sources
+
+How to use:
+- Write a clear prompt describing what information you need - the agent will search and fetch as needed
+- Account for "Today's date" in <env>. For example, if <env> says "Today's date: 2025-07-01", and the user wants the latest docs, do not use 2024 in the search query. Use 2025.
+- Optionally provide a `url` if you already know the target page
+- Use `output_format` (JSON Schema) to get structured data back from the agent
+
+What you receive:
+- The agent returns a text response summarizing its findings
+- With `output_format`, you receive structured JSON matching your schema
+- The response is the agent's analysis, not raw web content
+- Web content is saved to local files (paths included in Sources) - read them directly if you need full content\
+"""
+
+WEB_AGENT_PARAMETERS = {
+    "type": "object",
+    "properties": {
+        "description": {
+            "type": "string",
+            "description": "A short (3-5 word) description of the task",
+        },
+        "url": {
+            "type": "string",
+            "description": "The URL to fetch and analyze. If not provided, the agent will search the web first",
+        },
+        "prompt": {
+            "type": "string",
+            "description": "Instructions for searching, analyzing, or extracting content from the web page",
+        },
+        "output_format": {
+            "type": "object",
+            "description": "Optional JSON Schema for sub-agent structured output",
+        },
+    },
+    "required": ["description", "prompt"],
+    "additionalProperties": False,
+}
+
+
+def _web_agent_prompt_builder(args: dict[str, Any]) -> str:
+    """Build the WebAgent prompt from tool arguments."""
+    url = args.get("url", "")
+    prompt = args.get("prompt", "")
+    if url:
+        return f"URL to fetch: {url}\nTask: {prompt}"
+    return prompt
+
+
+register_sub_agent(
+    SubAgentProfile(
+        name="WebAgent",
+        description=WEB_AGENT_DESCRIPTION,
+        parameters=WEB_AGENT_PARAMETERS,
+        prompt_file="prompts/prompt-sub-agent-web.md",
+        tool_set=(tools.BASH, tools.READ, tools.WEB_FETCH, tools.WEB_SEARCH, tools.WRITE),
+        prompt_builder=_web_agent_prompt_builder,
+        active_form="Surfing",
+        output_schema_arg="output_format",
+    )
+)

klaude_code/protocol/tools.py

@@ -1,15 +1,17 @@
 BASH = "Bash"
 APPLY_PATCH = "apply_patch"
+MOVE = "Move"
 EDIT = "Edit"
-MULTI_EDIT = "MultiEdit"
+
 READ = "Read"
 WRITE = "Write"
 TODO_WRITE = "TodoWrite"
 UPDATE_PLAN = "update_plan"
 SKILL = "Skill"
 MERMAID = "Mermaid"
-MEMORY = "Memory"
 WEB_FETCH = "WebFetch"
+WEB_SEARCH = "WebSearch"
+REPORT_BACK = "report_back"
 
 # SubAgentType is just a string alias now; agent types are defined via SubAgentProfile
 SubAgentType = str

klaude_code/session/__init__.py

@@ -1,4 +1,4 @@
-from .selector import resume_select_session
+from .selector import SessionSelectOption, build_session_select_options, format_user_messages_display
 from .session import Session
 
-__all__ = ["Session", "resume_select_session"]
+__all__ = ["Session", "SessionSelectOption", "build_session_select_options", "format_user_messages_display"]

klaude_code/session/codec.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import json
+from typing import Any, TypeGuard, cast, get_args
+
+from pydantic import BaseModel
+
+from klaude_code.protocol import model
+
+
+def _is_basemodel_subclass(tp: object) -> TypeGuard[type[BaseModel]]:
+    return isinstance(tp, type) and issubclass(tp, BaseModel)
+
+
+def _flatten_union(tp: object) -> list[object]:
+    args = list(get_args(tp))
+    if not args:
+        return [tp]
+    flattened: list[object] = []
+    for arg in args:
+        flattened.extend(_flatten_union(arg))
+    return flattened
+
+
+def _build_type_registry() -> dict[str, type[BaseModel]]:
+    registry: dict[str, type[BaseModel]] = {}
+    for tp in _flatten_union(model.ConversationItem):
+        if not _is_basemodel_subclass(tp):
+            continue
+        registry[tp.__name__] = tp
+    return registry
+
+
+_CONVERSATION_ITEM_TYPES: dict[str, type[BaseModel]] = _build_type_registry()
+
+
+def encode_conversation_item(item: model.ConversationItem) -> dict[str, Any]:
+    return {"type": item.__class__.__name__, "data": item.model_dump(mode="json")}
+
+
+def decode_conversation_item(obj: dict[str, Any]) -> model.ConversationItem | None:
+    t = obj.get("type")
+    data = obj.get("data", {})
+    if not isinstance(t, str) or not isinstance(data, dict):
+        return None
+    cls = _CONVERSATION_ITEM_TYPES.get(t)
+    if cls is None:
+        return None
+    try:
+        item = cls(**data)
+    except TypeError:
+        return None
+    # pyright: ignore[reportReturnType]
+    return item  # type: ignore[return-value]
+
+
+def encode_jsonl_line(item: model.ConversationItem) -> str:
+    return json.dumps(encode_conversation_item(item), ensure_ascii=False) + "\n"
+
+
+def decode_jsonl_line(line: str) -> model.ConversationItem | None:
+    line = line.strip()
+    if not line:
+        return None
+    try:
+        obj = json.loads(line)
+    except json.JSONDecodeError:
+        return None
+    if not isinstance(obj, dict):
+        return None
+    return decode_conversation_item(cast(dict[str, Any], obj))