citrus-code 0.1.0__py3-none-any.whl

citrus_code/coding/tools/edit_text.py

@@ -0,0 +1,172 @@
+import unicodedata
+from bisect import bisect_right
+from dataclasses import dataclass
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class TextEdit(BaseModel):
+    model_config = ConfigDict(
+        extra="forbid",
+        strict=True,
+    )
+
+    old_text: str = Field(
+        alias="oldText",
+        min_length=1,
+        description=(
+            "Exact text to replace. It must match "
+            "a unique region of the original file."
+        ),
+    )
+
+    new_text: str = Field(
+        alias="newText",
+        description="The replacement text.",
+    )
+
+
+@dataclass(frozen=True)
+class MatchedEdit:
+    edit_index: int
+    start: int
+    end: int
+    new_text: str
+
+
+def normalize_to_lf(text: str) -> str:
+    return text.replace("\r\n", "\n").replace("\r", "\n")
+
+
+def normalize_for_match(text: str) -> str:
+    """Pi's fallback: normalize typography and trailing whitespace, not indentation."""
+    text = unicodedata.normalize("NFKC", text)
+    text = "\n".join(line.rstrip() for line in text.split("\n"))
+    replacements = {
+        **dict.fromkeys("\u2018\u2019\u201a\u201b", "'"),
+        **dict.fromkeys("\u201c\u201d\u201e\u201f", '"'),
+        **dict.fromkeys("\u2010\u2011\u2012\u2013\u2014\u2015\u2212", "-"),
+        **dict.fromkeys("\u00a0\u2002\u2003\u2004\u2005\u2006\u2007\u2008\u2009\u200a\u202f\u205f\u3000", " "),
+    }
+    return text.translate(str.maketrans(replacements))
+
+
+def apply_text_edits(
+    content: str,
+    edits: list[TextEdit],
+    path: str,
+) -> str:
+    if not edits:
+        raise ValueError(
+            "edits must contain at least one replacement"
+        )
+
+    original = normalize_to_lf(content)
+    prepared = [
+        (normalize_to_lf(edit.old_text), normalize_to_lf(edit.new_text))
+        for edit in edits
+    ]
+    use_fuzzy = any(old not in original for old, _ in prepared)
+    base = normalize_for_match(original) if use_fuzzy else original
+
+    matches: list[MatchedEdit] = []
+
+    # 所有匹配都针对同一份原始内容。
+    for edit_index, (old_text, new_text) in enumerate(prepared):
+        if use_fuzzy:
+            old_text = normalize_for_match(old_text)
+        if not old_text:
+            raise ValueError(f"edits[{edit_index}].oldText must not be empty in {path}")
+
+        start = base.find(old_text)
+
+        if start == -1:
+            raise ValueError(
+                f"Could not find edits[{edit_index}].oldText "
+                f"in {path}. The text must match exactly."
+            )
+
+        # Typography-equivalent duplicates are ambiguous even with one exact hit.
+        uniqueness_text = normalize_for_match(base)
+        uniqueness_old = normalize_for_match(old_text)
+        if not uniqueness_old:
+            uniqueness_text, uniqueness_old = base, old_text
+        first_start = uniqueness_text.find(uniqueness_old)
+        second_start = uniqueness_text.find(uniqueness_old, first_start + 1)
+
+        if second_start != -1:
+            raise ValueError(
+                f"edits[{edit_index}].oldText is not unique "
+                f"in {path}. Provide more surrounding text."
+            )
+
+        matches.append(
+            MatchedEdit(
+                edit_index=edit_index,
+                start=start,
+                end=start + len(old_text),
+                new_text=new_text,
+            )
+        )
+
+    matches.sort(key=lambda match: match.start)
+
+    # 排序之后，相邻区域就足以用于检测重叠。
+    for index in range(1, len(matches)):
+        previous = matches[index - 1]
+        current = matches[index]
+
+        if previous.end > current.start:
+            raise ValueError(
+                f"edits[{previous.edit_index}] and "
+                f"edits[{current.edit_index}] overlap "
+                f"in {path}. Merge them into one edit."
+            )
+
+    if not use_fuzzy:
+        result = original
+        # 从后往前替换，保持前面尚未处理的位置有效。
+        for match in reversed(matches):
+            result = result[:match.start] + match.new_text + result[match.end:]
+    else:
+        # Only rewrite lines touched by a normalized match. Untouched lines keep
+        # their original typography and trailing spaces.
+        original_lines = original.splitlines(keepends=True)
+        base_lines = base.splitlines(keepends=True)
+        if len(original_lines) != len(base_lines):
+            raise ValueError("Normalization changed the line structure")
+        offsets = [0]
+        for line in base_lines:
+            offsets.append(offsets[-1] + len(line))
+        groups: list[tuple[int, int, list[MatchedEdit]]] = []
+        for match in matches:
+            start_line = bisect_right(offsets, match.start) - 1
+            end_line = bisect_right(offsets, match.end - 1)
+            if groups and start_line < groups[-1][1]:
+                previous_start, previous_end, group_matches = groups[-1]
+                group_matches.append(match)
+                groups[-1] = (previous_start, max(previous_end, end_line), group_matches)
+            else:
+                groups.append((start_line, end_line, [match]))
+
+        pieces: list[str] = []
+        cursor = 0
+        for start_line, end_line, group_matches in groups:
+            pieces.extend(original_lines[cursor:start_line])
+            offset = offsets[start_line]
+            block = base[offset:offsets[end_line]]
+            for match in reversed(group_matches):
+                start, end = match.start - offset, match.end - offset
+                block = block[:start] + match.new_text + block[end:]
+            pieces.append(block)
+            cursor = end_line
+        pieces.extend(original_lines[cursor:])
+        result = "".join(pieces)
+
+    if result == original:
+        raise ValueError(
+            f"No changes made to {path}. "
+            "The replacements produced identical content."
+        )
+
+    return result

