loop.py

"""The ReAct loop: stream a model turn, run any tool calls, repeat until done.

One public entry point: run_turn(). Same loop serves the REPL and one-shot mode
(and, later, eval mode — see SEAM at the bottom).
"""
import json
import os
import re
import sys
import time

import httpx

import guardrails
import memory
import tools

# --- Exp 2: cross-turn tool-result cache -------------------------------------
# The per-turn loop-breaker resets each turn, so a model re-running the SAME
# expensive read across turns (measured: 6x identical grep) slips past it. This
# session-scoped cache serves the prior result for an idempotent repeat call,
# invalidated whenever anything mutates. Gate: PIT_TOOLCACHE=0 disables.
_result_cache = {}
_CACHEABLE_READ = {"read_file", "grep", "glob", "web_search", "web_fetch"}
# NOTE: bare '>' was REMOVED — it false-matched comparison operators in analysis
# code (`if count > 400`), wiping the cache every call (Exp 2 bug). Redirects to
# files aren't caught now; acceptable since the cache is opt-in and session reads
# are usually stable. Explicit mutation verbs + write_file/edit_file cover it.
_MUTATION_RE = re.compile(
    r"\bmv\b|\bcp\b|\bmkdir\b|\btee\b|sed\s+-i|\binstall\b|"
    r"\bnpm\b|\bpip\b|\bgit\s+(?:commit|add|push|merge|checkout|reset)|"
    r"\bINSERT\s+INTO\b|\bUPDATE\s+\w|\bDROP\s+TABLE\b", re.IGNORECASE)


def _is_mutating(name, args):
    if name in ("write_file", "edit_file"):
        return True
    if name in ("bash", "powershell", "python_eval"):
        return bool(_MUTATION_RE.search(args.get("command") or args.get("code") or ""))
    return False


def _is_cacheable(name, args):
    if name in _CACHEABLE_READ:
        return True
    if name in ("bash", "powershell", "python_eval"):
        return not _is_mutating(name, args)   # read-only shell/py is cacheable
    return False

STEP_CAP = 40           # max model<->tool round-trips per turn
TOKEN_WARN = 100_000    # est. tokens before we nag
PRUNE_TOOL_RESULTS = False  # KILLED by Experiment 0 (cache_probe.py): stubbing a
# tool result in place is a mid-context edit -> 0% KV-cache reuse, ~4s re-prefill
# every time it fires. The 90% fix is spill-to-file (append-only handle), not
# in-place editing. See CONTEXT_DESIGN_v2.md.


def _est_tokens(messages) -> int:
    """Rough char/4 estimate across all message text."""
    n = 0
    for m in messages:
        c = m.get("content")
        if isinstance(c, str):
            n += len(c)
        for tc in m.get("tool_calls", []) or []:
            n += len(tc.get("function", {}).get("arguments", ""))
    return n // 4


def _dim(s: str):
    sys.stdout.write(f"\033[2m{s}\033[0m\n")
    sys.stdout.flush()


def _payload(model_cfg, messages):
    p = {
        "model": model_cfg["model"],
        "messages": messages,
        "tools": tools.active_tools(),   # profile-filtered (see tools.PROFILES)
        "tool_choice": "auto",
        "stream": True,
    }
    p.update(model_cfg.get("extra_body", {}))
    return p


def _stream_turn(client, url, headers, payload, echo):
    """POST one streamed completion. Returns (content, tool_calls list).

    Accumulates OpenAI streamed tool_call deltas by index: id + function.name
    arrive once, function.arguments arrive in fragments to be concatenated.
    """
    content_parts = []
    accs = {}           # index -> {id, name, args}
    started_text = False
    with client.stream("POST", url, headers=headers, json=payload,
                       timeout=None) as resp:
        if resp.status_code != 200:
            body = resp.read().decode("utf-8", "replace")
            raise RuntimeError(f"model API {resp.status_code}: {body[:400]}")
        for line in resp.iter_lines():
            if not line or not line.startswith("data:"):
                continue
            data = line[5:].strip()
            if data == "[DONE]":
                break
            try:
                delta = json.loads(data)["choices"][0]["delta"]
            except (json.JSONDecodeError, KeyError, IndexError):
                continue
            piece = delta.get("content")
            if piece:
                content_parts.append(piece)
                if echo:
                    if not started_text:
                        started_text = True
                    sys.stdout.write(piece)
                    sys.stdout.flush()
            for tc in delta.get("tool_calls") or []:
                idx = tc.get("index", 0)
                a = accs.setdefault(idx, {"id": None, "name": "", "args": ""})
                if tc.get("id"):
                    a["id"] = tc["id"]
                fn = tc.get("function") or {}
                if fn.get("name"):
                    a["name"] = fn["name"]
                if fn.get("arguments"):
                    a["args"] += fn["arguments"]
    if echo and started_text:
        sys.stdout.write("\n")
        sys.stdout.flush()
    tool_calls = []
    for idx in sorted(accs):
        a = accs[idx]
        if not a["name"]:
            continue
        tool_calls.append({
            "id": a["id"] or f"call_{idx}",
            "type": "function",
            "function": {"name": a["name"], "arguments": a["args"] or "{}"},
        })
    return "".join(content_parts), tool_calls


