agent.py

"""Core agent loop: neutral message format, multi-provider streaming."""
from __future__ import annotations

import os
import queue
import threading
import time
import uuid
from pathlib import Path
from dataclasses import dataclass, field
from typing import Any, Generator

from tool_registry import get_tool_schemas, clear_last_output
from tools import execute_tool
import tools as _tools_init  # ensure built-in tools are registered on import
from providers import stream, AssistantTurn, TextChunk, ThinkingChunk, detect_provider
from compaction import maybe_compact
import governance

_SENTINEL = object()

def _interruptible_stream(gen):
    """Run a generator in a daemon thread, yield events via Queue.
    Ctrl+C (KeyboardInterrupt) is always deliverable because the main
    thread only blocks on queue.get(timeout=0.1) — never on a raw socket.
    """
    q: queue.Queue = queue.Queue(maxsize=64)

    def _producer():
        try:
            for event in gen:
                q.put(event)
        except Exception as exc:
            q.put(exc)
        finally:
            q.put(_SENTINEL)

    t = threading.Thread(target=_producer, daemon=True)
    t.start()
    while True:
        try:
            item = q.get(timeout=0.1)
        except queue.Empty:
            continue
        if item is _SENTINEL:
            break
        if isinstance(item, BaseException):
            raise item
        yield item

# ── Re-export event types (used by dulus) ─────────────────────────────────
__all__ = [
    "AgentState", "run",
    "TextChunk", "ThinkingChunk",
    "ToolStart", "ToolEnd", "TurnDone", "PermissionRequest",
]


@dataclass
class AgentState:
    """Mutable session state. messages use the neutral provider-independent format."""
    messages: list = field(default_factory=list)
    total_input_tokens:  int = 0
    total_output_tokens: int = 0
    total_cache_read_tokens: int = 0
    total_cache_creation_tokens: int = 0
    turn_count: int = 0
    _short_mem_counter: int = 0  # +1 per tool turn; nudges short_memory at 10
    # Why the last turn failed, or "" when it didn't. A provider can end a turn
    # with error=True (missing web-auth file, exhausted quota) — run() rolls the
    # history back and breaks, so the caller sees no new assistant message, no
    # TurnDone and no exception. Without this the failure is unobservable and a
    # non-interactive run reports success with whatever message was already last.
    last_error: str = ""
    # This turn's assistant text, or None when the turn never produced one.
    # Read by non-interactive callers that must report exactly what this turn
    # answered. Unlike scanning `messages` it survives mid-run auto-compaction
    # (which rebinds the list and invalidates any index) and cannot pick up the
    # Soul or gold-memory messages that are injected as role="assistant" at boot.
    last_reply: str | None = None


@dataclass
class ToolStart:
    name:   str
    inputs: dict

@dataclass
class ToolEnd:
    name:      str
    result:    str
    permitted: bool = True

@dataclass
class TurnDone:
    input_tokens:  int
    output_tokens: int
    cache_read_tokens:     int = 0
    cache_creation_tokens: int = 0

@dataclass
class PermissionRequest:
    description: str
    granted: bool = False


def _load_short_memory() -> str:
    """Load the current gold short_memory.md for the periodic nudge.

    Always re-seals ``gold: true`` via ``ensure_short_memory`` so a manual edit
    or bad save that stripped the flag cannot leave the scratchpad ungilded.
    """
    try:
        from config import CONFIG_DIR
        path = CONFIG_DIR / "memory" / "short_memory.md"
        try:
            from memory import ensure_short_memory
            ensure_short_memory(force_gold=True)
        except Exception:
            pass
        if path.exists():
            return path.read_text(encoding="utf-8", errors="replace")
    except Exception:
        pass
    return "(empty)"


# ── Agent loop ─────────────────────────────────────────────────────────────

