rockycode 0.1.0__py3-none-any.whl

rockycode/engine/compaction.py

@@ -0,0 +1,181 @@
+"""Context compaction: keep long sessions inside the model's window.
+
+Two stages, cheapest first:
+
+1. prune — deterministic, free. Old tool outputs (the bulk of any agent
+   context) are stubbed down to a short head; the assistant's own messages
+   and tool_calls survive, so the model still sees *what* it did, just not
+   every byte it once read.
+2. summarize — the recurrent state. If pruning can't get under the limit,
+   one non-streaming API call folds everything except a recent tail into a
+   dense state document, and history is rebuilt as [system, state, tail].
+   The call sends the *unmodified* history plus one instruction message, so
+   DeepSeek's automatic prefix cache makes its input nearly free.
+
+Token math is deliberately conservative (3 chars ≈ 1 token): compacting a
+little early costs a few cache misses; compacting late kills the turn.
+"""
+from __future__ import annotations
+
+CHARS_PER_TOKEN = 3
+MSG_OVERHEAD_TOKENS = 8
+
+KEEP_RECENT_TOKENS = 8_000     # protected tail budget (estimated tokens)
+PRUNE_KEEP_CHARS = 200         # head kept when a tool output is stubbed
+PRUNE_STUB = "… [old output dropped during compaction — re-run the tool if needed]"
+SUMMARY_MAX_TOKENS = 4096
+
+SUMMARIZE_INSTRUCTION = """\
+Context is nearly full. Before continuing, write a compact state document
+for this session so far; older messages will be dropped and replaced by it.
+
+Cover, with concrete file paths, names, and line references:
+1. Task — what we are trying to accomplish, in one or two sentences.
+2. State — what has been done and learned so far (files read/edited, key
+   findings, commands run and their results).
+3. Failures — approaches that did not work, so they are not retried.
+4. Next — the immediate plan.
+
+Write only the state document, no preamble. Be dense; facts over prose.
+"""
+
+
+def estimate_msg_tokens(msg: dict) -> int:
+    chars = len(msg.get("content") or "")
+    for tc in msg.get("tool_calls") or []:
+        fn = tc.get("function", {})
+        chars += len(fn.get("name", "")) + len(fn.get("arguments", ""))
+    return MSG_OVERHEAD_TOKENS + chars // CHARS_PER_TOKEN
+
+
+def estimate_tokens(messages: list[dict]) -> int:
+    return sum(estimate_msg_tokens(m) for m in messages)
+
+
+def tail_start(history: list[dict], budget_tokens: int) -> int:
+    """Index where the protected recent tail begins.
+
+    Never 0 (the system prompt is handled separately), always compresses at
+    least the older half of the conversation, and never starts on a `tool`
+    message — a tool result without its assistant tool_calls parent is an
+    invalid conversation, so the tail grows backward to include the parent.
+    """
+    i = len(history)
+    spent = 0
+    while i > 1:
+        cost = estimate_msg_tokens(history[i - 1])
+        if spent and spent + cost > budget_tokens:
+            break
+        spent += cost
+        i -= 1
+    floor = 1 + (len(history) - 1) // 2
+    if i < floor:
+        i = floor
+    while 1 < i < len(history) and history[i]["role"] == "tool":
+        i -= 1
+    return i
+
+
+def _prune_gain_chars(content: str) -> int:
+    return len(content) - PRUNE_KEEP_CHARS - len(PRUNE_STUB) - 1
+
+
+def prune_savings(history: list[dict], protect_from: int) -> int:
+    """Estimated tokens that prune_tool_outputs would free. No mutation."""
+    saved = 0
+    for msg in history[1:protect_from]:
+        if msg.get("role") == "tool":
+            gain = _prune_gain_chars(msg.get("content") or "")
+            if gain > 0:
+                saved += gain // CHARS_PER_TOKEN
+    return saved
+
+
+def prune_tool_outputs(history: list[dict], protect_from: int) -> int:
+    """Stub bulky tool outputs in history[1:protect_from], in place.
+
+    Returns the number of messages stubbed. Deterministic given the same
+    history and protect_from, so a trajectory's compaction record is enough
+    to replay it.
+    """
+    n = 0
+    for msg in history[1:protect_from]:
+        if msg.get("role") != "tool":
+            continue
+        content = msg.get("content") or ""
+        if _prune_gain_chars(content) <= 0:
+            continue
+        msg["content"] = content[:PRUNE_KEEP_CHARS] + "\n" + PRUNE_STUB
+        n += 1
+    return n
+
+
+MAX_MSG_CHARS = 40_000  # ~13k tokens; a single message above this is truncated
+                        # head+tail, so one giant paste / tool result can't defeat
+                        # compaction or overflow the window.
+
+
+def truncate_oversized(history: list[dict], max_chars: int = MAX_MSG_CHARS) -> int:
+    """Truncate any single message whose text content exceeds *max_chars* to a
+    head + tail with an elision marker, in place. Skips the system prompt.
+    Returns the count truncated.
+
+    The backstop for what prune and summarize miss: prune only shrinks `tool`
+    messages and summarize keeps the recent tail verbatim, so an oversized
+    NON-tool message (e.g. a pasted 300k-char log that tail_start always keeps)
+    survives both and keeps every step over the limit. This bounds it.
+    """
+    n = 0
+    keep = max_chars // 2 - 60
+    for msg in history[1:]:  # never the system prompt
+        content = msg.get("content")
+        if not isinstance(content, str) or len(content) <= max_chars:
+            continue
+        elided = len(content) - 2 * keep
+        msg["content"] = (
+            content[:keep]
+            + f"\n… [{elided:,} chars elided to fit the context window] …\n"
+            + content[-keep:]
+        )
+        n += 1
+    return n
+
+
+async def summarize(client, model: str, history: list[dict], tools: list[dict]) -> tuple[str, dict]:
+    """One non-streaming call: history + instruction → state document.
+
+    `tools` is passed through (with tool_choice="none") so the request body
+    matches the main loop's shape and the prefix cache can hit; thinking is
+    off — a summary needs no chain of thought.
+    """
+    resp = await client.chat.completions.create(
+        model=model,
+        messages=history + [{"role": "user", "content": SUMMARIZE_INSTRUCTION}],
+        tools=tools,
+        tool_choice="none",
+        max_tokens=SUMMARY_MAX_TOKENS,
+        stream=False,
+        extra_body={"thinking": {"type": "disabled"}},
+    )
+    summary = (resp.choices[0].message.content or "").strip()
+    usage: dict = {}
+    if resp.usage is not None:
+        try:
+            u = resp.usage.model_dump()
+        except AttributeError:
+            u = dict(resp.usage)
+        usage = {k: v for k, v in u.items() if isinstance(v, int)}
+    return summary, usage
+
+
+def state_message(summary: str) -> dict:
+    return {
+        "role": "user",
+        "content": (
+            "[context compacted] Older messages were compressed into this "
+            "state document:\n\n"
+            f"{summary}\n\n"
+            "Continue the task from this state. The most recent messages "
+            "follow unmodified — do not redo work recorded as done above."
+        ),
+    }

