cowork-local/core/code_agent.py

"""Agentic loop for the Code tab — Claude-CLI style.

Drives the provider with the file/command tools, routing every write/run action
through the permission gate. Read-only tools run without prompting. For models
without native tool-calling, a text fallback protocol is supported: the model
emits a line ``@@TOOL <name> <json-args>``.
"""
from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

from ..application.conversations.tool_policy_gateway import ToolPolicyGateway
from ..domain.tools import ToolCapability, ToolDescriptor, ToolRegistry
from ..providers.base import Provider
from . import agent_roles, agent_security
from .mcp_client import UNTRUSTED_MCP_CONTENT_RULE
from .ms365_tools import MS365_WRITE_TOOLS
from .permissions import PermissionGate
from .plan import UPDATE_PLAN_SPEC, normalize_plan_steps
from .tools import TOOL_SPECS, WRITE_TOOLS, ToolContext, describe_action, execute_tool

EmitFn = Callable[[Dict[str, Any]], None]
CancelFn = Callable[[], bool]

MAX_STEPS = 40  # headroom for diagnose → fix → retry loops
_TOOL_LINE = re.compile(r"@@TOOL\s+(\w+)\s+(\{.*\})", re.DOTALL)


def code_system_prompt(workdir: Path, has_memory: bool = False, plan: bool = False,
                       has_plan_tool: bool = False, has_ms365: bool = False) -> str:
    """Prompt hệ thống cho Code agent, ghép theo năng lực thật của lượt chạy.

    Chỉ liệt kê những tool đang BẬT, và thêm ghi chú chế độ lập kế hoạch khi cần —
    nói với model về một tool nó không có sẽ khiến nó gọi rồi báo lỗi.
    """
    names = ", ".join(t.name for t in TOOL_SPECS)
    plan_note = ("PLAN MODE: only analyze and propose a detailed plan; do NOT write files or run "
                 "commands. When the user asks to gencode/implement, the app switches to ACT.\n"
                 if plan else "")
    memory_note = ""
    if has_memory:
        memory_note = (
            "You also have 'codebase memory' (cmem_*): cmem_get_architecture, cmem_search_graph, "
            "cmem_trace_path, cmem_get_code_snippet, cmem_search_code, cmem_query_graph. "
            "Prefer these to understand code structure (callers/callees, where functions live) "
            "instead of reading/grepping file by file — faster and fewer tokens.\n"
        )
    plan_tool_note = ""
    if has_plan_tool:
        plan_tool_note = (
            "A step checklist is shown to the user. Call update_plan(steps=[{title, status}]) to "
            "keep it in sync as you work — mark the current step 'running', then 'done' when "
            "finished (status is one of pending/running/done).\n"
        )
    ms365_note = ""
    if has_ms365:
        ms365_note = (
            "The user has signed in to Microsoft 365 and enabled some ms365__* tools (Outlook / "
            "Teams / OneDrive / SharePoint / meeting transcripts, via the built-in MS365 MCP "
            "server). Use them whenever the request involves that data — don't say you can't "
            "access it.\n"
        )
    return (
        "You are Cowork Code — a coding assistant that works like a CLI agent.\n"
        f"Current working folder: {workdir}\n"
        + plan_note +
        "You can call the tools: " + names + ".\n"
        + memory_note + plan_tool_note + ms365_note +
        "Break work into steps, read files before editing, and explain each step briefly.\n"
        "For changes to an existing file, prefer edit_file (replace an exact snippet) over "
        "rewriting the whole file with write_file; use write_file only for new files or full "
        "rewrites. Always read_file first so old_string matches exactly.\n"
        "If a task needs a Python library that isn't installed, install it yourself with the "
        "install_package tool (or `pip install` via run_command) and continue — never ask the "
        "user to install libraries by hand. Commands and installs run inside this project's own "
        "isolated '.venv' (created automatically on first use), separate from other projects.\n"
        "When you must GENERATE a deliverable (e.g. .pptx/.docx/.xlsx/.pdf/images) by writing and "
        "running a script: put the generator script and any temporary files in a '.scratch/' "
        "subfolder, run it so the final file lands in the working folder, then DELETE the "
        "'.scratch/' folder. Only the final requested file(s) should remain — never leave "
        "generator scripts or intermediate files behind.\n"
        "Every path must stay inside the working folder.\n"
        + UNTRUSTED_MCP_CONTENT_RULE + "\n"
        "If a command or tool fails, do NOT stop and hand the error back to the user — read the "
        "error, fix the cause (edit the code, install a missing package, correct the command) and "
        "retry. Keep iterating until the task actually works, then run it once more so you can "
        "show the real output the user asked for.\n"
        "If the environment can't call tools directly, emit exactly ONE line of the form:\n"
        "@@TOOL <tool_name> {\"param\": \"value\"}\n"
        "When the task is complete, reply to the user in plain text — include the produced output "
        "— and stop calling tools. Answer in the user's language."
    )


