citrus-code 0.1.0__py3-none-any.whl

citrus_code/coding/session/manager.py

@@ -0,0 +1,200 @@
+"""Session tree and context reconstruction with optional JSONL persistence."""
+
+import time
+from pathlib import Path
+from threading import RLock
+from uuid import uuid4
+
+from ...ai.types import Message, UserMessage
+from .projection import project_context_entries
+from .storage import JsonlSessionStore, SessionFormatError
+from .types import (
+    ModelChangeEntry,
+    CompactionEntry,
+    ContextMessageEntry,
+    SessionContext,
+    SessionCursorEntry,
+    SessionEntry,
+    SessionHeader,
+    SessionLogEntry,
+    SessionMessageEntry,
+)
+
+
+class SessionManager:
+    def __init__(self, cwd: Path) -> None:
+        self._header = SessionHeader(
+            id=uuid4().hex,
+            timestamp=int(time.time() * 1000),
+            cwd=str(cwd.expanduser().resolve()),
+        )
+        self._entries: list[SessionEntry] = []
+        self._records: list[SessionLogEntry] = []
+        self._by_id: dict[str, SessionEntry] = {}
+        self._leaf_id: str | None = None
+        self._store: JsonlSessionStore | None = None
+        self._lock = RLock()
+
+    @classmethod
+    def create(cls, cwd: Path, file_path: Path) -> "SessionManager":
+        session = cls(cwd)
+        session.save(file_path)
+        return session
+
+    @classmethod
+    def load(cls, file_path: Path) -> "SessionManager":
+        store, header, records = JsonlSessionStore.load(file_path)
+        session = cls(Path(header.cwd))
+        session._header = header.model_copy(deep=True)
+
+        # Replay without attaching the store: reading must never append records.
+        for line_number, record in records:
+            try:
+                if isinstance(record, SessionCursorEntry):
+                    if record.entry_id is not None and record.entry_id not in session._by_id:
+                        raise ValueError(f"Unknown cursor target: {record.entry_id}")
+                    session._leaf_id = record.entry_id
+                    session._records.append(record.model_copy(deep=True))
+                else:
+                    session._append_entry(record)
+            except ValueError as exc:
+                raise SessionFormatError(store.path, line_number, str(exc)) from exc
+
+        session._store = store
+        return session
+
+    @property
+    def header(self) -> SessionHeader:
+        with self._lock:
+            return self._header.model_copy(deep=True)
+
+    @property
+    def leaf_id(self) -> str | None:
+        with self._lock:
+            return self._leaf_id
+
+    @property
+    def file_path(self) -> Path | None:
+        with self._lock:
+            return self._store.path if self._store is not None else None
+
+    def save(self, file_path: Path) -> Path:
+        """Attach a new file; subsequent changes are appended automatically."""
+        path = file_path.expanduser().resolve()
+        with self._lock:
+            if self._store is not None:
+                if path != self._store.path:
+                    raise ValueError("This session is already attached to another file")
+                # All committed changes have already been persisted.
+                return self._store.path
+            self._store = JsonlSessionStore.create(path, self._header, self._records)
+            return self._store.path
+
+    def get_entries(self) -> list[SessionEntry]:
+        with self._lock:
+            return [entry.model_copy(deep=True) for entry in self._entries]
+
+    def append_message(self, message: Message, *, task_start: bool = False) -> str:
+        with self._lock:
+            entry = SessionMessageEntry(
+                id=uuid4().hex,
+                parent_id=self._leaf_id,
+                timestamp=int(time.time() * 1000),
+                message=message,
+                task_start=task_start,
+            )
+            return self._append_entry(entry)
+
+    def append_model_change(self, provider: str, model_id: str) -> str:
+        with self._lock:
+            entry = ModelChangeEntry(
+                id=uuid4().hex,
+                parent_id=self._leaf_id,
+                timestamp=int(time.time() * 1000),
+                provider=provider,
+                model_id=model_id,
+            )
+            return self._append_entry(entry)
+
+    def append_compaction(self, entry: CompactionEntry) -> str:
+        with self._lock:
+            if entry.parent_id != self._leaf_id:
+                raise ValueError("Session changed while preparing compaction")
+            return self._append_entry(entry)
+
+    def _append_entry(self, entry: SessionEntry) -> str:
+        if entry.id in self._by_id:
+            raise ValueError(f"Duplicate session entry ID: {entry.id}")
+        if entry.parent_id is not None and entry.parent_id not in self._by_id:
+            raise ValueError(f"Parent entry does not exist: {entry.parent_id}")
+        if isinstance(entry, SessionMessageEntry) and entry.task_start:
+            if not isinstance(entry.message, UserMessage):
+                raise ValueError("Only a user message can start a task")
+        if isinstance(entry, CompactionEntry):
+            project_context_entries(self._get_branch(entry.parent_id) + [entry])
+
+        stored = entry.model_copy(deep=True)
+        # Commit to disk before reporting success or advancing the in-memory tip.
+        if self._store is not None:
+            self._store.append(stored)
+        self._entries.append(stored)
+        self._records.append(stored)
+        self._by_id[stored.id] = stored
+        self._leaf_id = stored.id
+        return stored.id
+
+    def checkout(self, entry_id: str | None) -> None:
+        with self._lock:
+            if entry_id is not None and entry_id not in self._by_id:
+                raise ValueError(f"Session entry does not exist: {entry_id}")
+            if entry_id == self._leaf_id:
+                return
+            record = SessionCursorEntry(
+                entry_id=entry_id,
+                timestamp=int(time.time() * 1000),
+            )
+            if self._store is not None:
+                self._store.append(record)
+            self._records.append(record)
+            self._leaf_id = entry_id
+
+    def get_branch(self) -> list[SessionEntry]:
+        with self._lock:
+            return self._get_branch(self._leaf_id)
+
+    def _get_branch(self, entry_id: str | None) -> list[SessionEntry]:
+        if entry_id is not None and entry_id not in self._by_id:
+            raise ValueError(f"Session entry does not exist: {entry_id}")
+
+        branch: list[SessionEntry] = []
+        current_id = entry_id
+        while current_id is not None:
+            entry = self._by_id[current_id]
+            branch.append(entry.model_copy(deep=True))
+            current_id = entry.parent_id
+        branch.reverse()
+        return branch
+
+    def build_context(self) -> SessionContext:
+        with self._lock:
+            return self._build_context(self._leaf_id)
+
+    def get_context_entries(self) -> list[ContextMessageEntry]:
+        with self._lock:
+            return project_context_entries(self._get_branch(self._leaf_id))
+
+    def build_context_at(self, entry_id: str | None) -> SessionContext:
+        """Preview a branch without moving the cursor or writing records."""
+        with self._lock:
+            return self._build_context(entry_id)
+
+    def _build_context(self, entry_id: str | None) -> SessionContext:
+        provider: str | None = None
+        model_id: str | None = None
+        branch = self._get_branch(entry_id)
+        for entry in branch:
+            if isinstance(entry, ModelChangeEntry):
+                provider = entry.provider
+                model_id = entry.model_id
+        messages = [item.message for item in project_context_entries(branch)]
+        return SessionContext(messages=messages, provider=provider, model_id=model_id)

