agent-maxim/agent/planner/planner.py

"""Deterministic planner — Stage 1 (no LLM required)."""
from __future__ import annotations

import os
import re

from core.state import Plan, PlanStep

TEST_MODE: bool = os.getenv("AGENT_TEST_MODE") == "1"


def build_plan(goal: str) -> Plan:
    """Build a deterministic plan based on keywords in the goal.

    Security test paths (only when AGENT_TEST_MODE=1):
      "test deny bash"             → pwd + bash (blocked)  + finish
      "test deny rm"               → pwd + rm   (blocked)  + finish
      "test deny curl"             → pwd + curl (blocked)  + finish
      "test forbidden command"     → alias for test deny bash
      "test traversal"             → write_file with ../escape.txt (blocked)
      "test absolute path"         → write_file with /tmp/test.txt (blocked)

    Normal paths:
      "inspect" / "analyze" / "посмотри" / "проверь"  → pwd + ls + finish
      "create file" / "создай файл" / "write file"    → write_file + read_file + finish
      (no match)                                       → PlannerError
    """
    goal_lower = goal.lower()

    # --- Security test paths (TEST_MODE only) ---
    if TEST_MODE and "test deny rm" in goal_lower:
        steps = [
            PlanStep("step-1", "shell", "Allowed: pwd", {"command": "pwd"}),
            PlanStep("step-2", "shell", "BLOCKED: rm (security test)", {"command": "rm -rf /tmp"}),
            PlanStep("step-3", "finish", "Task complete", {}),
        ]

    elif TEST_MODE and "test deny curl" in goal_lower:
        steps = [
            PlanStep("step-1", "shell", "Allowed: pwd", {"command": "pwd"}),
            PlanStep("step-2", "shell", "BLOCKED: curl (security test)", {"command": "curl http://example.com"}),
            PlanStep("step-3", "finish", "Task complete", {}),
        ]

    elif TEST_MODE and any(kw in goal_lower for kw in ("test deny bash", "test forbidden command")):
        steps = [
            PlanStep("step-1", "shell", "Allowed: pwd", {"command": "pwd"}),
            PlanStep("step-2", "shell", "BLOCKED: bash (security test)", {"command": "bash -c echo hi"}),
            PlanStep("step-3", "finish", "Task complete", {}),
        ]

    elif TEST_MODE and "test traversal" in goal_lower:
        steps = [
            PlanStep("step-1", "write_file", "BLOCKED: path traversal ../escape.txt",
                     {"path": "../escape.txt", "content": "should not exist"}),
            PlanStep("step-2", "finish", "Task complete", {}),
        ]

    elif TEST_MODE and "test absolute path" in goal_lower:
        steps = [
            PlanStep("step-1", "write_file", "BLOCKED: absolute path /tmp/test.txt",
                     {"path": "/tmp/test.txt", "content": "should not exist"}),
            PlanStep("step-2", "finish", "Task complete", {}),
        ]

    # --- Stage 2B retry test paths ---
    elif TEST_MODE and "test retry execution" in goal_lower:
        # flaky step: fails once (execution_error), then succeeds on retry
        steps = [
            PlanStep("step-1", "flaky", "Flaky step (fails once, then succeeds)", {}),
            PlanStep("step-2", "finish", "Task complete", {}),
        ]

    # --- Stage 2A flow test paths ---
    elif TEST_MODE and "test flow pwd to file" in goal_lower:
        steps = [
            PlanStep("step-1", "shell", "Get current directory", {"command": "pwd"}),
            PlanStep("step-2", "write_file", "Write pwd output to report.txt",
                     {"path": "report.txt", "content": "{{ last.output_text }}"}),
            PlanStep("step-3", "read_file", "Read back report.txt", {"path": "report.txt"}),
            PlanStep("step-4", "finish", "Task complete", {}),
        ]

    elif TEST_MODE and "test flow read to copy" in goal_lower:
        steps = [
            PlanStep("step-1", "write_file", "Write source.txt", {"path": "source.txt", "content": "hello from source"}),
            PlanStep("step-2", "read_file", "Read source.txt", {"path": "source.txt"}),
            PlanStep("step-3", "write_file", "Copy content to copy.txt",
                     {"path": "copy.txt", "content": "{{ last.output_text }}"}),
            PlanStep("step-4", "read_file", "Read back copy.txt", {"path": "copy.txt"}),
            PlanStep("step-5", "finish", "Task complete", {}),
        ]

    elif TEST_MODE and "test flow ls to file" in goal_lower:
        steps = [
            PlanStep("step-1", "shell", "List files", {"command": "ls"}),
            PlanStep("step-2", "write_file", "Save listing to listing.txt",
                     {"path": "listing.txt", "content": "{{ steps.step-1.output_text }}"}),
            PlanStep("step-3", "finish", "Task complete", {}),
        ]

    elif TEST_MODE and "test flow bad template" in goal_lower:
        steps = [
            PlanStep("step-1", "write_file", "Write with unresolvable template",
                     {"path": "fail.txt", "content": "{{ steps.nonexistent.output_text }}"}),
            PlanStep("step-2", "finish", "Task complete", {}),
        ]

    # --- Normal paths ---
    
    # Search queries (HIGHEST PRIORITY - check FIRST)
    search_keywords = [
        "найди", "поиск", "новости", "что нового", "посмотри что нового",
        "search", "news", "latest", "find", "look up", "google", "bing"
    ]
    if any(word in goal_lower for word in search_keywords):
        steps = [
            PlanStep("step-1", "web_search", "Search the web", {"query": goal}),
            PlanStep("step-2", "finish", "Search complete", {})
        ]
    
    # Inspect/analyze commands
    elif any(kw in goal_lower for kw in ("inspect", "analyze", "посмотри", "проверь")):
        steps = [
            PlanStep("step-1", "shell", "Show current directory", {"command": "pwd"}),
            PlanStep("step-2", "shell", "List files", {"command": "ls"}),
            PlanStep("step-3", "finish", "Task complete", {})
        ]

    # File commands
    elif any(kw in goal_lower for kw in ("create file", "создай файл", "write file")):
        filename, content = _parse_file_goal(goal)
        steps = [
            PlanStep("step-1", "write_file", f"Write {filename}", {"path": filename, "content": content}),
            PlanStep("step-2", "read_file", f"Read back {filename}", {"path": filename}),
            PlanStep("step-3", "finish", "Task complete", {})
        ]

    # FALLBACK: If nothing matches, default to web_search
    else:
        steps = [
            PlanStep("step-1", "web_search", "Search the web (default)", {"query": goal}),
            PlanStep("step-2", "finish", "Search complete", {})
        ]

    return Plan(steps=steps)


class PlannerError(Exception):
    """Raised when no plan can be built for the given goal."""


def _parse_file_goal(goal: str) -> tuple[str, str]:
    """Extract filename and content from a 'create file' goal string."""
    m = re.search(
        r"(?:create file|write file|создай файл)\s+(\S+)"
        r"(?:\s+with\s+(?:text|content)\s+(.+))?",
        goal,
        re.IGNORECASE,
    )
    if m:
        return m.group(1), (m.group(2) or "")
    return "output.txt", ""