def run(
    user_message: str,
    state: AgentState,
    config: dict,
    system_prompt: str,
    depth: int = 0,
    cancel_check=None,
) -> Generator:
    """
    Multi-turn agent loop (generator).
    Yields: TextChunk | ThinkingChunk | ToolStart | ToolEnd |
            PermissionRequest | TurnDone

    Args:
        depth: sub-agent nesting depth, 0 for top-level
        cancel_check: callable returning True to abort the loop early
    """
    from common import sanitize_text
    # Append user turn in neutral format (sanitize to kill Windows surrogates)
    user_msg: dict[str, Any] = {"role": "user", "content": sanitize_text(user_message)}
    # Attach pending image from /image command if present
    pending_img = config.pop("_pending_image", None)
    if pending_img:
        user_msg["images"] = [pending_img]
    # Attach pending video from /video command (Kimi K2.5/K2.6 only — the
    # provider layer drops it for non-multimodal models). Item may be a bare
    # base64 string or {"data": b64, "mime": "video/mp4"}.
    pending_vid = config.pop("_pending_video", None)
    if pending_vid:
        user_msg["videos"] = [pending_vid]
    
    initial_msg_count = len(state.messages)
    state.last_error = ""     # this turn's verdict starts clean
    state.last_reply = None
    # Same for the shared err() record: without this a stale reason from an
    # earlier turn could be blamed for this turn's outcome.
    try:
        import common as _common
        _common.clear_last_error()
    except Exception:
        pass
    state.messages.append(user_msg)

    # Inject runtime metadata into config so tools (e.g. Agent) can access it.
    config.update({"_depth": depth, "_system_prompt": system_prompt, "_state": state})

    # Governance layer (opt-in via config["governance"]). Built once per session
    # and cached so the budget/ledger persists across turns. None when disabled.
    if "_governance_obj" not in config:
        config["_governance_obj"] = governance.from_config(config)
    _gov = config["_governance_obj"]

    # Bounded retry for empty/stalled turns: a reasoning model (grok-4.5, k3,
    # DeepSeek) can emit only its <thinking> and end the stream — often on
    # finish_reason="length" (truncated) — before any answer or tool call.
    # 0 by design (Kev, 2026-09-03): a silent stall-retry that nudges the model
    # 3× invisibly is exactly what made turns "freeze out of nowhere". On the
    # first empty/stalled turn we now surface a clear message and stop, instead
    # of quietly re-prompting. The user re-sends or adjusts tokens/effort.
    _empty_retries = 0
    _MAX_EMPTY_RETRIES = 0
    while True:
        if cancel_check and cancel_check():
            return
        state.turn_count += 1
        assistant_turn: AssistantTurn | None = None

        # Compact context if approaching window limit
        maybe_compact(state, config)

        _api_source = state.messages

        # Strip GUI/REPL baseline display blobs (soul/gold/welcome). The model
        # already has them via system prompt; leaving them as role:assistant
        # made web-history consolidate drop everything before the last
        # assistant (= the gold blob itself) on first turn.
        try:
            from memory import is_baseline_display_message
            _api_source = [
                m for m in _api_source
                if not is_baseline_display_message(m)
            ]
        except Exception:
            pass

        # Sanitize message contents before sending to API (surrogate safety)
        _safe_messages = []
        for m in _api_source:
            _m = dict(m)
            _c = _m.get("content")
            if isinstance(_c, str):
                _m["content"] = sanitize_text(_c)
            _safe_messages.append(_m)

        # ── Short-memory nudge (every 10 tool turns) ──────────────────────
        # Visible check: after 10 tool turns we surface the current gold
        # short_memory.md to the user, ask if anything should be added, and
        # let the model update it via MemorySave if needed. The nudge resets
        # only after being shown, so it won't spam.
        _effective_system = system_prompt
        try:
            if getattr(state, "_short_mem_counter", 0) >= 10:
                state._short_mem_counter = 0
                _short_mem_text = _load_short_memory()
                _effective_system = system_prompt + (
                    "\n\n[SYSTEM: Short Memory Check — visible to user]\n\n"
                    "You have reached 10 tool turns. The current gold short_memory.md "
                    "is shown below. Read it, then do ONE of the following:\n"
                    "1. If recent conversation contains a new important fact, decision, "
                    "path, correction, or changed project state, call MemorySave with "
                    "name='short_memory' to update the file. Then briefly tell the user "
                    "what you added/updated.\n"
                    "2. If nothing new is worth saving, briefly say: "
                    "'Short memory check: nothing new to save.' and continue normally.\n\n"
                    "Current short_memory.md:\n---\n" + _short_mem_text + "\n---"
                )
        except Exception:
            pass

        # Stream from provider — wrapped so Ctrl+C always fires
        for event in _interruptible_stream(stream(
            model=config["model"],
            system=_effective_system,
            messages=_safe_messages,
            tool_schemas=get_tool_schemas(),
            config=config,
        )):
            if cancel_check and cancel_check():
                return
            if isinstance(event, (TextChunk, ThinkingChunk)):
                yield event
            elif isinstance(event, AssistantTurn):
                assistant_turn = event

        if assistant_turn is None:
            break

        if assistant_turn.error:
            # The turn's text IS the failure reason (provider layer puts it
            # there); record it before the rollback throws away the evidence.
            state.last_error = sanitize_text(assistant_turn.text or "").strip()
            # Rollback: remove anything added during this turn sequence to prevent corrupted history
            while len(state.messages) > initial_msg_count:
                state.messages.pop()
            break

        # Record assistant turn in neutral format
        _reply_text = sanitize_text(assistant_turn.text)
        state.messages.append({
            "role":       "assistant",
            "content":    _reply_text,
            "thinking":   sanitize_text(assistant_turn.thinking) if assistant_turn.thinking else "",
            "tool_calls": assistant_turn.tool_calls,
        })
        state.last_reply = _reply_text

        # Token accounting. Most providers only emit the usage chunk on the
        # FINAL (text) response — on intermediate tool-call responses they send
        # no usage, so assistant_turn.in_tokens is 0 there. That made every tool
        # turn print "+0 in / +0 out" and under-counted /cost. When the provider
        # reports nothing, estimate the real input we sent (windowed messages +
        # tool schemas + system) so the display and totals aren't a false zero.
        if assistant_turn.in_tokens <= 0:
            try:
                from compaction import estimate_tokens
                _est = estimate_tokens(_api_source, model=config.get("model", ""), config=config)
                try:
                    import json as _json
                    from tool_registry import get_all_tools
                    _est += len(_json.dumps([t.schema for t in get_all_tools()])) // 4
                except Exception:
                    pass
                if _est > 0:
                    assistant_turn.in_tokens = _est  # so TurnDone + the toolbar see it too
            except Exception:
                pass

        state.total_input_tokens  += assistant_turn.in_tokens
        state.total_output_tokens += assistant_turn.out_tokens
        c_read = getattr(assistant_turn, "cache_read_tokens", 0)
        c_create = getattr(assistant_turn, "cache_creation_tokens", 0)
        state.total_cache_read_tokens += c_read
        state.total_cache_creation_tokens += c_create
        yield TurnDone(
            assistant_turn.in_tokens,
            assistant_turn.out_tokens,
            cache_read_tokens=c_read,
            cache_creation_tokens=c_create,
        )

        # +1 per tool turn (this is where tokens surface). Drives the
        # short-memory nudge above: at 10 it fires and resets.
        state._short_mem_counter = getattr(state, "_short_mem_counter", 0) + 1

        # Charge the token budget. On warn → notify; on first hard breach →
        # fire on_breach and stop the loop gracefully (protects the wallet).
        if _gov is not None and _gov.ledger is not None:
            cr = _gov.ledger.charge("tokens", assistant_turn.in_tokens + assistant_turn.out_tokens)
            if cr.warned and _gov.hooks:
                _gov.hooks.fire("on_breach", kind="warn", dim=cr.dim, used=cr.used, granted=cr.granted)
            if cr.first_breach:
                if _gov.hooks:
                    _gov.hooks.fire("on_breach", kind="hard", dim=cr.dim, used=cr.used, granted=cr.granted)
                yield TextChunk(
                    f"\n[governance] token budget reached ({cr.used}/{cr.granted}). "
                    "Stopping this run.\n"
                )
                break

        if not assistant_turn.tool_calls:
            # A turn with real text is a genuine final answer — end here.
            # But NO text AND NO tool calls is a STALL, not a completion: the
            # model emitted only reasoning and the stream ended (often
            # finish_reason="length", i.e. truncated) before answering or
            # calling a tool. Breaking here is exactly what makes the agent
            # "suddenly stop" mid-task. Nudge it to actually act, bounded by a
            # retry cap so a persistently-empty model can't loop forever.
            if (assistant_turn.text or "").strip():
                break   # real final answer
            _empty_retries += 1
            if _empty_retries > _MAX_EMPTY_RETRIES:
                _fr = getattr(assistant_turn, "finish_reason", "") or ""
                _cut = " — the message was cut off by the token limit" if _fr == "length" else ""
                yield TextChunk(
                    "\n⚠️  The model produced only reasoning, with no answer and "
                    "no tool call" + _cut + " (usually truncated thinking). "
                    "Stopping — re-send, or try a higher max_tokens / lower "
                    "reasoning effort.\n"
                )
                break
            _fr = getattr(assistant_turn, "finish_reason", "") or ""
            state.messages.append({
                "role": "user",
                "content": (
                    "(You produced only reasoning — no answer and no tool call"
                    + (" and your message was cut off by the token limit"
                       if _fr == "length" else "")
                    + ". Continue now: call the next tool or give your final "
                    "answer.)"
                ),
            })
            continue
        _empty_retries = 0  # a real tool-executing turn — reset the stall streak

        # ── Execute tools ────────────────────────────────────────────────
        for tc in assistant_turn.tool_calls:
            if cancel_check and cancel_check():
                return
            yield ToolStart(tc["name"], tc["input"])

            # ── Governance gate (capabilities + pre_tool hook) ───────────────
            # Runs BEFORE the normal permission gate. A capability denial or a
            # vetoing pre_tool hook short-circuits the tool with a recorded
            # denial, so the model sees why and can adapt.
            if _gov is not None:
                _gov_deny = None
                if _gov.capabilities is not None and not _gov.capabilities.allows_tool(tc["name"]):
                    _gov_deny = f"tool '{tc['name']}' not permitted by capability policy"
                if _gov_deny is None and _gov.hooks is not None:
                    _ok, _why = _gov.hooks.fire(
                        "pre_tool", name=tc["name"], inputs=tc["input"], depth=depth)
                    if not _ok:
                        _gov_deny = _why
                if _gov_deny is not None:
                    result = f"Denied by governance: {_gov_deny}"
                    if _gov.ledger is not None:
                        _gov.ledger.charge("tool_calls", 1)
                    yield ToolEnd(tc["name"], result, False)
                    state.messages.append({
                        "role": "tool", "tool_call_id": tc["id"],
                        "name": tc["name"], "content": result,
                    })
                    continue

            # Permission gate
            permitted = _check_permission(tc, config)
            if not permitted:
                if config.get("permission_mode") == "plan":
                    # Plan mode: silently deny writes (no user prompt)
                    permitted = False
                else:
                    req = PermissionRequest(description=_permission_desc(tc))
                    yield req
                    permitted = req.granted

            if not permitted:
                if config.get("permission_mode") == "plan":
                    plan_file = config.get("_plan_file", "")
                    result = (
                        f"[Plan mode] Write operations are blocked except to the plan file: {plan_file}\n"
                        "Finish your analysis and write the plan to the plan file. "
                        "The user will run /plan done to exit plan mode and begin implementation."
                    )
                else:
                    result = "Denied: user rejected this operation"
            else:
                config["_turn_count"] = state.turn_count
                result = execute_tool(
                    tc["name"], tc["input"],
                    permission_mode="accept-all",  # already gate-checked above
                    config=config,
                )
                # time.sleep(1) # Removed delay as requested

            yield ToolEnd(tc["name"], result, permitted)

            # Governance: charge the tool-call budget + fire post_tool hook
            # (observational — audit, notify, metrics). Never blocks here.
            if _gov is not None:
                if _gov.ledger is not None:
                    _gov.ledger.charge("tool_calls", 1)
                if _gov.hooks is not None:
                    _gov.hooks.fire("post_tool", name=tc["name"], inputs=tc["input"],
                                    result=result, permitted=permitted, depth=depth)

            # Determine what the USER actually saw rendered, based on tool type +
            # auto_show + verbose. Inject a SYSTEM HINT when user saw nothing useful,
            # so the model can decide whether to PrintToConsole the content.
            from tool_registry import is_display_only
            display = is_display_only(tc["name"])
            auto_show_on = config.get("auto_show", True) if config else True
            verbose_on   = config.get("verbose",   False) if config else False

            # User-visibility rules (must match dulus.py print_tool_end logic):
            #   display tool   → user saw full output IF auto_show ON
            #   other tool     → user saw 500-char preview IF verbose ON
            if display:
                user_saw = auto_show_on
            else:
                user_saw = verbose_on

            if display and user_saw:
                # Display-only tool the user already saw: replace with placeholder to save tokens.
                result_summary = f"[Display output shown to user: {len(result)} characters]"
            else:
                result_summary = result

            # Inject the hint when (a) user did not see the content, (b) it's not a
            # purely internal tool, and (c) the call did not error out.
            _internal_tools = {
                "SearchLastOutput", "ReadJob", "TmuxOffload", "MemorySearch",
                "PrintToConsole", "AskUserQuestion", "Write", "Edit",
            }
            if (not user_saw
                    and tc["name"] not in _internal_tools
                    and not result.startswith(("Error", "Denied"))):
                state_desc = []
                if not auto_show_on: state_desc.append("auto_show OFF")
                if not verbose_on:   state_desc.append("verbose OFF")
                state_str = " + ".join(state_desc) or "user-display suppressed"
                result_summary = (
                    f"{result_summary}\n\n"
                    f"[SYSTEM HINT — {state_str}]\n"
                    "The user did NOT see this output rendered (only a brief [OK] line). "
                    "If this content is meant for the user (Bash output they asked for, file "
                    "they wanted to read, ASCII art they requested), call "
                    "PrintToConsole(content=...) or PrintToConsole(file_path=...) NOW to "
                    "show it. If this was just internal investigation, ignore this hint."
                )

            # Record tool result in neutral format
            state.messages.append({
                "role":         "tool",
                "tool_call_id": tc["id"],
                "name":         tc["name"],
                "content":      sanitize_text(result_summary),
            })

            # ── Truncation Awareness Reminder ────────────────────────────────
            # If the tool output was truncated, the model only saw a fragment.
            # Inject a hard reminder so it cannot honestly claim "X is missing"
            # without first using SearchLastOutput to actually search the file.
            # Skip this check for SearchLastOutput itself to avoid loops.
            if (tc["name"] != "SearchLastOutput"
                    and "[TRUNCATED" in result):
                try:
                    path = Path.home() / ".dulus" / "last_tool_output.txt"
                    if path.exists():
                        full_size = path.stat().st_size
                        seen_size = len(result)
                        if full_size > seen_size:
                            with path.open("rb") as _f:
                                full_lines = sum(1 for _ in _f)
                            state.messages.append({
                                "role": "user",
                                "content": (
                                    "[SYSTEM REMINDER — TRUNCATED OUTPUT]\n"
                                    f"The previous tool result was TRUNCATED. You only saw "
                                    f"~{seen_size} characters out of {full_size} total "
                                    f"({full_lines} lines). The full output is saved.\n\n"
                                    "RULE: You CANNOT claim that any item, font, key, name, "
                                    "match, or piece of data is missing, absent, or does not "
                                    "exist based on what you just saw. You only have a fragment.\n\n"
                                    "BEFORE answering the user's question, you MUST call:\n"
                                    "  SearchLastOutput(pattern=\"<the thing the user asked about>\")\n"
                                    "to verify against the full saved output. If the user asked "
                                    "about a specific name, search for that exact name. If they "
                                    "asked for a count or a list, use SearchLastOutput() with no "
                                    "pattern to get the full summary.\n\n"
                                    "Do not answer from memory or guess. Search first."
                                ),
                            })
                except Exception:
                    pass