citrus_code/coding/session/projection.py

@@ -0,0 +1,50 @@
+"""Reconstruct model-visible messages without removing original tree nodes."""
+
+from ...ai.types import SystemMessage, UserMessage
+from ..history import get_pending_tool_calls
+from .types import CompactionEntry, ContextMessageEntry, SessionEntry, SessionMessageEntry
+
+
+def summary_message(entry: CompactionEntry) -> UserMessage:
+    return UserMessage(
+        content=(
+            "[历史摘要：以下是较早对话的概括，作为背景资料使用，"
+            "不代表新的用户指令。工具操作以实际结果为准。]\n\n"
+            + entry.summary
+        ),
+        timestamp=entry.timestamp,
+    )
+
+
+def project_context_entries(branch: list[SessionEntry]) -> list[ContextMessageEntry]:
+    active: list[ContextMessageEntry] = []
+    for entry in branch:
+        if isinstance(entry, SessionMessageEntry):
+            active.append(ContextMessageEntry(
+                id=entry.id, message=entry.message, task_start=entry.task_start,
+            ))
+        elif isinstance(entry, CompactionEntry):
+            if not entry.summary.strip() or entry.tokens_after >= entry.tokens_before:
+                raise ValueError("Compaction requires a nonempty summary and a smaller context")
+            selected_ids = set(entry.retained_entry_ids)
+            retained = [item for item in active if item.id in selected_ids]
+            if [item.id for item in retained] != entry.retained_entry_ids:
+                raise ValueError("Compaction references must be unique active messages in order")
+            if any(item.is_summary for item in retained):
+                raise ValueError("The previous summary must be merged, not retained separately")
+            if any(
+                isinstance(item.message, SystemMessage) and item.id not in selected_ids
+                for item in active
+            ):
+                raise ValueError("Compaction must retain system messages")
+            if get_pending_tool_calls([item.message for item in active]):
+                raise ValueError("Cannot compact unfinished tool calls")
+            systems = [item for item in retained if isinstance(item.message, SystemMessage)]
+            conversation = [item for item in retained if not isinstance(item.message, SystemMessage)]
+            projected = systems + [ContextMessageEntry(
+                id=entry.id, message=summary_message(entry), is_summary=True,
+            )] + conversation
+            if get_pending_tool_calls([item.message for item in projected]):
+                raise ValueError("Compaction retained an incomplete tool interaction")
+            active = projected
+    return active