citrus_code/coding/tools/factory.py

@@ -0,0 +1,57 @@
+from collections.abc import Callable, Iterable
+from pathlib import Path
+
+from ...agent.types import AgentTool
+from .bash import BashTool
+from .edit import EditTool
+from .find import FindTool
+from .grep import GrepTool
+from .ls import LsTool
+from .powershell import PowerShellTool
+from .read import ReadTool
+from .write import WriteTool
+
+
+DEFAULT_CODING_TOOLS = ("read", "bash", "edit", "write")
+READ_ONLY_TOOLS = ("read", "grep", "find", "ls")
+
+_FACTORIES: dict[str, Callable[[Path], AgentTool]] = {
+    "read": ReadTool,
+    "bash": BashTool,
+    "powershell": PowerShellTool,
+    "edit": EditTool,
+    "write": WriteTool,
+    "grep": GrepTool,
+    "find": FindTool,
+    "ls": LsTool,
+}
+
+
+def create_tool(name: str, cwd: Path) -> AgentTool:
+    factory = _FACTORIES.get(name)
+    if factory is None:
+        raise ValueError(f"Unknown built-in tool: {name}")
+    return factory(cwd)
+
+
+def create_tools(cwd: Path, names: Iterable[str]) -> list[AgentTool]:
+    tools: list[AgentTool] = []
+    seen: set[str] = set()
+    for name in names:
+        if name in seen:
+            raise ValueError(f"Duplicate built-in tool: {name}")
+        seen.add(name)
+        tools.append(create_tool(name, cwd))
+    return tools
+
+
+def create_coding_tools(cwd: Path) -> list[AgentTool]:
+    return create_tools(cwd, DEFAULT_CODING_TOOLS)
+
+
+def create_read_only_tools(cwd: Path) -> list[AgentTool]:
+    return create_tools(cwd, READ_ONLY_TOOLS)
+
+
+def create_all_tools(cwd: Path) -> dict[str, AgentTool]:
+    return {name: factory(cwd) for name, factory in _FACTORIES.items()}

citrus_code/coding/tools/find.py