rockycode/engine/container.py

@@ -0,0 +1,225 @@
+"""Run engine tools inside a SWE-bench task container.
+
+A *session* is anything with `exec(script, stdin, timeout)`. DockerSession
+talks to a long-lived container via `docker exec`; LocalSession runs the
+same scripts in a local directory so the glue is testable without Docker.
+
+SWE-bench image layout (constant across the official images):
+  repo checked out at /testbed at the task's base_commit
+  conda env "testbed" at /opt/miniconda3 with the repo's deps installed
+"""
+from __future__ import annotations
+
+import asyncio
+import shlex
+from pathlib import Path
+from typing import Optional
+
+from rockycode.engine.tools import (
+    GLOB_MAX_PATHS,
+    GREP_MAX_MATCHES,
+    RISK,
+    SCHEMAS,
+    Tool,
+    _truncate,
+)
+
+# Generous: real test suites run under emulation on Apple Silicon.
+EXEC_TIMEOUT_S = 300
+
+# Activate the task env and land in the repo before every command.
+_WRAP = "source /opt/miniconda3/bin/activate testbed 2>/dev/null; cd /testbed 2>/dev/null; "
+
+# `git add -A` respects .gitignore, so build junk stays out; --cached diff
+# then captures edits AND new files the agent created.
+GIT_DIFF_SCRIPT = "git add -A >/dev/null 2>&1; git -c core.fileMode=false diff --cached"
+
+
+class DockerSession:
+    """One long-lived container per task; commands go through docker exec."""
+
+    def __init__(self, container_id: str) -> None:
+        self.container_id = container_id
+
+    @classmethod
+    async def start(cls, image: str, *, platform: str = "linux/amd64") -> "DockerSession":
+        proc = await asyncio.create_subprocess_exec(
+            "docker", "run", "-d", "--platform", platform, image, "tail", "-f", "/dev/null",
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        out, err = await proc.communicate()
+        if proc.returncode != 0:
+            raise RuntimeError(f"docker run failed for {image}: {err.decode(errors='replace').strip()}")
+        return cls(out.decode().strip())
+
+    async def exec(
+        self,
+        script: str,
+        *,
+        stdin: Optional[bytes] = None,
+        timeout: int = EXEC_TIMEOUT_S,
+    ) -> tuple[str, int]:
+        proc = await asyncio.create_subprocess_exec(
+            "docker", "exec", "-i", self.container_id, "bash", "-c", _WRAP + script,
+            stdin=asyncio.subprocess.PIPE if stdin is not None else asyncio.subprocess.DEVNULL,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+        )
+        try:
+            out, _ = await asyncio.wait_for(proc.communicate(input=stdin), timeout=timeout)
+        except asyncio.TimeoutError:
+            proc.kill()
+            # NB: kills the host-side `docker exec`; the in-container process
+            # may linger. Acceptable for v1 — the container dies after the task.
+            return f"[timeout] command exceeded {timeout}s and was killed", 124
+        return out.decode(errors="replace"), proc.returncode or 0
+
+    async def stop(self) -> None:
+        proc = await asyncio.create_subprocess_exec(
+            "docker", "rm", "-f", self.container_id,
+            stdout=asyncio.subprocess.DEVNULL,
+            stderr=asyncio.subprocess.DEVNULL,
+        )
+        await proc.wait()
+
+
+class LocalSession:
+    """Same interface, local bash in a directory. For tests."""
+
+    def __init__(self, workdir: Path) -> None:
+        self.workdir = workdir
+
+    async def exec(
+        self,
+        script: str,
+        *,
+        stdin: Optional[bytes] = None,
+        timeout: int = EXEC_TIMEOUT_S,
+    ) -> tuple[str, int]:
+        proc = await asyncio.create_subprocess_exec(
+            "bash", "-c", script,
+            cwd=self.workdir,
+            stdin=asyncio.subprocess.PIPE if stdin is not None else asyncio.subprocess.DEVNULL,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+        )
+        try:
+            out, _ = await asyncio.wait_for(proc.communicate(input=stdin), timeout=timeout)
+        except asyncio.TimeoutError:
+            proc.kill()
+            return f"[timeout] command exceeded {timeout}s and was killed", 124
+        return out.decode(errors="replace"), proc.returncode or 0
+
+    async def stop(self) -> None:
+        pass
+
+
+async def _bash(session, command: str) -> str:
+    out, code = await session.exec(command)
+    if out.startswith("[timeout]"):
+        return out
+    return _truncate(f"[exit {code}]\n{out}" if out else f"[exit {code}]")
+
+
+async def _read_file(session, path: str, offset=None, limit=None) -> str:
+    q = shlex.quote(path)
+    # cat -n BEFORE sed: numbering happens on the whole file, so a windowed
+    # read shows absolute line numbers the model can feed straight to edits.
+    window = ""
+    if offset or limit:
+        start = max(int(offset or 1), 1)
+        end = str(start + int(limit) - 1) if limit else "$"
+        window = f" | sed -n '{start},{end}p'"
+    out, _code = await session.exec(
+        f'if [ -d {q} ]; then echo "[directory] {path}"; ls -p {q}; '
+        f'elif [ -e {q} ]; then cat -n {q}{window}; '
+        f'else echo "[error] file not found: {path}"; exit 1; fi'
+    )
+    return _truncate(out.rstrip("\n"))
+
+
+async def _write_file(session, path: str, content: str) -> str:
+    q = shlex.quote(path)
+    out, code = await session.exec(
+        f'mkdir -p "$(dirname {q})" && cat > {q}', stdin=content.encode()
+    )
+    if code != 0:
+        return f"[error] write failed: {out.strip()}"
+    return f"[ok] wrote {len(content)} chars to {path}"
+
+
+async def _edit_file(session, path: str, old_string: str, new_string: str) -> str:
+    q = shlex.quote(path)
+    text, code = await session.exec(f"cat {q}")
+    if code != 0:
+        return f"[error] file not found: {path}"
+    n = text.count(old_string)
+    if n == 0:
+        return "[error] old_string not found in file. read the file again — it may have changed."
+    if n > 1:
+        return f"[error] old_string appears {n} times; include more surrounding context to make it unique."
+    return await _write_file(session, path, text.replace(old_string, new_string, 1))
+
+
+async def _grep(session, pattern: str, path: str = ".", include: str | None = None) -> str:
+    inc = f" --include={shlex.quote(include)}" if include else ""
+    # -I skips binaries; junk dirs excluded to match the local tool's behavior.
+    # No pipe here: it would mask grep's exit code (1 = no match, 2 = bad
+    # pattern); matches are capped host-side instead.
+    excludes = " ".join(f"--exclude-dir={d}" for d in (".git", "__pycache__", ".venv", "node_modules"))
+    cmd = f"grep -rnIE {excludes}{inc} -e {shlex.quote(pattern)} {shlex.quote(path)}"
+    out, code = await session.exec(cmd)
+    out = out.rstrip("\n")
+    if code == 1:
+        return "[no matches]"
+    if code != 0:
+        return f"[error] grep failed (bad pattern or path): {out[:200]}"
+    lines = out.splitlines()
+    if len(lines) > GREP_MAX_MATCHES:
+        lines = lines[:GREP_MAX_MATCHES]
+        lines.append(f"[truncated at {GREP_MAX_MATCHES} matches — narrow the pattern]")
+    return _truncate("\n".join(lines))
+
+
+_GLOB_PY = (
+    "import glob, sys; "
+    f"hits = sorted(glob.glob(sys.argv[1], recursive=True))[:{GLOB_MAX_PATHS}]; "
+    "print('\\n'.join(hits) if hits else '[no matches]')"
+)
+
+
+async def _glob(session, pattern: str) -> str:
+    out, code = await session.exec(f"python -c {shlex.quote(_GLOB_PY)} {shlex.quote(pattern)}")
+    if code != 0:
+        return f"[error] glob failed: {out.strip()[:200]}"
+    return _truncate(out.rstrip("\n"))
+
+
+def build_session_registry(session) -> dict[str, Tool]:
+    """The same tools the chat TUI has, executing inside the session."""
+    fns = {
+        "bash": lambda command: _bash(session, command),
+        "read_file": lambda path, offset=None, limit=None: _read_file(session, path, offset, limit),
+        "write_file": lambda path, content: _write_file(session, path, content),
+        "edit_file": lambda path, old_string, new_string: _edit_file(
+            session, path, old_string, new_string
+        ),
+        "grep": lambda pattern, path=".", include=None: _grep(session, pattern, path, include),
+        "glob": lambda pattern: _glob(session, pattern),
+    }
+    return {
+        name: Tool(name=name, schema=SCHEMAS[name], fn=fn, risk=RISK.get(name, "risky"))
+        for name, fn in fns.items()
+    }
+
+
+async def extract_patch(session) -> str:
+    """The agent's work as a git patch — this is the SWE-bench prediction."""
+    out, code = await session.exec(GIT_DIFF_SCRIPT)
+    if code != 0:
+        return ""
+    patch = out
+    if patch and not patch.endswith("\n"):
+        patch += "\n"
+    return patch

rockycode/engine/effort.py

@@ -0,0 +1,46 @@
+"""rocky's effort dial: off | high | xhigh | max.
+
+The dial is rocky-owned and provider-neutral. Providers name their tiers
+differently — DeepSeek V4 accepts only high|max, OpenAI-style models use
+low|medium|high — so each provider gets a clamp table here and the dial value
+never goes on the wire directly. Unknown values pass through untouched, so a
+future model registry can bring its own tiers without touching callers.
+
+`off` never reaches a request body: it means thinking disabled, and callers
+keep the last effort so switching back on restores it.
+"""
+from __future__ import annotations
+
+EFFORT_LEVELS = ("off", "high", "xhigh", "max")  # the user-facing dial
+CLI_EFFORTS = ("high", "xhigh", "max")  # flag values (`off` = --no-thinking)
+
+# Per-reasoning-policy clamp: the dial value → what this provider accepts.
+# DeepSeek V4 knows only high|max (xhigh clamps to max); OpenAI-style maps the
+# rocky dial to low|medium|high. A provider's profile names its policy.
+_CLAMP = {
+    "deepseek": {"high": "high", "xhigh": "max", "max": "max"},
+    "openai": {"high": "medium", "xhigh": "high", "max": "high"},
+}
+
+
+def to_deepseek(effort: str) -> str:  # kept for existing callers
+    return _CLAMP["deepseek"].get(effort, effort)
+
+
+def build_extra_body(thinking: bool, effort: str, reasoning: str = "deepseek") -> dict:
+    """The reasoning fields for this request, shaped by the provider's policy:
+
+    - "deepseek": `{thinking: {type}, reasoning_effort}` (DeepSeek's own knob)
+    - "openai":   bare `{reasoning_effort}` (OpenAI-style), only when thinking
+    - "none":     nothing — the provider has no reasoning param
+
+    The OpenAI SDK carries these via `extra_body`; unknown keys a provider
+    ignores are harmless, but "none" keeps the body clean for strict endpoints.
+    """
+    if not thinking or reasoning == "none":
+        return {"thinking": {"type": "disabled"}} if reasoning == "deepseek" else {}
+    clamp = _CLAMP.get(reasoning, {})
+    tier = clamp.get(effort, effort)
+    if reasoning == "deepseek":
+        return {"thinking": {"type": "enabled"}, "reasoning_effort": tier}
+    return {"reasoning_effort": tier}

rockycode/engine/events.py

@@ -0,0 +1,101 @@
+"""Engine events: the single contract between the agent core and everything
+that watches it (TUI transcript, Rocky pet, trajectory logger, game layer).
+
+Subscribers must tolerate unknown event types — match on isinstance and
+ignore what you don't handle, so new events never break old UIs.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+
+
+class AgentState(str, Enum):
+    """What Rocky is doing right now. Drives the pet animation."""
+
+    IDLE = "idle"
+    THINKING = "thinking"
+    RESPONDING = "responding"
+    TOOL = "tool"
+    COMPACTING = "compacting"
+    AMAZED = "amazed"
+    ERROR = "error"
+
+
+@dataclass
+class Event:
+    pass
+
+
+@dataclass
+class StateChanged(Event):
+    state: AgentState
+
+
+@dataclass
+class TurnStarted(Event):
+    user_message: str
+
+
+@dataclass
+class ThinkingDelta(Event):
+    """A chunk of DeepSeek's reasoning_content stream."""
+
+    text: str
+
+
+@dataclass
+class TextDelta(Event):
+    """A chunk of the assistant's visible reply."""
+
+    text: str
+
+
+@dataclass
+class ToolStarted(Event):
+    call_id: str
+    tool: str
+    args: dict = field(default_factory=dict)
+
+
+@dataclass
+class ToolFinished(Event):
+    call_id: str
+    tool: str
+    output: str
+    ok: bool
+    duration_s: float
+
+
+@dataclass
+class Compacted(Event):
+    """History was rewritten to fit the context window (mid-turn)."""
+
+    strategy: str  # "prune" (old tool outputs stubbed) | "summarize" (state doc rebuild)
+    tokens_before: int  # projected prompt tokens that tripped the limit
+    tokens_after: int  # conservative estimate of the rebuilt history
+    messages_before: int
+    messages_after: int
+
+
+@dataclass
+class ContextReminder(Event):
+    """Soft, non-blocking nudge: context passed the 'model degrades past here'
+    mark (DeepSeek V4 ~50%). The user decides whether to /clear; auto-compaction
+    only kicks in near the window ceiling."""
+
+    pct: float   # fraction of the window currently used
+    window: int
+
+
+@dataclass
+class TurnFinished(Event):
+    """End of a full user turn (all tool round-trips done)."""
+
+    steps: int
+    usage: dict = field(default_factory=dict)
+
+
+@dataclass
+class EngineError(Event):
+    message: str