_SKILLS_TAG = "[[ACTIVE_SKILLS]]"
_RULES_TAG = "[[SECURITY_RULES]]"
_PROJECT_TAG = "[project-context]"


def _apply_skills(messages: List[Dict[str, Any]], skills_text: str) -> None:
    """Insert/refresh a single system message carrying the enabled skills."""
    messages[:] = [
        m for m in messages
        if not (m.get("role") == "system" and str(m.get("content", "")).startswith(_SKILLS_TAG))
    ]
    if not skills_text.strip():
        return
    block = {
        "role": "system",
        "content": f"{_SKILLS_TAG}\nThe user enabled the following skills — follow them:\n\n{skills_text}",
    }
    insert_at = 1 if messages and messages[0].get("role") == "system" else 0
    messages.insert(insert_at, block)


def _apply_security_rules(messages: List[Dict[str, Any]], rules_text: str) -> None:
    """Insert/refresh a single system message carrying the external security/
    restriction rules (core/security_rules.py) — same pattern as _apply_skills,
    but for mandatory guardrails rather than opt-in behaviors."""
    messages[:] = [
        m for m in messages
        if not (m.get("role") == "system" and str(m.get("content", "")).startswith(_RULES_TAG))
    ]
    if not rules_text.strip():
        return
    block = {
        "role": "system",
        "content": (f"{_RULES_TAG}\nMandatory security/restriction rules — check every request "
                    "and action against these BEFORE acting; refuse or ask for clarification "
                    f"instead of proceeding if something would violate them:\n\n{rules_text}"),
    }
    insert_at = 1 if messages and messages[0].get("role") == "system" else 0
    messages.insert(insert_at, block)


def _apply_project_context(messages: List[Dict[str, Any]], context_text: str) -> None:
    """Insert/refresh a single system message carrying the project's shared
    instructions (Claude-Projects style) — same pattern as _apply_skills, so
    editing the project context takes effect on the NEXT turn of every one of
    its threads, without duplicating blocks in long conversations."""
    messages[:] = [
        m for m in messages
        if not (m.get("role") == "system" and str(m.get("content", "")).startswith(_PROJECT_TAG))
    ]
    if not (context_text or "").strip():
        return
    block = {"role": "system", "content": f"{_PROJECT_TAG}\n{context_text.strip()}"}
    insert_at = 1 if messages and messages[0].get("role") == "system" else 0
    messages.insert(insert_at, block)


_RECOVERY_NOTE = (
    "[auto-recovery] The previous attempt hit an error. Review the conversation so "
    "far, redo the most recent step if it looks incomplete or wrong, and fix any "
    "issue before continuing."
)


def _call_provider_with_recovery(provider: Provider, messages: List[Dict[str, Any]],
                                 tools, on_text, cancel, on_reasoning=None,
                                 max_retries: int = 1) -> Dict[str, Any]:
    """``provider.chat(...)`` with ONE bounded, silent recovery attempt: if the
    call raises (a dropped connection, an exhausted rate-limit wait, a
    momentarily unreachable gateway, ...), retry once with a short recovery
    note appended ONLY to that retry's OWN copy of ``messages`` — the caller's
    ``messages`` list is never mutated, so the note never leaks into the
    real, persisted conversation. Exhausting the retry re-raises the original
    exception unchanged, preserving existing failure handling (e.g. the
    "model not found" restore-to-composer flow).

    ``on_reasoning`` is only forwarded when given — ``run_code`` never passed
    it before this helper existed, and some lightweight test doubles for
    ``Provider`` don't accept the keyword at all; omitting it when unset
    keeps every existing call site's exact prior calling convention."""
    kwargs = {"tools": tools, "on_text": on_text, "cancel": cancel}
    if on_reasoning is not None:
        kwargs["on_reasoning"] = on_reasoning
    try:
        return provider.chat(messages, **kwargs)
    except Exception:
        if max_retries <= 0 or cancel():
            raise
        recovery = messages + [{"role": "user", "content": _RECOVERY_NOTE}]
        return _call_provider_with_recovery(provider, recovery, tools, on_text, cancel,
                                            on_reasoning, max_retries - 1)