citrus_code/coding/session/recovery.py

@@ -0,0 +1,33 @@
+"""Repair missing results at a valid session tail without executing tools."""
+
+import time
+
+from ...ai.types import TextContent, ToolResultMessage
+from ..history import get_pending_tool_calls
+from .manager import SessionManager
+
+
+def recover_pending_tools(manager: SessionManager) -> int:
+    """Call before starting an Agent, with exclusive session ownership."""
+    pending = get_pending_tool_calls(manager.build_context().messages)
+    for call in pending:
+        manager.append_message(
+            ToolResultMessage(
+                tool_call_id=call.id,
+                tool_name=call.name,
+                content=[
+                    TextContent(
+                        text=(
+                            "The session was interrupted before a result was recorded "
+                            "for this tool call. Its execution status is unknown: it "
+                            "may not have started, may have partially run, or may have "
+                            "completed. Inspect the actual state before deciding "
+                            "whether to retry."
+                        )
+                    )
+                ],
+                is_error=True,
+                timestamp=int(time.time() * 1000),
+            )
+        )
+    return len(pending)

citrus_code/coding/session/storage.py

@@ -0,0 +1,138 @@
+"""Synchronous JSONL persistence, intended for one writer per session file."""
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from threading import RLock
+
+from pydantic import TypeAdapter, ValidationError
+
+from .types import (
+    SessionFileEntry,
+    SessionHeader,
+    SessionLogEntry,
+)
+
+
+_RECORD_ADAPTER = TypeAdapter(SessionFileEntry)
+_IO_LOCK = RLock()
+
+
+class SessionFormatError(ValueError):
+    def __init__(self, path: Path, line: int, reason: str) -> None:
+        super().__init__(f"Invalid session file {path}, line {line}: {reason}")
+
+
+class SessionConflictError(RuntimeError):
+    """The stored file changed after this manager loaded or wrote it."""
+
+
+@dataclass(frozen=True)
+class _Revision:
+    device: int
+    inode: int
+    size: int
+    modified_ns: int
+
+    @classmethod
+    def from_stat(cls, stat: os.stat_result) -> "_Revision":
+        return cls(stat.st_dev, stat.st_ino, stat.st_size, stat.st_mtime_ns)
+
+
+class JsonlSessionStore:
+    def __init__(self, path: Path, revision: _Revision) -> None:
+        self.path = path
+        self._revision = revision
+        self._failed = False
+
+    @classmethod
+    def create(
+        cls,
+        path: Path,
+        header: SessionHeader,
+        records: list[SessionLogEntry],
+    ) -> "JsonlSessionStore":
+        path = path.expanduser().resolve()
+        with _IO_LOCK:
+            path.parent.mkdir(parents=True, exist_ok=True)
+            # Never overwrite an existing session, including a damaged file.
+            with path.open("xb") as file:
+                file.write(header.model_dump_json(by_alias=True).encode("utf-8") + b"\n")
+                for record in records:
+                    file.write(record.model_dump_json(by_alias=True).encode("utf-8") + b"\n")
+                file.flush()
+                os.fsync(file.fileno())
+                revision = _Revision.from_stat(os.fstat(file.fileno()))
+        return cls(path, revision)
+
+    @classmethod
+    def load(
+        cls,
+        path: Path,
+    ) -> tuple["JsonlSessionStore", SessionHeader, list[tuple[int, SessionLogEntry]]]:
+        path = path.expanduser().resolve()
+        header: SessionHeader | None = None
+        records: list[tuple[int, SessionLogEntry]] = []
+
+        with _IO_LOCK:
+            with path.open("rb") as file:
+                before = _Revision.from_stat(os.fstat(file.fileno()))
+                for line_number, line in enumerate(file, start=1):
+                    if not line.strip():
+                        continue
+                    try:
+                        record = _RECORD_ADAPTER.validate_json(line)
+                    except ValidationError as exc:
+                        raise SessionFormatError(
+                            path, line_number, "invalid JSON or record structure"
+                        ) from exc
+
+                    if header is None:
+                        if not isinstance(record, SessionHeader):
+                            raise SessionFormatError(path, line_number, "the first record must be a session header")
+                        header = record
+                    elif isinstance(record, SessionHeader):
+                        raise SessionFormatError(path, line_number, "duplicate session header")
+                    else:
+                        records.append((line_number, record))
+
+                after = _Revision.from_stat(os.fstat(file.fileno()))
+            if before != after or after != _Revision.from_stat(path.stat()):
+                raise SessionConflictError(f"Session changed while loading: {path}")
+
+        if header is None:
+            raise SessionFormatError(path, 1, "missing session header")
+        return cls(path, after), header, records
+
+    def append(self, record: SessionLogEntry) -> None:
+        payload = record.model_dump_json(by_alias=True).encode("utf-8") + b"\n"
+        with _IO_LOCK:
+            if self._failed:
+                raise SessionConflictError(
+                    "A previous write failed. Reload and inspect the session before writing again."
+                )
+            with self.path.open("r+b") as file:
+                current = _Revision.from_stat(os.fstat(file.fileno()))
+                if current != self._revision:
+                    raise SessionConflictError(f"Session was changed by another writer: {self.path}")
+                if current != _Revision.from_stat(self.path.stat()):
+                    raise SessionConflictError(f"Session path was replaced: {self.path}")
+
+                # Accept a valid final JSON record without LF, adding its separator
+                # before the new record. An incomplete record was rejected on load.
+                if current.size:
+                    file.seek(-1, os.SEEK_END)
+                    if file.read(1) != b"\n":
+                        payload = b"\n" + payload
+                file.seek(0, os.SEEK_END)
+                try:
+                    written = file.write(payload)
+                    if written != len(payload):
+                        raise OSError("Incomplete session write")
+                    file.flush()
+                    os.fsync(file.fileno())
+                    self._revision = _Revision.from_stat(os.fstat(file.fileno()))
+                except BaseException:
+                    # A partial tail is not silently truncated or skipped.
+                    self._failed = True
+                    raise

