nanobot/nanobot/agent/tools/self.py

"""MyTool: runtime state inspection and configuration for the agent loop."""

from __future__ import annotations

import time
from typing import Any

from loguru import logger

from nanobot.agent.subagent import SubagentStatus
from nanobot.agent.tools.base import Tool
from nanobot.agent.tools.context import ContextAware, RequestContext
from nanobot.agent.tools.runtime_state import RuntimeState
from nanobot.config.schema import Base


class MyToolConfig(Base):
    """Self-inspection tool configuration."""
    enable: bool = True
    allow_set: bool = False


def _has_real_attr(obj: Any, key: str) -> bool:
    """Check if obj has a real (explicitly set) attribute, not auto-generated by mock."""
    if isinstance(obj, dict):
        return key in obj
    d = getattr(obj, "__dict__", None)
    if d is not None and key in d:
        return True
    for cls in type(obj).__mro__:
        if key in cls.__dict__:
            return True
    return False


class MyTool(Tool, ContextAware):
    """Check and set the agent loop's runtime configuration."""

    _plugin_discoverable = False  # Requires AgentLoop reference; registered manually
    config_key = "my"

    @classmethod
    def config_cls(cls):
        return MyToolConfig

    @classmethod
    def enabled(cls, ctx: Any) -> bool:
        return ctx.config.my.enable

    BLOCKED = frozenset({
        # Core infrastructure
        "bus", "provider", "_running", "tools",
        # Config management
        "_runtime_vars",
        # Subsystems
        "runner", "sessions", "consolidator",
        "dream", "auto_compact", "context", "commands",
        # Sensitive runtime state (credentials, message routing, task tracking)
        "_mcp_servers", "_mcp_stacks", "_pending_queues",
        "_session_locks", "_active_tasks", "_background_tasks",
        # Security boundaries (inspect + modify both blocked)
        "restrict_to_workspace", "channels_config",
        "_concurrency_gate", "_unified_session", "_extra_hooks",
    })

    READ_ONLY = frozenset({
        "subagents",  # observable but replacing it would break the system
        "_current_iteration",  # updated by runner only
        "exec_config",  # inspect allowed (e.g. check sandbox), modify blocked
        "web_config",  # inspect allowed (e.g. check enable), modify blocked
    })

    _DENIED_ATTRS = frozenset({
        "__class__", "__dict__", "__bases__", "__subclasses__", "__mro__",
        "__init__", "__new__", "__reduce__", "__getstate__", "__setstate__",
        "__del__", "__call__", "__getattr__", "__setattr__", "__delattr__",
        "__code__", "__globals__", "func_globals", "func_code",
        "__wrapped__", "__closure__",
    })

    # Sub-field names that are sensitive regardless of parent path
    _SENSITIVE_NAMES = frozenset({
        "api_key", "secret", "password", "token", "credential",
        "private_key", "access_token", "refresh_token", "auth",
    })

    @classmethod
    def _is_sensitive_field_name(cls, name: str) -> bool:
        lowered = name.lower()
        return lowered in cls._SENSITIVE_NAMES or any(
            part in cls._SENSITIVE_NAMES for part in lowered.split("_")
        )

    RESTRICTED: dict[str, dict[str, Any]] = {
        "max_iterations":        {"type": int, "min": 1,   "max": 100},
        "context_window_tokens": {"type": int, "min": 4096, "max": 1_000_000},
        "model":                 {"type": str, "min_len": 1},
    }

    _MAX_RUNTIME_KEYS = 64

    def __init__(self, runtime_state: RuntimeState, modify_allowed: bool = True) -> None:
        self._runtime_state = runtime_state
        self._modify_allowed = modify_allowed
        self._channel = ""
        self._chat_id = ""

    def __deepcopy__(self, memo: dict[int, Any]) -> MyTool:
        cls = self.__class__
        result = cls.__new__(cls)
        memo[id(self)] = result
        result._runtime_state = self._runtime_state
        result._modify_allowed = self._modify_allowed
        result._channel = self._channel
        result._chat_id = self._chat_id
        return result

    def set_context(self, ctx: RequestContext) -> None:
        self._channel = ctx.channel
        self._chat_id = ctx.chat_id

    @property
    def name(self) -> str:
        return "my"

    @property
    def description(self) -> str:
        base = (
            "Check and set your own runtime state.\n"
            "Actions: check, set.\n"
            "- check (no key): full config overview — start here.\n"
            "- check (key): drill into a value. Dot-paths allowed "
            "(e.g. '_last_usage.prompt_tokens', 'web_config.enable').\n"
            "- set (key, value): change config or store notes in your scratchpad. "
            "Scratchpad keys persist across turns but not restarts.\n"
            "Key values: _current_iteration (current progress), "
            "max_iterations - _current_iteration = remaining iterations.\n"
            "Note: web_config and exec_config are readable but read-only.\n"
            "\n"
            "When to use:\n"
            "- User asks about your model, settings, or token usage → check that key.\n"
            "- A tool fails or behaves unexpectedly → check the related config to diagnose.\n"
            "- User asks you to remember a preference for this session → set to store it in your scratchpad.\n"
            "- About to start a large task → check context_window_tokens and max_iterations first."
        )
        if not self._modify_allowed:
            base += "\nREAD-ONLY MODE: set is disabled."
        else:
            base += (
                "\nIMPORTANT: Before setting state, predict the potential impact. "
                "If the operation could cause crashes or instability "
                "(e.g. changing model), warn the user first."
            )
        return base

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "action": {
                    "type": "string",
                    "enum": ["check", "set"],
                    "description": "Action to perform",
                },
                "key": {
                    "type": "string",
                    "description": "Dot-path for check/set. Examples: 'max_iterations', 'workspace', 'provider_retry_mode'. "
                    "For check without key, shows all config values.",
                },
                "value": {"description": "New value (for set). Type must match target (int for max_iterations/context_window_tokens, str for model)."},
            },
            "required": ["action"],
        }

    def _audit(self, action: str, detail: str) -> None:
        session = f"{self._channel}:{self._chat_id}" if self._channel else "unknown"
        logger.info("self.{} | {} | session:{}", action, detail, session)

    # ------------------------------------------------------------------
    # Path resolution
    # ------------------------------------------------------------------

    def _resolve_path(self, path: str) -> tuple[Any, str | None]:
        parts = path.split(".")
        obj = self._runtime_state
        for part in parts:
            if part in self._DENIED_ATTRS or part.startswith("__"):
                return None, f"'{part}' is not accessible"
            if part in self.BLOCKED:
                return None, f"'{part}' is not accessible"
            if part.lower() in self._SENSITIVE_NAMES:
                return None, f"'{part}' is not accessible"
            try:
                if isinstance(obj, dict):
                    if part in obj:
                        obj = obj[part]
                    else:
                        return None, f"'{part}' not found in dict"
                else:
                    obj = getattr(obj, part)
            except (KeyError, AttributeError) as e:
                return None, f"'{part}' not found: {e}"
        return obj, None

    @staticmethod
    def _validate_key(key: str | None, label: str = "key") -> str | None:
        if not key or not key.strip():
            return f"Error: '{label}' cannot be empty or whitespace"
        return None

    # ------------------------------------------------------------------
    # Smart formatting
    # ------------------------------------------------------------------

    @staticmethod
    def _format_status(st: SubagentStatus, indent: str = "  ") -> str:
        elapsed = time.monotonic() - st.started_at
        tool_summary = ", ".join(
            f"{e.get('name', '?')}({e.get('status', '?')})" for e in st.tool_events[-5:]
        ) or "none"
        lines = [
            f"{indent}phase: {st.phase}, iteration: {st.iteration}, elapsed: {elapsed:.1f}s",
            f"{indent}tools: {tool_summary}",
            f"{indent}usage: {st.usage or 'n/a'}",
        ]
        if st.error:
            lines.append(f"{indent}error: {st.error}")
        if st.stop_reason:
            lines.append(f"{indent}stop_reason: {st.stop_reason}")
        return "\n".join(lines)

    @staticmethod
    def _format_value(val: Any, key: str = "") -> str:
        if isinstance(val, SubagentStatus):
            header = f"Subagent [{val.task_id}] '{val.label}'"
            detail = MyTool._format_status(val, "  ")
            return f"{header}\n  task: {val.task_description}\n{detail}"
        # SubagentManager: delegate to its _task_statuses dict
        if hasattr(val, "_task_statuses") and isinstance(val._task_statuses, dict):
            return MyTool._format_value(val._task_statuses, key)
        if isinstance(val, dict) and val and isinstance(next(iter(val.values())), SubagentStatus):
            prefix = f"{key}: " if key else ""
            lines = [f"{prefix}{len(val)} subagent(s):"]
            for tid, st in val.items():
                detail = MyTool._format_status(st, "    ")
                lines.append(f"  [{tid}] '{st.label}'\n{detail}")
            return "\n".join(lines)
        if hasattr(val, "tool_names"):
            return f"tools: {len(val.tool_names)} registered — {val.tool_names}"
        # Scalar types — repr is fine
        if isinstance(val, (str, int, float, bool, type(None))):
            r = repr(val)
            return f"{key}: {r}" if key else r
        # Dict — small: show content; large: show keys for dot-path navigation
        if isinstance(val, dict):
            ks = list(val.keys())
            if not ks:
                return f"{key}: {{}}" if key else "{}"
            if len(ks) <= 5:
                r = repr(val)
                if len(r) <= 200:
                    return f"{key}: {r}" if key else r
            preview = ", ".join(str(k) for k in ks[:15])
            suffix = ", ..." if len(ks) > 15 else ""
            return f"{key}: {{{preview}{suffix}}}" if key else f"{{{preview}{suffix}}}"
        # List/tuple — count for large, repr for small
        if isinstance(val, (list, tuple)):
            if len(val) > 20:
                return f"{key}: [{len(val)} items]" if key else f"[{len(val)} items]"
            r = repr(val)
            return f"{key}: {r}" if key else r
        # Complex object — small Pydantic models: show values; others: show field names for navigation
        cls_name = type(val).__name__
        model_fields = getattr(type(val), "model_fields", None)
        if model_fields:
            fields = list(model_fields.keys())
            if len(fields) <= 8:
                # Small config objects: show field=value pairs
                pairs = []
                for f in fields:
                    fv = getattr(val, f, "?")
                    if MyTool._is_sensitive_field_name(f):
                        continue
                    if isinstance(fv, (str, int, float, bool, type(None))):
                        pairs.append(f"{f}={fv!r}")
                    else:
                        pairs.append(f"{f}=<{type(fv).__name__}>")
                preview = ", ".join(pairs)
                return f"{key}: {preview}" if key else preview
        else:
            fields = [a for a in getattr(val, "__dict__", {}) if not a.startswith("__")]
        if fields:
            preview = ", ".join(str(f) for f in fields[:20])
            suffix = ", ..." if len(fields) > 20 else ""
            return f"{key}: <{cls_name}> [{preview}{suffix}]" if key else f"<{cls_name}> [{preview}{suffix}]"
        r = repr(val)
        return f"{key}: {r}" if key else r

    # ------------------------------------------------------------------
    # Action dispatch
    # ------------------------------------------------------------------

    async def execute(
        self,
        action: str,
        key: str | None = None,
        value: Any = None,
        **_kwargs: Any,
    ) -> str:
        if action in ("inspect", "check"):
            return self._inspect(key)
        if not self._modify_allowed:
            return "Error: set is disabled (tools.my.allow_set is false)"
        if action in ("modify", "set"):
            return self._modify(key, value)
        return f"Unknown action: {action}"

    # -- inspect --

    def _inspect(self, key: str | None) -> str:
        if not key:
            return self._inspect_all()
        top = key.split(".")[0]
        if top in self._DENIED_ATTRS or top.startswith("__"):
            return f"Error: '{top}' is not accessible"
        obj, err = self._resolve_path(key)
        if err:
            # "scratchpad" alias for _runtime_vars
            if key == "scratchpad":
                rv = self._runtime_state._runtime_vars
                return self._format_value(rv, "scratchpad") if rv else "scratchpad is empty"
            # Fallback: check _runtime_vars for simple keys stored by modify
            if "." not in key and key in self._runtime_state._runtime_vars:
                return self._format_value(self._runtime_state._runtime_vars[key], key)
            return f"Error: {err}"
        # Guard against mock auto-generated attributes
        if "." not in key and not _has_real_attr(self._runtime_state, key):
            if key in self._runtime_state._runtime_vars:
                return self._format_value(self._runtime_state._runtime_vars[key], key)
            return f"Error: '{key}' not found"
        return self._format_value(obj, key)

    def _inspect_all(self) -> str:
        state = self._runtime_state
        parts: list[str] = []
        # RESTRICTED keys
        for k in self.RESTRICTED:
            parts.append(self._format_value(getattr(state, k, None), k))
        # Other useful top-level keys shown in description
        for k in ("workspace", "provider_retry_mode", "max_tool_result_chars", "_current_iteration", "web_config", "exec_config", "subagents"):
            if _has_real_attr(state, k):
                parts.append(self._format_value(getattr(state, k, None), k))
        # Token usage
        usage = state._last_usage
        if usage:
            parts.append(self._format_value(usage, "_last_usage"))
        rv = state._runtime_vars
        if rv:
            parts.append(self._format_value(rv, "scratchpad"))
        return "\n".join(parts)

    # -- modify --

    def _modify(self, key: str | None, value: Any) -> str:
        if err := self._validate_key(key):
            return err
        top = key.split(".")[0]
        if top in self.BLOCKED or top in self._DENIED_ATTRS or top.startswith("__") or top.lower() in self._SENSITIVE_NAMES:
            self._audit("modify", f"BLOCKED {key}")
            return f"Error: '{key}' is protected and cannot be modified"
        if top in self.READ_ONLY:
            self._audit("modify", f"READ_ONLY {key}")
            return f"Error: '{key}' is read-only and cannot be modified"
        if "." in key:
            parent_path, leaf = key.rsplit(".", 1)
            if leaf in self._DENIED_ATTRS or leaf.startswith("__"):
                self._audit("modify", f"BLOCKED leaf '{leaf}'")
                return f"Error: '{leaf}' is not accessible"
            if leaf.lower() in self._SENSITIVE_NAMES:
                self._audit("modify", f"BLOCKED sensitive leaf '{leaf}'")
                return f"Error: '{leaf}' is not accessible"
            parent, err = self._resolve_path(parent_path)
            if err:
                return f"Error: {err}"
            if isinstance(parent, dict):
                parent[leaf] = value
            else:
                setattr(parent, leaf, value)
            self._audit("modify", f"{key} = {value!r}")
            return f"Set {key} = {value!r}"
        if key in self.RESTRICTED:
            return self._modify_restricted(key, value)
        return self._modify_free(key, value)

    def _modify_restricted(self, key: str, value: Any) -> str:
        spec = self.RESTRICTED[key]
        expected = spec["type"]
        if expected is int and isinstance(value, bool):
            return f"Error: '{key}' must be {expected.__name__}, got bool"
        if not isinstance(value, expected):
            try:
                value = expected(value)
            except (ValueError, TypeError):
                return f"Error: '{key}' must be {expected.__name__}, got {type(value).__name__}"
        old = getattr(self._runtime_state, key)
        if "min" in spec and value < spec["min"]:
            return f"Error: '{key}' must be >= {spec['min']}"
        if "max" in spec and value > spec["max"]:
            return f"Error: '{key}' must be <= {spec['max']}"
        if "min_len" in spec and len(str(value)) < spec["min_len"]:
            return f"Error: '{key}' must be at least {spec['min_len']} characters"
        setattr(self._runtime_state, key, value)
        if key == "max_iterations" and hasattr(self._runtime_state, "_sync_subagent_runtime_limits"):
            self._runtime_state._sync_subagent_runtime_limits()
        self._audit("modify", f"{key}: {old!r} -> {value!r}")
        return f"Set {key} = {value!r} (was {old!r})"

    def _modify_free(self, key: str, value: Any) -> str:
        if _has_real_attr(self._runtime_state, key):
            old = getattr(self._runtime_state, key)
            if isinstance(old, (str, int, float, bool)):
                old_t, new_t = type(old), type(value)
                if old_t is float and new_t is int:
                    pass  # int → float coercion allowed
                elif old_t is not new_t:
                    self._audit(
                        "modify",
                        f"REJECTED type mismatch {key}: expects {old_t.__name__}, got {new_t.__name__}",
                    )
                    return f"Error: '{key}' expects {old_t.__name__}, got {new_t.__name__}"
            setattr(self._runtime_state, key, value)
            self._audit("modify", f"{key}: {old!r} -> {value!r}")
            return f"Set {key} = {value!r} (was {old!r})"
        if callable(value):
            self._audit("modify", f"REJECTED callable {key}")
            return "Error: cannot store callable values"
        err = self._validate_json_safe(value)
        if err:
            self._audit("modify", f"REJECTED {key}: {err}")
            return f"Error: {err}"
        if key not in self._runtime_state._runtime_vars and len(self._runtime_state._runtime_vars) >= self._MAX_RUNTIME_KEYS:
            self._audit("modify", f"REJECTED {key}: max keys ({self._MAX_RUNTIME_KEYS}) reached")
            return f"Error: scratchpad is full (max {self._MAX_RUNTIME_KEYS} keys). Remove unused keys first."
        old = self._runtime_state._runtime_vars.get(key)
        self._runtime_state._runtime_vars[key] = value
        self._audit("modify", f"scratchpad.{key}: {old!r} -> {value!r}")
        return f"Set scratchpad.{key} = {value!r}"

    @classmethod
    def _validate_json_safe(cls, value: Any, depth: int = 0) -> str | None:
        if depth > 10:
            return "value nesting too deep (max 10 levels)"
        if isinstance(value, (str, int, float, bool, type(None))):
            return None
        if isinstance(value, list):
            for i, item in enumerate(value):
                if err := cls._validate_json_safe(item, depth + 1):
                    return f"list[{i}] contains {err}"
            return None
        if isinstance(value, dict):
            for k, v in value.items():
                if not isinstance(k, str):
                    return f"dict key must be str, got {type(k).__name__}"
                if err := cls._validate_json_safe(v, depth + 1):
                    return f"dict key '{k}' contains {err}"
            return None
        return f"unsupported type {type(value).__name__}"