yycode 0.3.2__py3-none-any.whl

agent/change_snapshot.py

@@ -0,0 +1,186 @@
+"""Shared changed-files snapshot helpers."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from typing import Iterable
+
+
+@dataclass(frozen=True)
+class ChangedFile:
+    """One changed file and its diff stats."""
+
+    path: str
+    added: int
+    removed: int
+    diff: str = ""
+
+
+@dataclass(frozen=True)
+class ChangedFilesSnapshot:
+    """A stable changed-files snapshot for TUI/protocol adapters."""
+
+    files: list[ChangedFile] = field(default_factory=list)
+    total_added: int = 0
+    total_removed: int = 0
+    source: str = ""
+
+
+def extract_diff_text(content: str) -> str:
+    """Return the unified diff portion from a tool result, if present."""
+    if not content:
+        return ""
+    lines = content.splitlines()
+    for index, line in enumerate(lines):
+        if line.startswith(("diff --git ", "--- ")):
+            return "\n".join(lines[index:])
+    return ""
+
+
+def build_changed_files_snapshot(diff: str, *, source: str = "diff") -> ChangedFilesSnapshot:
+    """Return per-file stats and diff sections from a unified diff blob."""
+    sections: list[dict] = []
+    current: dict | None = None
+    for line in diff.splitlines():
+        if line.startswith("diff --git "):
+            if current is not None:
+                sections.append(current)
+            current = {"path": _path_from_diff_header(line), "added": 0, "removed": 0, "lines": [line]}
+            continue
+        if line.startswith("--- "):
+            if current is not None and _diff_section_has_changes(current):
+                sections.append(current)
+                current = {
+                    "path": _strip_diff_prefix(line[4:].split("\t", 1)[0].strip()),
+                    "added": 0,
+                    "removed": 0,
+                    "lines": [line],
+                }
+            elif current is None:
+                current = {
+                    "path": _strip_diff_prefix(line[4:].split("\t", 1)[0].strip()),
+                    "added": 0,
+                    "removed": 0,
+                    "lines": [line],
+                }
+            else:
+                current["lines"].append(line)
+            continue
+        if current is None:
+            continue
+        current["lines"].append(line)
+        if line.startswith("+++ "):
+            path = _strip_diff_prefix(line[4:].split("\t", 1)[0].strip())
+            if path != "/dev/null":
+                current["path"] = path
+            continue
+        if line.startswith("--- ") or line.startswith("@@") or line.startswith("index "):
+            continue
+        if line.startswith("+"):
+            current["added"] += 1
+        elif line.startswith("-"):
+            current["removed"] += 1
+    if current is not None:
+        sections.append(current)
+    files = _merge_changed_file_sections(sections)
+    return ChangedFilesSnapshot(
+        files=files,
+        total_added=sum(item.added for item in files),
+        total_removed=sum(item.removed for item in files),
+        source=source,
+    )
+
+
+def changed_files_as_dicts(snapshot: ChangedFilesSnapshot) -> list[dict]:
+    """Return snapshot files in the legacy dict shape used by TUI state."""
+    return [
+        {
+            "path": file.path,
+            "added": file.added,
+            "removed": file.removed,
+            "diff": file.diff,
+        }
+        for file in snapshot.files
+    ]
+
+
+def merge_changed_paths(items: Iterable[object]) -> list[str]:
+    """Return unique changed paths from timeline-like items."""
+    paths: list[str] = []
+    for item in items:
+        event_type = getattr(item, "event_type", "")
+        tool_name = getattr(item, "tool_name", "")
+        if event_type == "file_changed":
+            candidates = list(getattr(item, "file_paths", []) or [])
+        elif event_type in {"tool_start", "tool_end"} and tool_name in {
+            "apply_patch",
+            "write_file",
+            "edit_file",
+        }:
+            candidates = list(getattr(item, "file_paths", []) or [])
+        else:
+            candidates = []
+        for path in candidates:
+            if path and path not in paths:
+                paths.append(path)
+    return paths
+
+
+def turn_had_successful_write(items: Iterable[object]) -> bool:
+    """Return whether timeline-like items include a successful write tool."""
+    return any(
+        getattr(item, "event_type", "") == "tool_end"
+        and getattr(item, "status", None) != "failed"
+        and getattr(item, "tool_name", "") in {"apply_patch", "write_file", "edit_file"}
+        for item in items
+    )
+
+
+def _merge_changed_file_sections(sections: list[dict]) -> list[ChangedFile]:
+    merged: dict[str, dict] = {}
+    order: list[str] = []
+    for item in sections:
+        if not item.get("added") and not item.get("removed"):
+            continue
+        path = str(item.get("path", ""))
+        if not path:
+            continue
+        if path not in merged:
+            merged[path] = {"path": path, "added": 0, "removed": 0, "diffs": []}
+            order.append(path)
+        merged_item = merged[path]
+        merged_item["added"] += int(item.get("added", 0) or 0)
+        merged_item["removed"] += int(item.get("removed", 0) or 0)
+        merged_item["diffs"].append("\n".join(item.get("lines", [])))
+    return [
+        ChangedFile(
+            path=merged[path]["path"],
+            added=merged[path]["added"],
+            removed=merged[path]["removed"],
+            diff="\n\n".join(diff for diff in merged[path]["diffs"] if diff),
+        )
+        for path in order
+    ]
+
+
+def _diff_section_has_changes(section: dict) -> bool:
+    return bool(
+        section.get("added")
+        or section.get("removed")
+        or any(str(line).startswith("@@") for line in section.get("lines", []))
+    )
+
+
+def _path_from_diff_header(line: str) -> str:
+    match = re.match(r"diff --git a/(.+?) b/(.+)$", line)
+    if match:
+        return match.group(2)
+    parts = line.split()
+    return _strip_diff_prefix(parts[-1]) if parts else "file"
+
+
+def _strip_diff_prefix(path: str) -> str:
+    if path.startswith("a/") or path.startswith("b/"):
+        return path[2:]
+    return path

agent/context_compressor.py

@@ -0,0 +1,116 @@
+"""Lightweight context compression for long-running sessions."""
+
+from dataclasses import dataclass
+
+from langchain_core.messages import BaseMessage, ToolMessage
+
+
+DEFAULT_COMPRESSION_RATIO = 0.7
+DEFAULT_KEEP_RECENT_MESSAGES = 20
+DEFAULT_MAX_TOOL_CHARS = 2_000
+AUTO_COMPRESSION_REASON = "context window usage crossed the compression threshold."
+MANUAL_COMPRESSION_REASON = "manually compressed by Message Token Manager."
+
+
+@dataclass(frozen=True)
+class CompressionResult:
+    """Result of a context compression pass."""
+
+    messages: list[BaseMessage]
+    did_compress: bool
+    original_tokens: int
+    compressed_tokens: int
+    compressed_messages: int
+
+
+class ContextCompressor:
+    """Conservative compressor that trims old tool outputs."""
+
+    def __init__(
+        self,
+        *,
+        context_window_tokens: int,
+        compression_ratio: float = DEFAULT_COMPRESSION_RATIO,
+        keep_recent_messages: int = DEFAULT_KEEP_RECENT_MESSAGES,
+        max_tool_chars: int = DEFAULT_MAX_TOOL_CHARS,
+    ):
+        self.context_window_tokens = context_window_tokens
+        self.compression_ratio = compression_ratio
+        self.keep_recent_messages = keep_recent_messages
+        self.max_tool_chars = max_tool_chars
+
+    def maybe_compress(
+        self,
+        messages: list[BaseMessage],
+        estimated_tokens: int,
+        estimate_tokens,
+    ) -> CompressionResult:
+        """Compress old tool outputs when context usage crosses the threshold."""
+        if estimated_tokens < self.threshold_tokens:
+            return CompressionResult(messages, False, estimated_tokens, estimated_tokens, 0)
+
+        compressed_messages = []
+        compressed_count = 0
+        cutoff = max(len(messages) - self.keep_recent_messages, 0)
+        for index, message in enumerate(messages):
+            if index < cutoff and self._should_trim_tool_message(message):
+                compressed_messages.append(self._trim_tool_message(message))
+                compressed_count += 1
+            else:
+                compressed_messages.append(message)
+
+        if compressed_count == 0:
+            return CompressionResult(messages, False, estimated_tokens, estimated_tokens, 0)
+
+        compressed_tokens = estimate_tokens(compressed_messages)
+        return CompressionResult(
+            compressed_messages,
+            compressed_tokens < estimated_tokens,
+            estimated_tokens,
+            compressed_tokens,
+            compressed_count,
+        )
+
+    @property
+    def threshold_tokens(self) -> int:
+        """Token threshold that triggers compression."""
+        return int(self.context_window_tokens * self.compression_ratio)
+
+    def _should_trim_tool_message(self, message: BaseMessage) -> bool:
+        if not isinstance(message, ToolMessage):
+            return False
+        content = message.content
+        return isinstance(content, str) and len(content) > self.max_tool_chars
+
+    def _trim_tool_message(self, message: ToolMessage) -> ToolMessage:
+        return compress_tool_message(message, reason=AUTO_COMPRESSION_REASON)
+
+
+def compress_tool_message(
+    message: ToolMessage,
+    *,
+    reason: str,
+    estimated_original_tokens: int | None = None,
+) -> ToolMessage:
+    """Return a compact marker ToolMessage while preserving tool linkage."""
+    original_chars = len(str(message.content))
+    name = message.name or "unknown"
+    lines = [
+        "[Compressed old tool output]",
+        f"tool: {name}",
+        f"original_chars: {original_chars}",
+    ]
+    if estimated_original_tokens is not None:
+        lines.append(f"estimated_original_tokens: {estimated_original_tokens}")
+    lines.append(f"reason: {reason}")
+    trimmed = ToolMessage(
+        content="\n".join(lines),
+        tool_call_id=message.tool_call_id,
+        name=message.name,
+    )
+    trimmed.additional_kwargs.update(getattr(message, "additional_kwargs", {}) or {})
+    trimmed.additional_kwargs["context_compressed"] = True
+    trimmed.additional_kwargs["original_chars"] = original_chars
+    if estimated_original_tokens is not None:
+        trimmed.additional_kwargs["estimated_original_tokens"] = estimated_original_tokens
+    return trimmed

agent/graph.py

@@ -0,0 +1,137 @@
+"""Agent graph definition using LangGraph with provider abstraction."""
+
+from pathlib import Path
+
+from langgraph.graph import StateGraph, START
+
+from agent.approval import ApprovalCallback
+from agent.nodes.llm_node import create_llm_node as _create_llm_node
+from agent.nodes.state import AgentState
+from agent.nodes.task_guard_node import (
+    create_task_guard_node,
+    route_after_llm,
+    route_after_tools,
+    route_after_task_guard,
+)
+from agent.nodes.tools_node import create_tools_node as _create_tools_node
+from agent.providers.base import LLMProvider
+from agent.runtime.context import AgentRuntimeContext, WorkflowState
+from agent.streaming import StreamEventCallback
+from agent.tool_retry import async_run_tool_with_retry
+from tools import TOOL_HANDLERS, TOOLS
+
+
+def create_runtime(
+    provider: LLMProvider,
+    system_prompt: str,
+    todo_manager,
+    workdir: Path,
+    session_id: str,
+    skill_dirs: list[str] | None = None,
+    app_root: Path | None = None,
+    stream_callback: StreamEventCallback = None,
+    approval_callback: ApprovalCallback = None,
+) -> AgentRuntimeContext:
+    """Create runtime context for a graph run."""
+    return AgentRuntimeContext(
+        provider=provider,
+        system_prompt=system_prompt,
+        todo_manager=todo_manager,
+        workdir=workdir,
+        session_id=session_id,
+        skill_dirs=skill_dirs,
+        app_root=app_root,
+        stream_callback=stream_callback,
+        approval_callback=approval_callback,
+        tools=TOOLS,
+        tool_handlers=TOOL_HANDLERS,
+        workflow_state=WorkflowState(),
+        run_tool=async_run_tool_with_retry,
+    )
+
+
+def create_llm_node(
+    provider: LLMProvider,
+    system_prompt: str,
+    session_id: str,
+    stream_callback: StreamEventCallback = None,
+):
+    """Backward-compatible LLM node factory."""
+    runtime = AgentRuntimeContext(
+        provider=provider,
+        system_prompt=system_prompt,
+        todo_manager=None,
+        workdir=Path.cwd(),
+        session_id=session_id,
+        stream_callback=stream_callback,
+        tools=TOOLS,
+        tool_handlers=TOOL_HANDLERS,
+    )
+    return _create_llm_node(runtime)
+
+
+def create_tools_node(
+    provider: LLMProvider,
+    system_prompt: str,
+    todo_manager,
+    workdir: Path,
+    session_id: str,
+    skill_dirs: list[str] | None = None,
+    app_root: Path | None = None,
+    stream_callback: StreamEventCallback = None,
+    approval_callback: ApprovalCallback = None,
+):
+    """Backward-compatible tools node factory."""
+    return _create_tools_node(
+        create_runtime(
+            provider=provider,
+            system_prompt=system_prompt,
+            todo_manager=todo_manager,
+            workdir=workdir,
+            session_id=session_id,
+            skill_dirs=skill_dirs,
+            app_root=app_root,
+            stream_callback=stream_callback,
+            approval_callback=approval_callback,
+        )
+    )
+
+
+def build_graph(
+    provider: LLMProvider,
+    system_prompt: str,
+    todo_manager,
+    workdir: Path,
+    session_id: str,
+    skill_dirs: list[str] | None = None,
+    app_root: Path | None = None,
+    stream_callback: StreamEventCallback = None,
+    approval_callback: ApprovalCallback = None,
+):
+    """Build the agent graph."""
+    runtime = create_runtime(
+        provider=provider,
+        system_prompt=system_prompt,
+        todo_manager=todo_manager,
+        workdir=workdir,
+        session_id=session_id,
+        skill_dirs=skill_dirs,
+        app_root=app_root,
+        stream_callback=stream_callback,
+        approval_callback=approval_callback,
+    )
+
+    builder = StateGraph(AgentState)
+    builder.add_node("llm", _create_llm_node(runtime))
+    builder.add_node("tools", _create_tools_node(runtime))
+    builder.add_node("task_guard", create_task_guard_node(todo_manager))
+
+    builder.add_edge(START, "llm")
+    builder.add_conditional_edges("llm", route_after_llm)
+    builder.add_conditional_edges("tools", route_after_tools)
+    builder.add_conditional_edges(
+        "task_guard",
+        lambda state: route_after_task_guard(state, todo_manager),
+    )
+
+    return builder.compile()