usestrix__strix/strix/tools/notes/tools.py

"""Per-run notes storage — mirrored to {state_dir}/notes.json."""

from __future__ import annotations

import asyncio
import json
import logging
import tempfile
import threading
import uuid
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from agents import RunContextWrapper, function_tool

from strix.tools.nullish import clean_optional


logger = logging.getLogger(__name__)


_notes_storage: dict[str, dict[str, Any]] = {}
_VALID_NOTE_CATEGORIES = ["general", "findings", "methodology", "questions", "plan", "wiki"]
_notes_lock = threading.RLock()
_DEFAULT_CONTENT_PREVIEW_CHARS = 280
_NOTE_ID_GENERATION_ATTEMPTS = 1024

_notes_path: Path | None = None


def _caller_identity(ctx: RunContextWrapper) -> tuple[str | None, str | None]:
    """Return the (agent_id, agent_name) of the agent invoking this tool."""
    inner = ctx.context if isinstance(ctx.context, dict) else {}
    raw_agent_id = inner.get("agent_id")
    agent_id = raw_agent_id if isinstance(raw_agent_id, str) else None
    agent_name: str | None = None
    coordinator = inner.get("coordinator")
    if agent_id is not None and coordinator is not None:
        names = getattr(coordinator, "names", {})
        if isinstance(names, dict):
            raw_agent_name = names.get(agent_id)
            agent_name = raw_agent_name if isinstance(raw_agent_name, str) else None
    return agent_id, agent_name


def _generate_note_id() -> str | None:
    for _ in range(_NOTE_ID_GENERATION_ATTEMPTS):
        note_id = uuid.uuid4().hex[:6]
        if note_id not in _notes_storage:
            return note_id
    return None