@@ -0,0 +1,86 @@
+import os
+import sys
+from contextlib import aclosing
+from pathlib import Path
+
+from pydantic import BaseModel, ConfigDict, Field, JsonValue
+
+from ...agent.types import ToolExecutionResult
+from ...ai.types import TextContent, ToolDefinition
+from .paths import resolve_to_cwd
+from .processes import command_lines, find_executable
+from .truncate import truncate_text
+
+
+class FindArguments(BaseModel):
+    model_config = ConfigDict(extra="forbid", strict=True)
+
+    pattern: str = Field(min_length=1, description="Glob pattern, e.g. '*.py' or 'src/**/*.py'.")
+    path: str = Field(default=".", description="Directory to search.")
+    limit: int = Field(default=1000, ge=1, description="Maximum results to return.")
+
+
+class FindTool:
+    name: str = "find"
+    description: str = (
+        "Find paths by glob pattern, respecting .gitignore and including hidden "
+        "files. Results are relative to the search directory. Defaults to 1000 "
+        "results, with a 50KB output cap."
+    )
+
+    def __init__(self, cwd: Path) -> None:
+        self._cwd = cwd.resolve()
+
+    def get_definition(self) -> ToolDefinition:
+        return ToolDefinition(
+            name=self.name,
+            description=self.description,
+            parameters=FindArguments.model_json_schema(),
+        )
+
+    async def execute(
+        self, tool_call_id: str, arguments: dict[str, JsonValue]
+    ) -> ToolExecutionResult:
+        args = FindArguments.model_validate(arguments)
+        root = resolve_to_cwd(args.path, self._cwd)
+        if not root.is_dir():
+            raise NotADirectoryError(f"Not a directory: {root}")
+        command = [find_executable("fd"), "--glob", "--color=never", "--hidden", "--max-results", str(args.limit)]
+        if not any((parent / ".git").exists() for parent in (root, *root.parents)):
+            command.append("--no-require-git")
+        pattern = args.pattern
+        if "/" in pattern:
+            command.append("--full-path")
+            if not pattern.startswith(("/", "**/")) and pattern != "**":
+                pattern = "**/" + pattern
+            if os.name == "nt":
+                pattern = pattern.replace("/", r"[/\\]")
+        command.extend(["--", pattern, str(root)])
+        results: list[str] = []
+        async with aclosing(command_lines(command, self._cwd)) as stream:
+            async for line in stream:
+                trailing = line.endswith(("/", "\\"))
+                path = Path(line)
+                relative = path.relative_to(root) if path.is_absolute() else path
+                label = relative.as_posix()
+                if trailing and not label.endswith("/"):
+                    label += "/"
+                results.append(label)
+                if len(results) >= args.limit:
+                    break
+
+        if not results:
+            return ToolExecutionResult(content=[TextContent(text="No files found matching pattern")])
+        truncation = truncate_text("\n".join(results), max_lines=sys.maxsize)
+        text = truncation.content
+        details: dict[str, JsonValue] = {}
+        notices: list[str] = []
+        if len(results) >= args.limit:
+            details["resultLimitReached"] = args.limit
+            notices.append(f"{args.limit} results limit reached. Refine the pattern or increase limit")
+        if truncation.truncated:
+            details["truncation"] = truncation.model_dump(mode="json", exclude={"content"})
+            notices.append("50KB output limit reached")
+        if notices:
+            text += "\n\n[" + ". ".join(notices) + "]"
+        return ToolExecutionResult(content=[TextContent(text=text)], details=details)

citrus_code/coding/tools/grep.py