citrus_code/coding/session/types.py

@@ -0,0 +1,108 @@
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from ...ai.types import Message
+
+
+class SessionRecord(BaseModel):
+    model_config = ConfigDict(
+        extra="forbid",
+        strict=True,
+        populate_by_name=True,
+    )
+
+
+class SessionHeader(SessionRecord):
+    type: Literal["session"] = "session"
+
+    version: Literal[1] = 1
+
+    id: str = Field(min_length=1)
+    timestamp: int = Field(ge=0)
+    cwd: str = Field(min_length=1)
+
+
+class SessionEntryBase(SessionRecord):
+    id: str = Field(min_length=1)
+
+    parent_id: str | None = Field(
+        alias="parentId",
+    )
+
+    timestamp: int = Field(ge=0)
+
+
+class SessionMessageEntry(SessionEntryBase):
+    type: Literal["message"] = "message"
+
+    message: Message
+    task_start: bool = Field(default=False, alias="taskStart")
+
+
+class CompactionEntry(SessionEntryBase):
+    type: Literal["compaction"] = "compaction"
+    summary: str = Field(min_length=1)
+    retained_entry_ids: list[str] = Field(alias="retainedEntryIds", min_length=1)
+    tokens_before: int = Field(alias="tokensBefore", ge=0)
+    tokens_after: int = Field(alias="tokensAfter", ge=0)
+
+
+class ContextMessageEntry(BaseModel):
+    """A message in the active projection, including a synthetic summary."""
+
+    model_config = ConfigDict(extra="forbid")
+    id: str
+    message: Message
+    task_start: bool = False
+    is_summary: bool = False
+
+
+class ModelChangeEntry(SessionEntryBase):
+    type: Literal["model_change"] = "model_change"
+
+    provider: str = Field(min_length=1)
+
+    model_id: str = Field(
+        alias="modelId",
+        min_length=1,
+    )
+
+
+class SessionCursorEntry(SessionRecord):
+    """Persist the selected branch tip without creating a new tree node."""
+
+    type: Literal["cursor"] = "cursor"
+    entry_id: str | None = Field(alias="entryId")
+    timestamp: int = Field(ge=0)
+
+
+SessionEntry = Annotated[
+    SessionMessageEntry | ModelChangeEntry | CompactionEntry,
+    Field(discriminator="type"),
+]
+
+
+SessionLogEntry = Annotated[
+    SessionMessageEntry | ModelChangeEntry | CompactionEntry | SessionCursorEntry,
+    Field(discriminator="type"),
+]
+
+
+SessionFileEntry = Annotated[
+    SessionHeader
+    | SessionMessageEntry
+    | ModelChangeEntry
+    | CompactionEntry
+    | SessionCursorEntry,
+    Field(discriminator="type"),
+]
+
+
+class SessionContext(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    messages: list[Message]
+
+    provider: str | None = None
+    model_id: str | None = None

citrus_code/coding/tools/__init__.py

@@ -0,0 +1,35 @@
+from .bash import BashTool
+from .edit import EditTool
+from .factory import (
+    DEFAULT_CODING_TOOLS,
+    READ_ONLY_TOOLS,
+    create_all_tools,
+    create_coding_tools,
+    create_read_only_tools,
+    create_tool,
+    create_tools,
+)
+from .find import FindTool
+from .grep import GrepTool
+from .ls import LsTool
+from .powershell import PowerShellTool
+from .read import ReadTool
+from .write import WriteTool
+
+__all__ = [
+    "BashTool",
+    "DEFAULT_CODING_TOOLS",
+    "EditTool",
+    "FindTool",
+    "GrepTool",
+    "LsTool",
+    "PowerShellTool",
+    "READ_ONLY_TOOLS",
+    "ReadTool",
+    "WriteTool",
+    "create_all_tools",
+    "create_coding_tools",
+    "create_read_only_tools",
+    "create_tool",
+    "create_tools",
+]

citrus_code/coding/tools/bash.py

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

citrus_code/coding/tools/edit.py

@@ -0,0 +1,101 @@
+import asyncio
+import difflib
+from pathlib import Path
+
+from pydantic import BaseModel, ConfigDict, Field, JsonValue
+
+from ...agent.types import ToolExecutionResult
+from ...ai.types import TextContent, ToolDefinition
+from .edit_text import TextEdit, apply_text_edits, normalize_to_lf
+from .mutation import file_mutation
+from .paths import resolve_to_cwd
+
+
+class EditArguments(BaseModel):
+    model_config = ConfigDict(extra="forbid", strict=True)
+
+    path: str = Field(min_length=1, description="File to edit, relative or absolute.")
+    edits: list[TextEdit] = Field(
+        min_length=1,
+        description=(
+            "Targeted replacements matched against the original file. "
+            "Each oldText must be unique. Regions must not overlap; "
+            "merge nearby changes into one edit."
+        ),
+    )
+
+
+class EditTool:
+    name: str = "edit"
+    description: str = (
+        "Edit one file using exact text replacement. Each edits[].oldText must "
+        "match a unique, non-overlapping region of the original file. "
+        "Use multiple edits for separate locations; merge overlapping changes."
+    )
+
+    def __init__(self, cwd: Path) -> None:
+        self._cwd = cwd.resolve()
+
+    def get_definition(self) -> ToolDefinition:
+        return ToolDefinition(
+            name=self.name,
+            description=self.description,
+            parameters=EditArguments.model_json_schema(by_alias=True),
+        )
+
+    async def execute(
+        self,
+        tool_call_id: str,
+        arguments: dict[str, JsonValue],
+    ) -> ToolExecutionResult:
+        args = EditArguments.model_validate(arguments)
+        path = resolve_to_cwd(args.path, self._cwd)
+        return await asyncio.to_thread(self._edit, path, args)
+
+    def _edit(self, path: Path, args: EditArguments) -> ToolExecutionResult:
+        with file_mutation(path):
+            raw = path.read_bytes().decode("utf-8")
+            bom = "\ufeff" if raw.startswith("\ufeff") else ""
+            text = raw[len(bom):]
+            first_newline = text.find("\n")
+            ending = (
+                "\r\n"
+                if first_newline > 0 and text[first_newline - 1] == "\r"
+                else "\n"
+            )
+            original = normalize_to_lf(text)
+            updated = apply_text_edits(original, args.edits, args.path)
+
+            old_lines = original.splitlines(keepends=True)
+            new_lines = updated.splitlines(keepends=True)
+            patch_parts: list[str] = []
+            for line in difflib.unified_diff(
+                old_lines, new_lines, fromfile=args.path, tofile=args.path, n=4
+            ):
+                patch_parts.append(line)
+                if not line.endswith("\n"):
+                    patch_parts.append("\n\\n")
+            patch = "".join(patch_parts)
+
+            first_changed_line = 1
+            matcher = difflib.SequenceMatcher(a=old_lines, b=new_lines, autojunk=False)
+            for tag, _, _, new_start, _ in matcher.get_opcodes():
+                if tag != "equal":
+                    first_changed_line = new_start + 1
+                    break
+
+            # Match and prepare the report before touching the file.
+            final = bom + updated.replace("\n", ending)
+            path.write_bytes(final.encode("utf-8"))
+
+        return ToolExecutionResult(
+            content=[TextContent(
+                text=f"Successfully replaced {len(args.edits)} block(s) in {args.path}."
+            )],
+            details={
+                "path": str(path),
+                "diff": patch,
+                "patch": patch,
+                "firstChangedLine": first_changed_line,
+            },
+        )