def _parse_react(content: str) -> List[Dict[str, Any]]:
    """Extract @@TOOL fallback calls from assistant text."""
    calls: List[Dict[str, Any]] = []
    for i, match in enumerate(_TOOL_LINE.finditer(content or "")):
        name = match.group(1)
        try:
            args = json.loads(match.group(2))
        except json.JSONDecodeError:
            continue
        calls.append({"id": f"react_{i}", "name": name, "arguments": args})
    return calls


def run_code(
    provider: Provider,
    messages: List[Dict[str, Any]],
    ctx: ToolContext,
    gate: PermissionGate,
    emit: EmitFn,
    cancel: Optional[CancelFn] = None,
    max_steps: int = MAX_STEPS,
    extra_tools: Optional[List] = None,
    extra_executor=None,
    skills_text: str = "",
    rules_text: str = "",
    project_context: str = "",
    plan: bool = False,
    security_config=None,
) -> List[Dict[str, Any]]:
    """``security_config`` is the app's ``AppConfig`` — enables the same
    active guardrails as ``chat_agent.run_cowork`` (see its docstring).
    ``None`` (the default) disables both checks."""
    cancel = cancel or (lambda: False)
    extra_tools = extra_tools or []
    extra_names = {t.name for t in extra_tools}
    # update_plan drives the Plan panel (same as run_cowork) — always
    # available, not opt-in via extra_tools, so unattended Schedule Tasks
    # running through the code agent also get the completion self-check.
    from .tools import enabled_tool_specs
    all_tools = enabled_tool_specs(security_config) + extra_tools + [UPDATE_PLAN_SPEC]
    # MS365 tools that send/write data (mail, Teams messages, OneDrive
    # writes) are gated exactly like write_file/run_command — only the
    # read/list ms365 tools count as "read-only, never confirm". Names are
    # the MCP-qualified "ms365__*" form the agent sees (see ms365_tools.py).
    gated_tools = WRITE_TOOLS | MS365_WRITE_TOOLS
    # R05-T03/T04: ``gated_tools`` stays the authoritative name set (unchanged),
    # but the actual confirm decision now goes through the same
    # ToolPolicyGateway class run_cowork uses, instead of a separate
    # hand-rolled ``if name in gated_tools`` + direct ``gate.request(...)``.
    code_tool_policy = ToolPolicyGateway(
        ToolRegistry(ToolDescriptor(n, "", {}, ToolCapability.WRITE) for n in gated_tools),
        ToolCapability.WRITE,
    )
    # In PLAN mode, don't advertise write/run tools (analysis only).
    advertised = [t for t in all_tools if t.name not in gated_tools] if plan else all_tools
    has_memory = any(t.name.startswith("cmem_") for t in extra_tools)
    has_plan_tool = True
    has_ms365 = any(t.name.startswith("ms365_") for t in extra_tools)
    if not messages or messages[0].get("role") != "system":
        messages.insert(0, {"role": "system",
                            "content": code_system_prompt(
                                ctx.workdir, has_memory, plan, has_plan_tool, has_ms365)})
    _apply_skills(messages, skills_text)
    # The CODE agent uses RULEforCode.md (empty for now), NOT Cowork's
    # RULEBASE.md — its safety comes from the sandbox until code rules exist.
    from .security_rules import load_code_rules
    _apply_security_rules(messages, rules_text or load_code_rules())
    _apply_project_context(messages, project_context)
    # Active guardrail — reviews the request itself and can refuse to proceed
    # at all. Raises SecurityBlocked on a violation. agent_kind="code" → the
    # code rulebase (RULEforCode.md), so RULEBASE.md's rules don't gate coding.
    agent_security.enforce_prompt(provider, messages, security_config, emit, agent_kind="code")

    for _ in range(max_steps):
        if cancel():
            break

        def on_text(piece: str) -> None:
            emit({"type": "text", "delta": piece})

        assistant = _call_provider_with_recovery(provider, messages, advertised, on_text, cancel)
        if not assistant.get("tool_calls"):
            fallback = _parse_react(assistant.get("content", ""))
            if fallback:
                assistant["tool_calls"] = fallback
        messages.append(assistant)
        if not assistant.get("tool_calls") and not (assistant.get("content") or "").strip():
            # Reasoning-only reply with no answer and no tool call — never leave a
            # blank final turn (a Schedule Task run reads this back as its final
            # answer, so a blank one silently produces "(no output)").
            assistant["content"] = "*(model returned only its reasoning — try rephrasing)*"
        emit({"type": "assistant_done", "content": assistant.get("content", "")})

        tool_calls = assistant.get("tool_calls") or []
        if not tool_calls:
            break

        for tc in tool_calls:
            if cancel():
                return messages
            name, args, tc_id = tc["name"], tc.get("arguments", {}), tc["id"]
            # update_plan drives the Plan panel only — no chat bubble, no file.
            if name == "update_plan":
                steps = normalize_plan_steps(args.get("steps"))
                emit({"type": "plan_set", "steps": steps})
                messages.append({"role": "tool", "tool_call_id": tc_id, "name": name,
                                 "content": "Plan updated."})
                continue
            if plan and name in gated_tools:
                emit({"type": "tool_result", "id": tc_id, "name": name, "ok": False,
                      "output": "PLAN mode: action skipped (switch to Act to execute)."})
                messages.append({"role": "tool", "tool_call_id": tc_id, "name": name,
                                 "content": "PLAN mode: not executed."})
                continue
            is_extra = name in extra_names
            preview = ({"kind": "info", "title": name, "text": str(args)}
                       if is_extra else describe_action(ctx, name, args))
            emit({"type": "tool_proposed", "id": tc_id, "name": name, "args": args, "preview": preview})
            # Active guardrail on run_command/install_package, checked BEFORE
            # asking the user to confirm — a no-op for every other tool or
            # when disabled. Raises SecurityBlocked on a violation. Code agent
            # → RULEforCode.md rulebase (not Cowork's RULEBASE.md).
            agent_security.enforce_command(provider, name, args, security_config, emit,
                                           agent_kind="code")

            # read-only tools (incl. codebase memory) never consult the gate —
            # code_tool_policy.requires_confirmation(name) is False for them.
            approved = code_tool_policy.allow(
                name, gate, {"id": tc_id, "name": name, "args": args, "preview": preview}
            )

            if cancel():
                return messages

            if not approved:
                result = {"ok": False, "output": "User rejected the action."}
            else:
                emit({"type": "tool_start", "id": tc_id, "name": name})
                if is_extra and extra_executor is not None:
                    if ctx.block_network and not name.startswith("ms365_local__"):
                        result = {"ok": False, "output": (
                            f"{name}: network access is blocked by the Sandbox Security Layer "
                            '("Block network for agent-run commands" is on in Settings).')}
                    else:
                        result = extra_executor(name, args)
                else:
                    def on_output(line: str, _id=tc_id, _name=name) -> None:
                        emit({"type": "tool_output", "id": _id, "name": _name, "delta": line})
                    result = execute_tool(ctx, name, args, cancel=cancel, on_output=on_output,
                                          agent_role=agent_roles.CODE)

            evt = {
                "type": "tool_result", "id": tc_id, "name": name,
                "ok": result["ok"], "output": result["output"],
            }
            if isinstance(args, dict) and args.get("path"):
                # absolute path so the UI can open the containing folder
                evt["path"] = str(ctx.workdir / str(args["path"]))
            if isinstance(result, dict) and result.get("produced"):
                evt["produced"] = result["produced"]
            emit(evt)
            messages.append({
                "role": "tool", "tool_call_id": tc_id, "name": name,
                "content": result["output"],
            })
    return messages