def _recall(query, messages):
    """Exp 3: keyword search over THIS session's own prior findings — the model's
    earlier assistant conclusions and tool-result previews — so it can retrieve
    what it already learned instead of re-running the query. Retrieval side of the
    re-fetch problem (Exp 2 showed caching the source was the wrong fix). Precision
    over recall (research): return only the few best-overlapping snippets."""
    terms = {w for w in re.findall(r"[a-z0-9]{3,}", (query or "").lower())}
    if not terms:
        return "[recall] give a query with some keywords."
    scored = []
    # skip the last 2 messages (the current recall call itself + its assistant)
    for m in messages[:-2]:
        role = m.get("role")
        if role not in ("assistant", "tool"):
            continue
        text = m.get("content") or ""
        if not text or text.startswith("[pruned"):
            continue
        low = text.lower()
        hits = sum(1 for t in terms if t in low)
        if hits:
            # assistant conclusions are higher-value than raw tool output
            weight = hits * (2 if role == "assistant" else 1)
            scored.append((weight, role, text))
    if not scored:
        return (f"[recall] nothing in this session matches {sorted(terms)}. "
                "You have not looked at this yet — go fetch it.")
    scored.sort(key=lambda x: -x[0])
    out = [f"[recall: what you already found this session about "
           f"'{query}' — use this instead of re-fetching]"]
    for _, role, text in scored[:3]:
        snip = text.strip()[:600]
        out.append(f"\n({'your finding' if role == 'assistant' else 'a tool result'}): {snip}")
    return "\n".join(out)


def _exec_tool(call, echo, messages=None):
    name = call["function"]["name"]
    raw = call["function"]["arguments"] or "{}"
    try:
        args = json.loads(raw)
    except json.JSONDecodeError:
        return f"[error] could not parse arguments as JSON: {raw[:200]}"
    if name not in tools.ACTIVE:                  # profile-disabled or hallucinated
        return (f"[error] tool '{name}' is not enabled in this run. Available: "
                + ", ".join(t["function"]["name"] for t in tools.active_tools()))
    if name == "recall":                          # Exp 3: handled here (needs history)
        if echo:
            _dim(f"[recall] {str(args.get('query',''))[:80]}")
        return _recall(args.get("query", ""), messages or [])
    fn = tools.TOOL_FUNCS.get(name)
    if fn is None:
        return f"[error] unknown tool: {name}"
    if echo:
        label = args.get("command") or args.get("path") or \
            args.get("pattern") or ""
        _dim(f"[{name}] {str(label)[:100]}")

    cache_on = os.environ.get("PIT_TOOLCACHE") == "1"   # opt-in: Exp 2 inconclusive
    sig = name + "|" + (raw or "")
    if cache_on and _is_cacheable(name, args) and sig in _result_cache:
        if echo:
            _dim("  -> [cross-turn cache hit — not re-running]")
        cached = ("[cross-turn cache: you already ran this exact call earlier "
                  "this session and nothing has changed since — cached result "
                  "below; do not keep re-running it]\n" + _result_cache[sig])
        return tools.spill_result(name, cached)

    try:
        result = fn(**args)
    except TypeError as e:
        return f"[error] bad arguments for {name}: {e}"
    except Exception as e:                       # noqa: BLE001
        return f"[error] {name} raised: {e}"

    if cache_on:
        if _is_mutating(name, args):
            _result_cache.clear()                # state changed → stale reads out
        elif _is_cacheable(name, args) and not result.startswith("[error]"):
            _result_cache[sig] = result
    # Spill only UNBOUNDED shell output (bash/powershell/python_eval). Exp 1: a big
    # grep dump spilled = −62% win. Exp 7/8: spilling read_file broke extraction (the
    # model needs the whole doc; read_file already has its own cap+pagination). grep/
    # glob/web_fetch are already capped. And never re-spill a .scratch read (cascade).
    if name in ("bash", "powershell", "python_eval") and ".scratch" not in (raw or ""):
        result = tools.spill_result(name, result)
    if echo:
        first = result.splitlines()[0] if result else ""
        n = len(result.splitlines())
        tag = "error" if result.startswith("[error]") else f"{n} line(s)"
        _dim(f"  -> {tag}: {first[:100]}")
    return result