# ── Helpers ───────────────────────────────────────────────────────────────

def _check_permission(tc: dict, config: dict) -> bool:
    """Return True if operation is auto-approved (no need to ask user)."""
    perm_mode = config.get("permission_mode", "auto")
    name = tc["name"]

    # Plan mode tools are always auto-approved
    if name in ("EnterPlanMode", "ExitPlanMode"):
        return True

    if perm_mode == "accept-all":
        return True
    if perm_mode == "manual":
        return False   # always ask

    if perm_mode == "plan":
        # Allow writes ONLY to the plan file
        if name in ("Write", "Edit"):
            plan_file = config.get("_plan_file", "")
            target = tc["input"].get("file_path", "")
            if plan_file and target and \
               os.path.normpath(target) == os.path.normpath(plan_file):
                return True
            return False
        if name == "NotebookEdit":
            return False
        if name == "Bash":
            from tools import _is_safe_bash
            return _is_safe_bash(tc["input"].get("command", ""))
        return True  # reads are fine

    # "auto" mode: only ask for writes and non-safe bash
    if name in ("Read", "Glob", "Grep", "WebFetch", "WebSearch"):
        return True
    if name == "Bash":
        from tools import _is_safe_bash
        return _is_safe_bash(tc["input"].get("command", ""))
    return False   # Write, Edit → ask


def _permission_desc(tc: dict) -> str:
    name = tc["name"]
    inp  = tc["input"]
    if name == "Bash":   return f"Run: {inp.get('command', '')}"
    if name == "Write":  return f"Write to: {inp.get('file_path', '')}"
    if name == "Edit":   return f"Edit: {inp.get('file_path', '')}"
    return f"{name}({list(inp.values())[:1]})"