@@ -0,0 +1,151 @@
+import asyncio
+import base64
+import os
+import sys
+from contextlib import aclosing
+from pathlib import Path
+
+from pydantic import BaseModel, ConfigDict, Field, JsonValue
+
+from ...agent.types import ToolExecutionResult
+from ...ai.types import TextContent, ToolDefinition
+from .paths import resolve_to_cwd
+from .processes import command_lines, find_executable
+from .truncate import DEFAULT_MAX_BYTES, GREP_MAX_LINE_LENGTH, truncate_text
+
+
+class GrepArguments(BaseModel):
+    model_config = ConfigDict(extra="forbid", strict=True)
+
+    pattern: str = Field(description="Search pattern: regex, or a literal string when literal is true.")
+    path: str = Field(default=".", description="File or directory to search.")
+    glob: str | None = Field(default=None, description="Filter files, e.g. '*.py'.")
+    ignore_case: bool = Field(default=False, alias="ignoreCase")
+    literal: bool = False
+    context: int = Field(default=0, ge=0, description="Context lines before and after each match.")
+    limit: int = Field(default=100, ge=1, description="Maximum matching lines to return.")
+
+
+class _RgText(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+    text: str | None = None
+    bytes: str | None = None
+
+    def value(self) -> str:
+        if self.text is not None:
+            return self.text
+        if self.bytes is not None:
+            return os.fsdecode(base64.b64decode(self.bytes))
+        raise ValueError("ripgrep returned an empty text field")
+
+
+class _RgData(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+    path: _RgText | None = None
+    lines: _RgText | None = None
+    line_number: int | None = None
+
+
+class _RgEvent(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+    type: str
+    data: _RgData | None = None
+
+
+class GrepTool:
+    name: str = "grep"
+    description: str = (
+        "Search file contents using regex or literal text. Respects .gitignore, "
+        "includes hidden files, and returns paths with line numbers. Default limit "
+        "is 100 matching lines, with a 50KB output cap and 500 characters per line."
+    )
+
+    def __init__(self, cwd: Path) -> None:
+        self._cwd = cwd.resolve()
+
+    def get_definition(self) -> ToolDefinition:
+        return ToolDefinition(
+            name=self.name,
+            description=self.description,
+            parameters=GrepArguments.model_json_schema(by_alias=True),
+        )
+
+    async def execute(
+        self, tool_call_id: str, arguments: dict[str, JsonValue]
+    ) -> ToolExecutionResult:
+        args = GrepArguments.model_validate(arguments)
+        root = resolve_to_cwd(args.path, self._cwd)
+        if not root.exists():
+            raise FileNotFoundError(f"Path not found: {root}")
+        command = [find_executable("rg"), "--json", "--line-number", "--color=never", "--hidden", "--sort=path"]
+        if args.ignore_case:
+            command.append("--ignore-case")
+        if args.literal:
+            command.append("--fixed-strings")
+        if args.glob is not None:
+            command.extend(["--glob", args.glob])
+        command.extend(["--", args.pattern, str(root)])
+
+        matches: list[tuple[Path, int, str]] = []
+        async with aclosing(command_lines(command, self._cwd, success_codes=(0, 1))) as stream:
+            async for raw in stream:
+                event = _RgEvent.model_validate_json(raw)
+                if event.type != "match":
+                    continue
+                data = event.data
+                if data is None or data.path is None or data.lines is None or data.line_number is None:
+                    raise ValueError("ripgrep returned an invalid match")
+                matches.append((Path(data.path.value()), data.line_number, data.lines.value()))
+                if len(matches) >= args.limit:
+                    break
+
+        if not matches:
+            return ToolExecutionResult(content=[TextContent(text="No matches found")])
+
+        output: list[str] = []
+        cache: dict[Path, list[str]] = {}
+        lines_truncated = False
+        output_bytes = 0
+        for path, line_number, match_text in matches:
+            label = path.relative_to(root).as_posix() if root.is_dir() else path.name
+            if args.context:
+                if path not in cache:
+                    text = await asyncio.to_thread(path.read_text, encoding="utf-8", errors="replace")
+                    cache[path] = text.splitlines()
+                lines = cache[path]
+                start = max(1, line_number - args.context)
+                end = min(len(lines), line_number + args.context)
+            else:
+                lines = [match_text.rstrip("\r\n")]
+                start = end = line_number
+
+            for number in range(start, end + 1):
+                text = lines[number - 1] if args.context else lines[0]
+                if len(text) > GREP_MAX_LINE_LENGTH:
+                    text = text[:GREP_MAX_LINE_LENGTH] + "... [truncated]"
+                    lines_truncated = True
+                separator = ":" if number == line_number else "-"
+                row = f"{label}{separator}{number}{separator} {text}"
+                output.append(row)
+                output_bytes += len(row.encode("utf-8")) + 1
+                if output_bytes > DEFAULT_MAX_BYTES:
+                    break
+            if output_bytes > DEFAULT_MAX_BYTES:
+                break
+
+        truncation = truncate_text("\n".join(output), max_lines=sys.maxsize)
+        text = truncation.content
+        details: dict[str, JsonValue] = {}
+        notices: list[str] = []
+        if len(matches) >= args.limit:
+            details["matchLimitReached"] = args.limit
+            notices.append(f"{args.limit} matches limit reached. Refine the pattern or increase limit")
+        if truncation.truncated:
+            details["truncation"] = truncation.model_dump(mode="json", exclude={"content"})
+            notices.append("50KB output limit reached")
+        if lines_truncated:
+            details["linesTruncated"] = True
+            notices.append("Some lines were truncated to 500 characters. Use read for full lines")
+        if notices:
+            text += "\n\n[" + ". ".join(notices) + "]"
+        return ToolExecutionResult(content=[TextContent(text=text)], details=details)

citrus_code/coding/tools/ls.py

@@ -0,0 +1,61 @@
+import asyncio
+import sys
+from pathlib import Path
+
+from pydantic import BaseModel, ConfigDict, Field, JsonValue
+
+from ...agent.types import ToolExecutionResult
+from ...ai.types import TextContent, ToolDefinition
+from .paths import resolve_to_cwd
+from .truncate import truncate_text
+
+
+class LsArguments(BaseModel):
+    model_config = ConfigDict(extra="forbid", strict=True)
+
+    path: str = Field(default=".", description="Directory to list.")
+    limit: int = Field(default=500, ge=1, description="Maximum entries to return.")
+
+
+class LsTool:
+    name: str = "ls"
+    description: str = (
+        "List directory entries alphabetically, including dotfiles. Directories "
+        "have a '/' suffix. Defaults to 500 entries, with a 50KB output cap."
+    )
+
+    def __init__(self, cwd: Path) -> None:
+        self._cwd = cwd.resolve()
+
+    def get_definition(self) -> ToolDefinition:
+        return ToolDefinition(
+            name=self.name,
+            description=self.description,
+            parameters=LsArguments.model_json_schema(),
+        )
+
+    async def execute(
+        self, tool_call_id: str, arguments: dict[str, JsonValue]
+    ) -> ToolExecutionResult:
+        args = LsArguments.model_validate(arguments)
+        path = resolve_to_cwd(args.path, self._cwd)
+        return await asyncio.to_thread(self._list, path, args.limit)
+
+    def _list(self, path: Path, limit: int) -> ToolExecutionResult:
+        if not path.is_dir():
+            raise NotADirectoryError(f"Not a directory: {path}")
+        entries = sorted(path.iterdir(), key=lambda entry: (entry.name.casefold(), entry.name))
+        labels = [entry.name + ("/" if entry.is_dir() else "") for entry in entries[:limit]]
+        truncation = truncate_text("\n".join(labels), max_lines=sys.maxsize)
+        text = truncation.content or "(empty directory)"
+        details: dict[str, JsonValue] = {}
+        notices: list[str] = []
+        if len(entries) > limit:
+            details["entryLimitReached"] = limit
+            notices.append(f"{limit} entries limit reached. Increase limit for more")
+        if truncation.truncated:
+            details["truncation"] = truncation.model_dump(mode="json", exclude={"content"})
+            notices.append("50KB output limit reached")
+        if notices:
+            text += "\n\n[" + ". ".join(notices) + "]"
+        return ToolExecutionResult(content=[TextContent(text=text)], details=details)

citrus_code/coding/tools/mutation.py

@@ -0,0 +1,39 @@
+import os
+from _thread import LockType
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from threading import Lock
+
+
+@dataclass
+class _Entry:
+    lock: LockType = field(default_factory=Lock)
+    users: int = 0
+
+
+_entries: dict[str, _Entry] = {}
+_registry_lock = Lock()
+
+
+@contextmanager
+def file_mutation(path: Path) -> Iterator[None]:
+    """Serialize writes to one path inside the actual worker thread.
+
+    Cancelling the awaiting coroutine cannot release this lock while the
+    filesystem operation is still running. Different paths use different locks.
+    """
+    key = os.path.normcase(str(path.resolve()))
+    with _registry_lock:
+        entry = _entries.setdefault(key, _Entry())
+        entry.users += 1
+
+    try:
+        with entry.lock:
+            yield
+    finally:
+        with _registry_lock:
+            entry.users -= 1
+            if entry.users == 0:
+                del _entries[key]

citrus_code/coding/tools/paths.py

@@ -0,0 +1,11 @@
+from pathlib import Path
+
+
+def resolve_to_cwd(value: str, cwd: Path) -> Path:
+    """Resolve model-supplied paths against the tool's working directory."""
+    if value.startswith("@"):
+        value = value[1:]
+    path = Path(value).expanduser()
+    if not path.is_absolute():
+        path = cwd / path
+    return path.resolve()

citrus_code/coding/tools/powershell.py

@@ -0,0 +1,5 @@
+from .shell import ShellTool
+
+
+class PowerShellTool(ShellTool):
+    name: str = "powershell"

citrus_code/coding/tools/processes.py

@@ -0,0 +1,103 @@
+import asyncio
+import os
+import shutil
+import signal
+import subprocess
+from collections.abc import AsyncIterator
+from pathlib import Path
+from typing import BinaryIO
+
+
+def find_executable(name: str) -> str:
+    """Use PATH or pi's already-installed helpers; never download executables."""
+    found = shutil.which(name)
+    if found:
+        return found
+    suffix = ".exe" if os.name == "nt" else ""
+    bundled = Path.home() / ".pi" / "agent" / "bin" / f"{name}{suffix}"
+    if bundled.is_file():
+        return str(bundled)
+    raise FileNotFoundError(f"{name} is required. Install it and add it to PATH.")
+
+
+async def spawn_process(
+    command: list[str],
+    cwd: Path,
+    *,
+    merge_stderr: bool = False,
+    output: BinaryIO | None = None,
+) -> asyncio.subprocess.Process:
+    return await asyncio.create_subprocess_exec(
+        *command,
+        cwd=cwd,
+        stdin=asyncio.subprocess.DEVNULL,
+        stdout=output if output is not None else asyncio.subprocess.PIPE,
+        stderr=(asyncio.subprocess.STDOUT if merge_stderr else asyncio.subprocess.PIPE),
+        start_new_session=os.name != "nt",
+        creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0,
+        limit=8 * 1024 * 1024,
+    )
+
+
+async def stop_process(process: asyncio.subprocess.Process) -> None:
+    """Stop a running command and its process tree on cancellation or timeout."""
+    if process.returncode is not None:
+        return
+    if os.name == "nt":
+        killer = await asyncio.create_subprocess_exec(
+            "taskkill", "/PID", str(process.pid), "/T", "/F",
+            stdout=asyncio.subprocess.DEVNULL,
+            stderr=asyncio.subprocess.DEVNULL,
+            creationflags=subprocess.CREATE_NO_WINDOW,
+        )
+        await killer.wait()
+        if process.returncode is None:
+            try:
+                process.kill()
+            except ProcessLookupError:
+                pass
+    else:
+        try:
+            os.killpg(process.pid, signal.SIGKILL)
+        except ProcessLookupError:
+            pass
+
+
+async def command_lines(
+    command: list[str],
+    cwd: Path,
+    *,
+    success_codes: tuple[int, ...] = (0,),
+) -> AsyncIterator[str]:
+    """Stream search output; closing the generator also reaps its child process."""
+    process = await spawn_process(command, cwd)
+    stderr = bytearray()
+
+    async def drain_errors() -> None:
+        if process.stderr is None:
+            return
+        while chunk := await process.stderr.read(8192):
+            if len(stderr) < 50 * 1024:
+                stderr.extend(chunk[:50 * 1024 - len(stderr)])
+
+    error_task = asyncio.create_task(drain_errors())
+    try:
+        if process.stdout is None:
+            raise RuntimeError("The process has no stdout pipe")
+        async for line in process.stdout:
+            yield line.decode("utf-8", errors="replace").rstrip("\r\n")
+        code = await process.wait()
+        await error_task
+        if code not in success_codes:
+            message = stderr.decode("utf-8", errors="replace").strip()
+            raise RuntimeError(message or f"Command exited with code {code}")
+    finally:
+        await stop_process(process)
+        if process.stdout is not None:
+            # A killed child may leave unread bytes in a paused pipe. Drain them
+            # before waiting for asyncio's process transport to close.
+            await process.stdout.read()
+        await process.wait()
+        if not error_task.done():
+            error_task.cancel()
+        await asyncio.gather(error_task, return_exceptions=True)