Files
strix/strix/tools/todo/tools.py
T
2026-04-26 15:01:35 -07:00

667 lines
22 KiB
Python

"""Per-agent todo tools.
Per-agent in-memory dict, scoped via ``ctx.context['agent_id']``. The
table is mirrored to ``{state_dir}/todos.json`` after every mutation so a
process restart can ``hydrate_todos_from_disk`` and each respawned
agent finds its prior list intact. The persistence is best-effort —
errors are logged and swallowed so a disk failure can't kill the agent
mid-call. Bulk forms are preserved so the prompt-template documentation
still works (``todos`` / ``updates`` / ``todo_ids`` accept JSON strings
or comma-separated strings).
"""
from __future__ import annotations
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
logger = logging.getLogger(__name__)
VALID_PRIORITIES = ["low", "normal", "high", "critical"]
VALID_STATUSES = ["pending", "in_progress", "done"]
_PRIORITY_RANK = {"critical": 0, "high": 1, "normal": 2, "low": 3}
_STATUS_RANK = {"done": 0, "in_progress": 1, "pending": 2}
def _todo_sort_key(todo: dict[str, Any]) -> tuple[int, int, str]:
return (
_STATUS_RANK.get(todo.get("status", "pending"), 99),
_PRIORITY_RANK.get(todo.get("priority", "normal"), 99),
todo.get("created_at", ""),
)
# Per-agent silo: ``_todos_storage[agent_id][todo_id] = todo_dict``.
# Keyed by ``ctx.context['agent_id']`` so two agents in the same scan
# don't see each other's lists.
_todos_storage: dict[str, dict[str, dict[str, Any]]] = {}
# On-disk mirror path. Set by ``hydrate_todos_from_disk`` once per scan;
# unset means "no persistence" (e.g. unit tests). All writes go through
# ``_persist`` which is a no-op until the path is set.
_todos_path: Path | None = None
_todos_io_lock = threading.RLock()
def hydrate_todos_from_disk(state_dir: Path) -> None:
"""Wire the on-disk mirror at ``{state_dir}/todos.json`` and reload it.
Called by :func:`run_strix_scan` once at scan setup. Subsequent CRUD
calls auto-persist after every mutation. Idempotent on missing file.
Tolerant of corruption — logs and starts empty rather than failing
the scan over a broken sidecar artifact.
"""
global _todos_path # noqa: PLW0603
_todos_path = state_dir / "todos.json"
with _todos_io_lock:
_todos_storage.clear()
if not _todos_path.exists():
return
try:
data = json.loads(_todos_path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError):
logger.exception(
"todos.json at %s is unreadable; starting with empty todos",
_todos_path,
)
return
if not isinstance(data, dict):
return
loaded = 0
for aid, by_id in data.items():
if not isinstance(aid, str) or not isinstance(by_id, dict):
continue
cleaned = {
str(tid): t
for tid, t in by_id.items()
if isinstance(tid, str) and isinstance(t, dict)
}
if cleaned:
_todos_storage[aid] = cleaned
loaded += len(cleaned)
logger.info(
"todos hydrated from %s (%d agent(s), %d todo(s))",
_todos_path,
len(_todos_storage),
loaded,
)
def _persist() -> None:
"""Atomic-rename mirror of ``_todos_storage`` → ``{state_dir}/todos.json``.
No-op when ``_todos_path`` isn't wired (tests). Errors are logged
and swallowed.
"""
path = _todos_path
if path is None:
return
try:
payload = json.dumps(_todos_storage, ensure_ascii=False, default=str)
path.parent.mkdir(parents=True, exist_ok=True)
with (
_todos_io_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("todos persist to %s failed", path)
def _agent_id_from(ctx: RunContextWrapper) -> str:
inner = ctx.context if isinstance(ctx.context, dict) else {}
return str(inner.get("agent_id") or "default")
def _get_agent_todos(agent_id: str) -> dict[str, dict[str, Any]]:
return _todos_storage.setdefault(agent_id, {})
def _normalize_priority(priority: str | None, default: str = "normal") -> str:
candidate = (priority or default or "normal").lower()
if candidate not in VALID_PRIORITIES:
raise ValueError(f"Invalid priority. Must be one of: {', '.join(VALID_PRIORITIES)}")
return candidate
def _sorted_todos(agent_id: str) -> list[dict[str, Any]]:
todos_list = [
{**todo, "todo_id": todo_id} for todo_id, todo in _get_agent_todos(agent_id).items()
]
todos_list.sort(key=_todo_sort_key)
return todos_list
def _normalize_todo_ids(raw_ids: Any) -> list[str]:
if raw_ids is None:
return []
if isinstance(raw_ids, str):
stripped = raw_ids.strip()
if not stripped:
return []
try:
data = json.loads(stripped)
except json.JSONDecodeError:
data = stripped.split(",") if "," in stripped else [stripped]
if isinstance(data, list):
return [str(item).strip() for item in data if str(item).strip()]
return [str(data).strip()]
if isinstance(raw_ids, list):
return [str(item).strip() for item in raw_ids if str(item).strip()]
return [str(raw_ids).strip()]
def _normalize_bulk_updates(raw_updates: Any) -> list[dict[str, Any]]:
if raw_updates is None:
return []
data: Any = raw_updates
if isinstance(raw_updates, str):
stripped = raw_updates.strip()
if not stripped:
return []
try:
data = json.loads(stripped)
except json.JSONDecodeError as e:
raise ValueError("Updates must be valid JSON") from e
if isinstance(data, dict):
data = [data]
if not isinstance(data, list):
raise TypeError("Updates must be a list of update objects")
normalized: list[dict[str, Any]] = []
for item in data:
if not isinstance(item, dict):
raise TypeError("Each update must be an object with todo_id")
todo_id = item.get("todo_id") or item.get("id")
if not todo_id:
raise ValueError("Each update must include 'todo_id'")
normalized.append(
{
"todo_id": str(todo_id).strip(),
"title": item.get("title"),
"description": item.get("description"),
"priority": item.get("priority"),
"status": item.get("status"),
},
)
return normalized
def _normalize_bulk_todos(raw_todos: Any) -> list[dict[str, Any]]:
if raw_todos is None:
return []
data: Any = raw_todos
if isinstance(raw_todos, str):
stripped = raw_todos.strip()
if not stripped:
return []
try:
data = json.loads(stripped)
except json.JSONDecodeError:
entries = [line.strip(" -*\t") for line in stripped.splitlines() if line.strip(" -*\t")]
return [{"title": entry} for entry in entries]
if isinstance(data, dict):
data = [data]
if not isinstance(data, list):
raise TypeError("Todos must be provided as a list, dict, or JSON string")
normalized: list[dict[str, Any]] = []
for item in data:
if isinstance(item, str):
title = item.strip()
if title:
normalized.append({"title": title})
continue
if not isinstance(item, dict):
raise TypeError("Each todo entry must be a string or object with a title")
title = item.get("title", "")
if not isinstance(title, str) or not title.strip():
raise ValueError("Each todo entry must include a non-empty 'title'")
normalized.append(
{
"title": title.strip(),
"description": (item.get("description") or "").strip() or None,
"priority": item.get("priority"),
},
)
return normalized
def _apply_single_update(
agent_todos: dict[str, dict[str, Any]],
todo_id: str,
title: str | None = None,
description: str | None = None,
priority: str | None = None,
status: str | None = None,
) -> dict[str, Any] | None:
if todo_id not in agent_todos:
return {"todo_id": todo_id, "error": f"Todo with ID '{todo_id}' not found"}
todo = agent_todos[todo_id]
if title is not None:
if not title.strip():
return {"todo_id": todo_id, "error": "Title cannot be empty"}
todo["title"] = title.strip()
if description is not None:
todo["description"] = description.strip() if description else None
if priority is not None:
try:
todo["priority"] = _normalize_priority(priority, str(todo.get("priority", "normal")))
except ValueError as exc:
return {"todo_id": todo_id, "error": str(exc)}
if status is not None:
status_candidate = status.lower()
if status_candidate not in VALID_STATUSES:
return {
"todo_id": todo_id,
"error": f"Invalid status. Must be one of: {', '.join(VALID_STATUSES)}",
}
todo["status"] = status_candidate
todo["completed_at"] = datetime.now(UTC).isoformat() if status_candidate == "done" else None
todo["updated_at"] = datetime.now(UTC).isoformat()
return None
# --- public tools ---------------------------------------------------------
@function_tool(timeout=30)
async def create_todo(
ctx: RunContextWrapper,
title: str | None = None,
description: str | None = None,
priority: str = "normal",
todos: str | None = None,
) -> str:
"""Create one or many todos for the current agent.
Each agent (including subagents) has its **own private todo list** —
your todos don't leak to other agents and vice versa.
When to use:
- Planning multi-step assessments with parallel workstreams.
- Tracking work you'll come back to later.
- Breaking down complex scopes (per-endpoint, per-target, per-vuln-class).
When NOT to use:
- Simple linear workflows where progress is obvious.
- Single quick task — just do it.
Batch related todos in one call via the ``todos`` bulk parameter
rather than firing many ``create_todo`` calls.
Args:
title: Short, actionable title (e.g., "Test /api/admin for IDOR").
description: Optional details / context for the single todo.
priority: ``"low"`` / ``"normal"`` / ``"high"`` / ``"critical"``.
todos: Bulk create — either JSON array of
``{"title": "...", "description": "...", "priority": "..."}``
objects, or a newline-separated bullet list (``- item\\n- item``).
"""
agent_id = _agent_id_from(ctx)
try:
default_priority = _normalize_priority(priority)
tasks: list[dict[str, Any]] = []
if todos is not None:
tasks.extend(_normalize_bulk_todos(todos))
if title and title.strip():
tasks.append(
{
"title": title.strip(),
"description": description.strip() if description else None,
"priority": default_priority,
},
)
if not tasks:
return json.dumps(
{
"success": False,
"error": "Provide a title or 'todos' list to create.",
"todo_id": None,
},
ensure_ascii=False,
default=str,
)
agent_todos = _get_agent_todos(agent_id)
created: list[dict[str, Any]] = []
for task in tasks:
task_priority = _normalize_priority(task.get("priority"), default_priority)
todo_id = str(uuid.uuid4())[:6]
timestamp = datetime.now(UTC).isoformat()
agent_todos[todo_id] = {
"title": task["title"],
"description": task.get("description"),
"priority": task_priority,
"status": "pending",
"created_at": timestamp,
"updated_at": timestamp,
"completed_at": None,
}
created.append({"todo_id": todo_id, "title": task["title"], "priority": task_priority})
except (ValueError, TypeError) as e:
return json.dumps(
{"success": False, "error": f"Failed to create todo: {e}", "todo_id": None},
ensure_ascii=False,
default=str,
)
_persist()
return json.dumps(
{
"success": True,
"created": created,
"count": len(created),
"todos": _sorted_todos(agent_id),
"total_count": len(_get_agent_todos(agent_id)),
},
ensure_ascii=False,
default=str,
)
@function_tool(timeout=30)
async def list_todos(
ctx: RunContextWrapper,
status: str | None = None,
priority: str | None = None,
) -> str:
"""List the current agent's todos, sorted by status then priority.
Sort order: status (done → in_progress → pending), then priority
within each status (critical → high → normal → low).
Args:
status: Filter — ``"pending"`` / ``"in_progress"`` / ``"done"``.
priority: Filter — ``"low"`` / ``"normal"`` / ``"high"`` /
``"critical"``.
"""
agent_id = _agent_id_from(ctx)
try:
agent_todos = _get_agent_todos(agent_id)
status_filter = status.lower() if isinstance(status, str) else None
priority_filter = priority.lower() if isinstance(priority, str) else None
todos_list: list[dict[str, Any]] = []
for todo_id, todo in agent_todos.items():
if status_filter and todo.get("status") != status_filter:
continue
if priority_filter and todo.get("priority") != priority_filter:
continue
entry = todo.copy()
entry["todo_id"] = todo_id
todos_list.append(entry)
todos_list.sort(key=_todo_sort_key)
summary: dict[str, int] = {"pending": 0, "in_progress": 0, "done": 0}
for todo in todos_list:
sv = todo.get("status", "pending")
summary[sv] = summary.get(sv, 0) + 1
except (ValueError, TypeError) as e:
return json.dumps(
{
"success": False,
"error": f"Failed to list todos: {e}",
"todos": [],
"total_count": 0,
"summary": {"pending": 0, "in_progress": 0, "done": 0},
},
ensure_ascii=False,
default=str,
)
return json.dumps(
{
"success": True,
"todos": todos_list,
"total_count": len(todos_list),
"summary": summary,
},
ensure_ascii=False,
default=str,
)
@function_tool(timeout=30)
async def update_todo(
ctx: RunContextWrapper,
todo_id: str | None = None,
title: str | None = None,
description: str | None = None,
priority: str | None = None,
status: str | None = None,
updates: str | None = None,
) -> str:
"""Update one or many todos. Prefer the bulk form for multiple updates.
For toggling status only, use the dedicated ``mark_todo_done`` /
``mark_todo_pending`` tools — they're simpler and accept bulk
``todo_ids``.
Args:
todo_id: Single-todo target.
title / description / priority / status: New values for the
single todo. Omit to leave unchanged.
updates: Bulk form — JSON array like
``[{"todo_id": "abc", "status": "done"}, ...]``.
"""
agent_id = _agent_id_from(ctx)
try:
agent_todos = _get_agent_todos(agent_id)
updates_to_apply: list[dict[str, Any]] = []
if updates is not None:
updates_to_apply.extend(_normalize_bulk_updates(updates))
if todo_id is not None:
updates_to_apply.append(
{
"todo_id": todo_id,
"title": title,
"description": description,
"priority": priority,
"status": status,
},
)
if not updates_to_apply:
return json.dumps(
{"success": False, "error": "Provide todo_id or 'updates' list to update."},
ensure_ascii=False,
default=str,
)
updated: list[str] = []
errors: list[dict[str, Any]] = []
for upd in updates_to_apply:
err = _apply_single_update(
agent_todos,
upd["todo_id"],
upd.get("title"),
upd.get("description"),
upd.get("priority"),
upd.get("status"),
)
if err:
errors.append(err)
else:
updated.append(upd["todo_id"])
except (ValueError, TypeError) as e:
return json.dumps({"success": False, "error": str(e)}, ensure_ascii=False, default=str)
if updated:
_persist()
response: dict[str, Any] = {
"success": len(errors) == 0,
"updated": updated,
"updated_count": len(updated),
"todos": _sorted_todos(agent_id),
"total_count": len(agent_todos),
}
if errors:
response["errors"] = errors
return json.dumps(response, ensure_ascii=False, default=str)
def _mark(
*,
agent_id: str,
todo_id: str | None,
todo_ids: str | None,
new_status: str,
) -> str:
try:
agent_todos = _get_agent_todos(agent_id)
ids: list[str] = []
if todo_ids is not None:
ids.extend(_normalize_todo_ids(todo_ids))
if todo_id is not None:
ids.append(todo_id)
if not ids:
msg = f"Provide todo_id or todo_ids to mark as {new_status}."
return json.dumps({"success": False, "error": msg}, ensure_ascii=False, default=str)
marked: list[str] = []
errors: list[dict[str, Any]] = []
timestamp = datetime.now(UTC).isoformat()
for tid in ids:
if tid not in agent_todos:
errors.append({"todo_id": tid, "error": f"Todo with ID '{tid}' not found"})
continue
todo = agent_todos[tid]
todo["status"] = new_status
todo["completed_at"] = timestamp if new_status == "done" else None
todo["updated_at"] = timestamp
marked.append(tid)
except (ValueError, TypeError) as e:
return json.dumps({"success": False, "error": str(e)}, ensure_ascii=False, default=str)
if marked:
_persist()
key = "marked_done" if new_status == "done" else "marked_pending"
response: dict[str, Any] = {
"success": len(errors) == 0,
key: marked,
"marked_count": len(marked),
"todos": _sorted_todos(agent_id),
"total_count": len(agent_todos),
}
if errors:
response["errors"] = errors
return json.dumps(response, ensure_ascii=False, default=str)
@function_tool(timeout=30)
async def mark_todo_done(
ctx: RunContextWrapper,
todo_id: str | None = None,
todo_ids: str | None = None,
) -> str:
"""Mark one or many todos as done.
Args:
todo_id: Single todo's ID.
todo_ids: Bulk form — JSON array, comma-separated string, or
single ID. Combinable with ``todo_id`` for one-off plus bulk.
"""
return _mark(
agent_id=_agent_id_from(ctx),
todo_id=todo_id,
todo_ids=todo_ids,
new_status="done",
)
@function_tool(timeout=30)
async def mark_todo_pending(
ctx: RunContextWrapper,
todo_id: str | None = None,
todo_ids: str | None = None,
) -> str:
"""Reset one or many todos to pending (e.g., to retry a failed task).
Args:
todo_id: Single todo's ID.
todo_ids: Bulk form — JSON array, comma-separated, or single ID.
"""
return _mark(
agent_id=_agent_id_from(ctx),
todo_id=todo_id,
todo_ids=todo_ids,
new_status="pending",
)
@function_tool(timeout=30)
async def delete_todo(
ctx: RunContextWrapper,
todo_id: str | None = None,
todo_ids: str | None = None,
) -> str:
"""Delete one or many todos. Removes them entirely (no soft-delete).
Args:
todo_id: Single todo's ID.
todo_ids: Bulk form — JSON array, comma-separated, or single ID.
"""
agent_id = _agent_id_from(ctx)
try:
agent_todos = _get_agent_todos(agent_id)
ids: list[str] = []
if todo_ids is not None:
ids.extend(_normalize_todo_ids(todo_ids))
if todo_id is not None:
ids.append(todo_id)
if not ids:
return json.dumps(
{"success": False, "error": "Provide todo_id or todo_ids to delete."},
ensure_ascii=False,
default=str,
)
deleted: list[str] = []
errors: list[dict[str, Any]] = []
for tid in ids:
if tid not in agent_todos:
errors.append({"todo_id": tid, "error": f"Todo with ID '{tid}' not found"})
continue
del agent_todos[tid]
deleted.append(tid)
except (ValueError, TypeError) as e:
return json.dumps({"success": False, "error": str(e)}, ensure_ascii=False, default=str)
if deleted:
_persist()
response: dict[str, Any] = {
"success": len(errors) == 0,
"deleted": deleted,
"deleted_count": len(deleted),
"todos": _sorted_todos(agent_id),
"total_count": len(agent_todos),
}
if errors:
response["errors"] = errors
return json.dumps(response, ensure_ascii=False, default=str)