def hydrate_notes_from_disk(state_dir: Path) -> None:
    global _notes_path  # noqa: PLW0603
    _notes_path = state_dir / "notes.json"
    with _notes_lock:
        _notes_storage.clear()
        if not _notes_path.exists():
            return
        try:
            data = json.loads(_notes_path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            logger.exception(
                "notes.json at %s is unreadable; starting with empty notes",
                _notes_path,
            )
            return
        if not isinstance(data, dict):
            return
        _notes_storage.update(
            {
                nid: note
                for nid, note in data.items()
                if isinstance(nid, str) and isinstance(note, dict)
            }
        )
        logger.info(
            "notes hydrated from %s (%d note(s))",
            _notes_path,
            len(_notes_storage),
        )


def _persist() -> None:
    path = _notes_path
    if path is None:
        return
    try:
        payload = json.dumps(_notes_storage, ensure_ascii=False, default=str)
        path.parent.mkdir(parents=True, exist_ok=True)
        with (
            _notes_lock,
            tempfile.NamedTemporaryFile(
                mode="w",
                encoding="utf-8",
                dir=str(path.parent),
                prefix=f".{path.name}.",
                suffix=".tmp",
                delete=False,
            ) as tmp,
        ):
            tmp.write(payload)
            tmp_path = Path(tmp.name)
        tmp_path.replace(path)
    except Exception:
        logger.exception("notes persist to %s failed", path)


def _filter_notes(
    category: str | None = None,
    tags: list[str] | None = None,
    search_query: str | None = None,
) -> list[dict[str, Any]]:
    category = clean_optional(category)
    search_query = clean_optional(search_query)

    filtered: list[dict[str, Any]] = []
    for note_id, note in _notes_storage.items():
        if category and note.get("category") != category:
            continue
        if tags:
            note_tags = note.get("tags", [])
            if not any(tag in note_tags for tag in tags):
                continue
        if search_query:
            search_lower = search_query.lower()
            title_match = search_lower in note.get("title", "").lower()
            content_match = search_lower in note.get("content", "").lower()
            if not (title_match or content_match):
                continue
        entry = note.copy()
        entry["note_id"] = note_id
        filtered.append(entry)
    filtered.sort(key=lambda x: x.get("created_at", ""), reverse=True)
    return filtered


def _mark_authorship(
    entry: dict[str, Any], note: dict[str, Any], caller_agent_id: str | None
) -> dict[str, Any]:
    """Attach the note's author and flag whether the caller wrote it."""
    agent_name = note.get("agent_name")
    if agent_name:
        entry["agent_name"] = agent_name
    agent_id = note.get("agent_id")
    if agent_id:
        entry["agent_id"] = agent_id
    if caller_agent_id is not None and agent_id == caller_agent_id:
        entry["by_you"] = True
    return entry


def _to_note_listing_entry(
    note: dict[str, Any],
    *,
    include_content: bool = False,
    caller_agent_id: str | None = None,
) -> dict[str, Any]:
    entry = {
        "note_id": note.get("note_id"),
        "title": note.get("title", ""),
        "category": note.get("category", "general"),
        "tags": note.get("tags", []),
        "created_at": note.get("created_at", ""),
        "updated_at": note.get("updated_at", ""),
    }
    content = str(note.get("content", ""))
    if include_content:
        entry["content"] = content
    elif content:
        if len(content) > _DEFAULT_CONTENT_PREVIEW_CHARS:
            entry["content_preview"] = f"{content[:_DEFAULT_CONTENT_PREVIEW_CHARS].rstrip()}..."
        else:
            entry["content_preview"] = content
    return _mark_authorship(entry, note, caller_agent_id)


def _create_note_impl(
    title: str,
    content: str,
    category: str = "general",
    tags: list[str] | None = None,
    agent_id: str | None = None,
    agent_name: str | None = None,
) -> dict[str, Any]:
    with _notes_lock:
        try:
            if not title or not title.strip():
                return {"success": False, "error": "Title cannot be empty", "note_id": None}
            if not content or not content.strip():
                return {"success": False, "error": "Content cannot be empty", "note_id": None}
            if category not in _VALID_NOTE_CATEGORIES:
                return {
                    "success": False,
                    "error": (
                        f"Invalid category. Must be one of: {', '.join(_VALID_NOTE_CATEGORIES)}"
                    ),
                    "note_id": None,
                }

            note_id = _generate_note_id()
            if note_id is None:
                return {
                    "success": False,
                    "error": "Failed to generate a unique note ID",
                    "note_id": None,
                }

            timestamp = datetime.now(UTC).isoformat()
            note = {
                "title": title.strip(),
                "content": content.strip(),
                "category": category,
                "tags": tags or [],
                "created_at": timestamp,
                "updated_at": timestamp,
            }
            if agent_id:
                note["agent_id"] = agent_id
            if agent_name:
                note["agent_name"] = agent_name
            _notes_storage[note_id] = note
        except (ValueError, TypeError) as e:
            return {"success": False, "error": f"Failed to create note: {e}", "note_id": None}
        else:
            _persist()
            return {
                "success": True,
                "note_id": note_id,
                "message": f"Note '{title}' created successfully",
                "total_count": len(_notes_storage),
            }


def _list_notes_impl(
    category: str | None = None,
    tags: list[str] | None = None,
    search: str | None = None,
    include_content: bool = False,
    caller_agent_id: str | None = None,
) -> dict[str, Any]:
    with _notes_lock:
        try:
            filtered = _filter_notes(category=category, tags=tags, search_query=search)
            notes = [
                _to_note_listing_entry(
                    n, include_content=include_content, caller_agent_id=caller_agent_id
                )
                for n in filtered
            ]
        except (ValueError, TypeError) as e:
            return {
                "success": False,
                "error": f"Failed to list notes: {e}",
                "notes": [],
                "filtered_count": 0,
                "total_count": 0,
            }
        return {
            "success": True,
            "notes": notes,
            "filtered_count": len(notes),
            "total_count": len(_notes_storage),
        }


def _get_note_impl(note_id: str, caller_agent_id: str | None = None) -> dict[str, Any]:
    with _notes_lock:
        try:
            if not note_id or not note_id.strip():
                return {"success": False, "error": "Note ID cannot be empty", "note": None}
            note = _notes_storage.get(note_id)
            if note is None:
                return {
                    "success": False,
                    "error": f"Note with ID '{note_id}' not found",
                    "note": None,
                }
            note_with_id = note.copy()
            note_with_id["note_id"] = note_id
            _mark_authorship(note_with_id, note, caller_agent_id)
        except (ValueError, TypeError) as e:
            return {"success": False, "error": f"Failed to get note: {e}", "note": None}
        else:
            return {"success": True, "note": note_with_id}


def _update_note_impl(
    note_id: str,
    title: str | None = None,
    content: str | None = None,
    tags: list[str] | None = None,
) -> dict[str, Any]:
    with _notes_lock:
        try:
            if note_id not in _notes_storage:
                return {"success": False, "error": f"Note with ID '{note_id}' not found"}
            note = _notes_storage[note_id]
            if title is not None:
                if not title.strip():
                    return {"success": False, "error": "Title cannot be empty"}
                note["title"] = title.strip()
            if content is not None:
                if not content.strip():
                    return {"success": False, "error": "Content cannot be empty"}
                note["content"] = content.strip()
            if tags is not None:
                note["tags"] = tags
            note["updated_at"] = datetime.now(UTC).isoformat()
        except (ValueError, TypeError) as e:
            return {"success": False, "error": f"Failed to update note: {e}"}
        else:
            _persist()
            return {
                "success": True,
                "note_id": note_id,
                "message": f"Note '{note['title']}' updated successfully",
                "total_count": len(_notes_storage),
            }


def _delete_note_impl(note_id: str) -> dict[str, Any]:
    with _notes_lock:
        try:
            if note_id not in _notes_storage:
                return {"success": False, "error": f"Note with ID '{note_id}' not found"}
            note = _notes_storage[note_id]
            note_title = note["title"]
            del _notes_storage[note_id]
        except (ValueError, TypeError) as e:
            return {"success": False, "error": f"Failed to delete note: {e}"}
        else:
            _persist()
            return {
                "success": True,
                "note_id": note_id,
                "message": f"Note '{note_title}' deleted successfully",
                "total_count": len(_notes_storage),
            }


@function_tool(timeout=30)
async def create_note(
    ctx: RunContextWrapper,
    title: str,
    content: str,
    category: str = "general",
    tags: list[str] | None = None,
) -> str:
    """Document an observation, finding, methodology step, or research note.

    Notes are visible to every agent in the same scan for the lifetime
    of the run; they live in-memory only and are cleared when the
    process exits. Each note records the agent that wrote it, so
    ``list_notes`` / ``get_note`` show the author (``agent_name``) and
    flag your own notes with ``by_you``.

    For actionable tasks, use ``todo`` instead — notes are for capturing
    information, todos are for tracking work.

    Categories:

    - ``general`` — default, anything that doesn't fit elsewhere.
    - ``findings`` — confirmed vulnerabilities or weaknesses (write
      these up promptly; you'll cite them when filing reports).
    - ``methodology`` — what you tried, what worked, what didn't —
      useful for the final scan report.
    - ``questions`` — open questions / things to come back to.
    - ``plan`` — multi-step plans you want to track.
    - ``wiki`` — long-form repository or target maps.

    Tags are free-form (e.g. ``["sqli", "auth", "critical"]``) — useful
    for later ``list_notes(tags=...)`` filtering.

    Args:
        title: Short headline.
        content: Full note body. Markdown is preserved.
        category: One of the categories above. Default ``"general"``.
        tags: Optional free-form tags.
    """
    agent_id, agent_name = _caller_identity(ctx)
    return json.dumps(
        await asyncio.to_thread(
            _create_note_impl, title, content, category, tags, agent_id, agent_name
        ),
        ensure_ascii=False,
        default=str,
    )


@function_tool(timeout=30)
async def list_notes(
    ctx: RunContextWrapper,
    category: str | None = None,
    tags: list[str] | None = None,
    search: str | None = None,
    include_content: bool = False,
) -> str:
    """List existing notes — metadata-first by default.

    Filters compose: passing ``category="findings"`` and
    ``tags=["sqli"]`` returns notes that are *both* in the findings
    category AND have at least one of those tags.

    By default each entry includes a ``content_preview`` (first 280
    chars). Set ``include_content=True`` to get full bodies — useful
    when you need to scan many notes; expensive in tokens for large
    notes.

    Each entry also carries the author (``agent_name``) and, for notes
    you wrote yourself, ``by_you: true``.

    Args:
        category: Filter by category.
        tags: Filter to notes that have any of these tags.
        search: Substring match against title and content.
        include_content: When False (default) entries have a preview;
            when True the full ``content`` is included.
    """
    caller_agent_id, _ = _caller_identity(ctx)
    return json.dumps(
        await asyncio.to_thread(
            _list_notes_impl,
            category=category,
            tags=tags,
            search=search,
            include_content=include_content,
            caller_agent_id=caller_agent_id,
        ),
        ensure_ascii=False,
        default=str,
    )


@function_tool(timeout=30)
async def get_note(ctx: RunContextWrapper, note_id: str) -> str:
    """Fetch one note by its 6-char ID. Returns the full content.

    Args:
        note_id: Note id from ``create_note`` or a ``list_notes`` entry.
    """
    caller_agent_id, _ = _caller_identity(ctx)
    return json.dumps(
        await asyncio.to_thread(_get_note_impl, note_id, caller_agent_id),
        ensure_ascii=False,
        default=str,
    )


@function_tool(timeout=30)
async def update_note(
    ctx: RunContextWrapper,
    note_id: str,
    title: str | None = None,
    content: str | None = None,
    tags: list[str] | None = None,
) -> str:
    """Update a note's title, content, or tags.

    Pass ``None`` for any field you want left unchanged. Replacing
    ``content`` is a full overwrite — to append, fetch first with
    ``get_note``, concat, and pass the result.

    Args:
        note_id: Target note's 6-char ID.
        title: New title, or ``None`` to keep.
        content: New content, or ``None`` to keep.
        tags: New tags list, or ``None`` to keep.
    """
    return json.dumps(
        await asyncio.to_thread(
            _update_note_impl,
            note_id=note_id,
            title=title,
            content=content,
            tags=tags,
        ),
        ensure_ascii=False,
        default=str,
    )


@function_tool(timeout=30)
async def delete_note(ctx: RunContextWrapper, note_id: str) -> str:
    """Delete a note.

    Args:
        note_id: Note id to delete.
    """
    return json.dumps(
        await asyncio.to_thread(_delete_note_impl, note_id), ensure_ascii=False, default=str
    )