"""Per-agent todo tools — mirrored to {state_dir}/todos.json.""" 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", ""), ) _todos_storage: dict[str, dict[str, dict[str, Any]]] = {} _todos_path: Path | None = None _todos_io_lock = threading.RLock() def hydrate_todos_from_disk(state_dir: Path) -> None: 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: 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 @function_tool(timeout=30) async def create_todo(ctx: RunContextWrapper, todos: str) -> str: """Create one or many todos for the current agent. Always pass a list, even for a single todo (wrap it in a one-item array). 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. Args: todos: JSON array of todo objects. For one todo, pass a one-item list. Each object's fields: - ``title`` (str, **required**): short actionable title, e.g. ``"Test /api/admin for IDOR"``. - ``description`` (str, optional): extra context or acceptance criteria. - ``priority`` (str, optional): one of ``"low"`` / ``"normal"`` / ``"high"`` / ``"critical"``. Defaults to ``"normal"``. Example: ``[{"title": "Probe /admin", "priority": "high"}, {"title": "Check JWT alg=none"}]``. """ agent_id = _agent_id_from(ctx) try: tasks = _normalize_bulk_todos(todos) if not tasks: return json.dumps( {"success": False, "error": "Provide a non-empty 'todos' list to create"}, 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")) 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}"}, ensure_ascii=False, default=str, ) _persist() return json.dumps( { "success": True, "created": 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": [], "filtered_count": 0, "total_count": 0, "summary": {"pending": 0, "in_progress": 0, "done": 0}, }, ensure_ascii=False, default=str, ) return json.dumps( { "success": True, "todos": todos_list, "filtered_count": len(todos_list), "total_count": len(agent_todos), "summary": summary, }, ensure_ascii=False, default=str, ) @function_tool(timeout=30) async def update_todo(ctx: RunContextWrapper, updates: str) -> str: """Update one or many todos. Always pass a list, even for a single update (wrap it in a one-item array). For toggling status only, prefer the dedicated ``mark_todo_done`` / ``mark_todo_pending`` tools — they're simpler and accept the same list-of-ids form. Args: updates: JSON array of update objects. For one update, pass a one-item list. Each object's fields: - ``todo_id`` (str, **required**): ID returned by ``create_todo``. - ``title`` (str, optional): new title. - ``description`` (str, optional): new description (empty string clears it). - ``priority`` (str, optional): one of ``"low"`` / ``"normal"`` / ``"high"`` / ``"critical"``. - ``status`` (str, optional): one of ``"pending"`` / ``"in_progress"`` / ``"done"``. Omitted fields stay unchanged. Example: ``[{"todo_id": "abc", "status": "in_progress", "priority": "high"}]``. """ agent_id = _agent_id_from(ctx) try: agent_todos = _get_agent_todos(agent_id) updates_to_apply = _normalize_bulk_updates(updates) if not updates_to_apply: return json.dumps( {"success": False, "error": "Provide a non-empty 'updates' list"}, 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_ids: str, new_status: str) -> str: try: agent_todos = _get_agent_todos(agent_id) ids = _normalize_todo_ids(todo_ids) if not ids: msg = f"Provide a non-empty 'todo_ids' list 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() response: dict[str, Any] = { "success": len(errors) == 0, "marked": marked, "marked_count": len(marked), "new_status": new_status, "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_ids: str) -> str: """Mark one or many todos as done. Always pass a list, even for a single ID (wrap it in a one-item array). Args: todo_ids: JSON array of todo IDs to mark done. For one todo, pass a one-item list. """ return _mark(agent_id=_agent_id_from(ctx), todo_ids=todo_ids, new_status="done") @function_tool(timeout=30) async def mark_todo_pending(ctx: RunContextWrapper, todo_ids: str) -> str: """Reset one or many todos to pending (e.g., to retry a failed task). Always pass a list, even for a single ID (wrap it in a one-item array). Args: todo_ids: JSON array of todo IDs to reset to pending. For one todo, pass a one-item list. """ return _mark(agent_id=_agent_id_from(ctx), todo_ids=todo_ids, new_status="pending") @function_tool(timeout=30) async def delete_todo(ctx: RunContextWrapper, todo_ids: str) -> str: """Delete one or many todos. Removes them entirely (no soft-delete). Always pass a list, even for a single ID (wrap it in a one-item array). Args: todo_ids: JSON array of todo IDs to delete. For one todo, pass a one-item list. """ agent_id = _agent_id_from(ctx) try: agent_todos = _get_agent_todos(agent_id) ids = _normalize_todo_ids(todo_ids) if not ids: return json.dumps( {"success": False, "error": "Provide a non-empty 'todo_ids' list 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)