def run_turn(messages, model_cfg, log=None, echo=True, step_cap=STEP_CAP):
    """Drive one agentic turn to completion. Mutates `messages` in place, so the
    caller keeps full history. `log` is a callable(msg_dict) for jsonl logging.
    Returns the final assistant text."""
    url = model_cfg["base_url"].rstrip("/") + "/chat/completions"
    headers = {"Content-Type": "application/json"}
    key_env = model_cfg.get("api_key_env")
    if key_env and os.environ.get(key_env):
        headers["Authorization"] = f"Bearer {os.environ[key_env]}"

    call_counts = {}        # (name|args) -> times seen, for loop-breaking
    # progress guardrails (guardrails.py): catch reworded dead ends / confabulation
    # / general no-progress spinning that the exact-repeat breaker misses. Gate:
    # PIT_GUARDRAILS=0 disables (baseline arm of the A/B).
    mon = None
    if os.environ.get("PIT_GUARDRAILS", "1") != "0":
        req = next((m["content"] for m in reversed(messages)
                    if m.get("role") == "user" and isinstance(m.get("content"), str)), "")
        mon = guardrails.ProgressMonitor(req)
    t0 = time.perf_counter()
    gen_seconds = 0.0       # wall-clock spent in model calls (the cost metric)
    n_steps = 0
    with httpx.Client() as client:
        for _ in range(step_cap):
            n_steps += 1
            if PRUNE_TOOL_RESULTS:
                saved = memory.prune_stale_tool_results(messages)
                if saved > 4000 and echo:
                    _dim(f"  [context: pruned ~{saved//1000}k chars of stale "
                         "tool output]")
            if _est_tokens(messages) > TOKEN_WARN and echo:
                _dim(f"[context ~{_est_tokens(messages)//1000}k tokens — "
                     "consider /compact]")
            _tg = time.perf_counter()
            content, tool_calls = _stream_turn(
                client, url, headers, _payload(model_cfg, messages), echo)
            gen_seconds += time.perf_counter() - _tg
            assistant = {"role": "assistant", "content": content or None}
            if tool_calls:
                assistant["tool_calls"] = tool_calls
            messages.append(assistant)
            if log:
                log(assistant)
            if not tool_calls:
                if log:                          # cost telemetry for this turn
                    log({"role": "meta", "kind": "timing",
                         "gen_seconds": round(gen_seconds, 2),
                         "wall_seconds": round(time.perf_counter() - t0, 2),
                         "steps": n_steps,
                         "ctx_tokens": _est_tokens(messages)})
                return content
            pending = None      # progress-guardrail action for this batch
            for call in tool_calls:
                name = call["function"]["name"]
                raw = call["function"]["arguments"] or ""
                sig = name + "|" + raw
                call_counts[sig] = call_counts.get(sig, 0) + 1
                if call_counts[sig] >= 3:
                    # Same exact call 3+ times: stop executing, break the loop.
                    result = (f"[harness] You have already run this exact call "
                              f"{call_counts[sig]} times and gotten the same result. "
                              "Repeating it will not help. Change approach: if a "
                              "search keeps finding nothing, the thing may be ABSENT "
                              "(and may need to be added). If you have enough "
                              "information, stop searching and act on it now. "
                              "If you are genuinely stuck after multiple different "
                              "approaches, consult_oracle is available.")
                    if echo:
                        _dim(f"  [loop-break: '{name}' x{call_counts[sig]} suppressed]")
                else:
                    result = _exec_tool(call, echo, messages)
                tmsg = {"role": "tool", "tool_call_id": call["id"],
                        "content": result}
                messages.append(tmsg)
                if log:
                    log(tmsg)
                if mon:                          # progress guardrail (post-result)
                    act = mon.record(name, raw, result)
                    if act == "stop" or (act == "reground" and pending != "stop"):
                        pending = act

            if pending:     # every tool_call has its response; safe to intervene now
                if log:
                    log({"role": "meta", "kind": "guardrail", "action": pending,
                         "offender": mon.last_offender, "step": n_steps})
                if pending == "stop":
                    if echo:
                        _dim(f"  [guardrail: hard stop — {mon.last_offender or 'no progress'}]")
                    return mon.message
                if echo:
                    _dim("  [guardrail: re-grounding — stalled, nudging back on task]")
                reg = {"role": "user", "content": mon.message}
                messages.append(reg)
                if log:
                    log(reg)
    if echo:
        _dim(f"[step cap {step_cap} reached — stopping]")
    return "[stopped: step cap reached]"

# SEAM: eval mode reuses run_turn() unchanged — a runner copies a fixture dir,
# builds messages=[system, task_prompt], calls run_turn(echo=False), then diffs
# the workdir and scores the trace. No loop changes needed.