| 1 | from __future__ import annotations |
| 2 | |
| 3 | import json |
| 4 | import re |
| 5 | import time |
| 6 | from datetime import datetime, timezone |
| 7 | from pathlib import Path |
| 8 | from typing import Any |
| 9 | |
| 10 | from helpers import files |
| 11 | from helpers.tool import Response, Tool |
| 12 | |
| 13 | |
| 14 | PLUGIN_NAME = "_goal" |
| 15 | GOALS_DIR = "goals" |
| 16 | ACTIVE_STATUSES = {"active", "paused"} |
| 17 | FINAL_STATUSES = {"complete", "blocked"} |
| 18 | VALID_STATUSES = ACTIVE_STATUSES | FINAL_STATUSES |
| 19 | |
| 20 | |
| 21 | class GoalTool(Tool): |
| 22 | async def execute( |
| 23 | self, |
| 24 | action: str = "", |
| 25 | objective: str = "", |
| 26 | status: str = "", |
| 27 | note: str = "", |
| 28 | token_budget: int | None = None, |
| 29 | **kwargs, |
| 30 | ) -> Response: |
| 31 | action = str(action or self.args.get("action") or "get").strip().lower() |
| 32 | context_id = self.agent.context.id |
| 33 | |
| 34 | try: |
| 35 | if action in {"get", "show", "status"}: |
| 36 | return Response(message=summarize_goal(get_goal(context_id)), break_loop=False) |
| 37 | if action in {"create", "set"}: |
| 38 | goal = create_goal( |
| 39 | context_id, |
| 40 | objective, |
| 41 | created_by="model", |
| 42 | token_budget=token_budget, |
| 43 | ) |
| 44 | return Response(message=f"Goal created: {goal['objective']}", break_loop=False) |
| 45 | if action in {"complete", "blocked"}: |
| 46 | status = action |
| 47 | if action in {"update", "complete", "blocked"}: |
| 48 | status = str(status).strip().lower() |
| 49 | if status not in FINAL_STATUSES: |
| 50 | return Response( |
| 51 | message="Model-managed goal updates may only mark goals complete or blocked.", |
| 52 | break_loop=False, |
| 53 | ) |
| 54 | goal = update_goal( |
| 55 | context_id, |
| 56 | status=status, |
| 57 | objective=objective or None, |
| 58 | note=note or None, |
| 59 | ) |
| 60 | return Response(message=summarize_goal(goal), break_loop=False) |
| 61 | except (FileNotFoundError, ValueError) as error: |
| 62 | return Response(message=str(error), break_loop=False) |
| 63 | |
| 64 | return Response( |
| 65 | message="Unknown goal action. Supported actions: get, create, update.", |
| 66 | break_loop=False, |
| 67 | ) |
| 68 | |
| 69 | |
| 70 | def get_goal(context_id: str) -> dict[str, Any] | None: |
| 71 | context_id = _require_context_id(context_id) |
| 72 | path = _goal_path(context_id) |
| 73 | if not Path(path).is_file(): |
| 74 | return None |
| 75 | |
| 76 | try: |
| 77 | raw = json.loads(files.read_file(path)) |
| 78 | except (OSError, json.JSONDecodeError): |
| 79 | return None |
| 80 | if not isinstance(raw, dict): |
| 81 | return None |
| 82 | |
| 83 | goal = _normalize_goal(raw, context_id=context_id) |
| 84 | return goal if goal.get("objective") else None |
| 85 | |
| 86 | |
| 87 | def create_goal( |
| 88 | context_id: str, |
| 89 | objective: str, |
| 90 | *, |
| 91 | created_by: str = "user", |
| 92 | token_budget: int | None = None, |
| 93 | ) -> dict[str, Any]: |
| 94 | context_id = _require_context_id(context_id) |
| 95 | objective = _clean_objective(objective) |
| 96 | if not objective: |
| 97 | raise ValueError("Goal objective is required") |
| 98 | |
| 99 | now = _now() |
| 100 | existing = get_goal(context_id) |
| 101 | goal = { |
| 102 | "context_id": context_id, |
| 103 | "objective": objective, |
| 104 | "status": "active", |
| 105 | "created_by": _clean_created_by(created_by), |
| 106 | "token_budget": _clean_token_budget(token_budget), |
| 107 | "created_at": existing.get("created_at") if existing else now, |
| 108 | "active_since": now, |
| 109 | "elapsed_seconds": 0, |
| 110 | "updated_at": now, |
| 111 | "note": "", |
| 112 | } |
| 113 | _write_goal(goal) |
| 114 | return goal |
| 115 | |
| 116 | |
| 117 | def update_goal( |
| 118 | context_id: str, |
| 119 | *, |
| 120 | objective: str | None = None, |
| 121 | status: str | None = None, |
| 122 | note: str | None = None, |
| 123 | token_budget: int | None = None, |
| 124 | ) -> dict[str, Any]: |
| 125 | context_id = _require_context_id(context_id) |
| 126 | goal = get_goal(context_id) |
| 127 | if not goal: |
| 128 | raise FileNotFoundError("Goal not found") |
| 129 | |
| 130 | if objective is not None: |
| 131 | goal["objective"] = _required_objective(objective) |
| 132 | if status is not None: |
| 133 | _apply_status(goal, _normalize_status(status)) |
| 134 | if note is not None: |
| 135 | goal["note"] = str(note or "").strip() |
| 136 | if token_budget is not None: |
| 137 | goal["token_budget"] = _clean_token_budget(token_budget) |
| 138 | |
| 139 | goal["updated_at"] = _now() |
| 140 | _write_goal(goal) |
| 141 | return goal |
| 142 | |
| 143 | |
| 144 | def delete_goal(context_id: str) -> None: |
| 145 | context_id = _require_context_id(context_id) |
| 146 | files.delete_file(_goal_path(context_id)) |
| 147 | _notify_goal_changed(context_id) |
| 148 | |
| 149 | |
| 150 | def public_goal(goal: dict[str, Any] | None) -> dict[str, Any] | None: |
| 151 | if not goal: |
| 152 | return None |
| 153 | return { |
| 154 | "context_id": str(goal.get("context_id") or ""), |
| 155 | "objective": str(goal.get("objective") or ""), |
| 156 | "status": str(goal.get("status") or "active"), |
| 157 | "created_by": str(goal.get("created_by") or "user"), |
| 158 | "token_budget": goal.get("token_budget"), |
| 159 | "created_at": str(goal.get("created_at") or ""), |
| 160 | "active_since": str(goal.get("active_since") or ""), |
| 161 | "elapsed_seconds": _clean_elapsed_seconds(goal.get("elapsed_seconds")), |
| 162 | "updated_at": str(goal.get("updated_at") or ""), |
| 163 | "note": str(goal.get("note") or ""), |
| 164 | } |
| 165 | |
| 166 | |
| 167 | def summarize_goal(goal: dict[str, Any] | None) -> str: |
| 168 | if not goal: |
| 169 | return "No goal is set for this chat." |
| 170 | |
| 171 | lines = [ |
| 172 | f"Status: {goal.get('status') or 'active'}", |
| 173 | f"Goal: {str(goal.get('objective') or '').strip()}", |
| 174 | f"Active time: {_format_elapsed(_elapsed_seconds(goal))}", |
| 175 | ] |
| 176 | if updated := str(goal.get("updated_at") or "").strip(): |
| 177 | lines.append(f"Updated: {updated}") |
| 178 | if note := str(goal.get("note") or "").strip(): |
| 179 | lines.append(f"Note: {note}") |
| 180 | return "\n".join(lines) |
| 181 | |
| 182 | |
| 183 | def _write_goal(goal: dict[str, Any]) -> None: |
| 184 | files.write_file( |
| 185 | _goal_path(str(goal["context_id"])), |
| 186 | json.dumps(public_goal(goal), indent=2, ensure_ascii=False) + "\n", |
| 187 | ) |
| 188 | _notify_goal_changed(str(goal["context_id"])) |
| 189 | |
| 190 | |
| 191 | def _notify_goal_changed(context_id: str) -> None: |
| 192 | from agent import AgentContext |
| 193 | |
| 194 | context = AgentContext.get(context_id) |
| 195 | if context is None: |
| 196 | return |
| 197 | context.set_output_data("_goal_revision", time.time()) |
| 198 | |
| 199 | try: |
| 200 | from helpers.state_monitor_integration import mark_dirty_for_context |
| 201 | |
| 202 | mark_dirty_for_context(context_id, reason="plugins._goal") |
| 203 | except Exception: |
| 204 | pass |
| 205 | |
| 206 | |
| 207 | def _normalize_goal(raw: dict[str, Any], *, context_id: str) -> dict[str, Any]: |
| 208 | status = str(raw.get("status") or "active").strip().lower() |
| 209 | status = status if status in VALID_STATUSES else "active" |
| 210 | created_at = str(raw.get("created_at") or "") |
| 211 | updated_at = str(raw.get("updated_at") or "") |
| 212 | active_since = str(raw.get("active_since") or "") |
| 213 | elapsed_seconds = _clean_elapsed_seconds(raw.get("elapsed_seconds")) |
| 214 | if "elapsed_seconds" not in raw and status != "active": |
| 215 | elapsed_seconds = _seconds_between(created_at, updated_at) |
| 216 | if status == "active" and not active_since: |
| 217 | active_since = created_at or updated_at |
| 218 | |
| 219 | return { |
| 220 | "context_id": context_id, |
| 221 | "objective": _clean_objective(str(raw.get("objective") or "")), |
| 222 | "status": status, |
| 223 | "created_by": _clean_created_by(str(raw.get("created_by") or "user")), |
| 224 | "token_budget": _clean_token_budget(raw.get("token_budget")), |
| 225 | "created_at": created_at, |
| 226 | "active_since": active_since, |
| 227 | "elapsed_seconds": elapsed_seconds, |
| 228 | "updated_at": updated_at, |
| 229 | "note": str(raw.get("note") or "").strip(), |
| 230 | } |
| 231 | |
| 232 | |
| 233 | def _goal_path(context_id: str) -> str: |
| 234 | directory = files.get_abs_path(files.USER_DIR, files.PLUGINS_DIR, PLUGIN_NAME, GOALS_DIR) |
| 235 | safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", context_id).strip("._")[:180] |
| 236 | if not safe: |
| 237 | raise ValueError("A chat context is required") |
| 238 | return files.get_abs_path(directory, f"{safe}.json") |
| 239 | |
| 240 | |
| 241 | def _require_context_id(context_id: str) -> str: |
| 242 | context_id = str(context_id or "").strip() |
| 243 | if not context_id: |
| 244 | raise ValueError("A chat context is required") |
| 245 | return context_id |
| 246 | |
| 247 | |
| 248 | def _required_objective(value: str) -> str: |
| 249 | objective = _clean_objective(value) |
| 250 | if not objective: |
| 251 | raise ValueError("Goal objective is required") |
| 252 | return objective |
| 253 | |
| 254 | |
| 255 | def _clean_objective(value: str) -> str: |
| 256 | return re.sub(r"\s+", " ", str(value or "")).strip() |
| 257 | |
| 258 | |
| 259 | def _normalize_status(status: str) -> str: |
| 260 | status = str(status or "").strip().lower() |
| 261 | if status not in VALID_STATUSES: |
| 262 | raise ValueError("Goal status must be active, paused, complete, or blocked") |
| 263 | return status |
| 264 | |
| 265 | |
| 266 | def _clean_created_by(created_by: str) -> str: |
| 267 | created_by = str(created_by or "").strip().lower() |
| 268 | return created_by if created_by in {"user", "model"} else "user" |
| 269 | |
| 270 | |
| 271 | def _clean_token_budget(token_budget: Any) -> int | None: |
| 272 | try: |
| 273 | token_budget = int(token_budget) |
| 274 | except (TypeError, ValueError): |
| 275 | return None |
| 276 | return token_budget if token_budget > 0 else None |
| 277 | |
| 278 | |
| 279 | def _apply_status(goal: dict[str, Any], status: str) -> None: |
| 280 | current = str(goal.get("status") or "active") |
| 281 | if current == "active" and status != "active": |
| 282 | goal["elapsed_seconds"] = _elapsed_seconds(goal) |
| 283 | goal["active_since"] = "" |
| 284 | elif current != "active" and status == "active": |
| 285 | goal["active_since"] = _now() |
| 286 | goal["status"] = status |
| 287 | |
| 288 | |
| 289 | def _elapsed_seconds(goal: dict[str, Any]) -> int: |
| 290 | elapsed = _clean_elapsed_seconds(goal.get("elapsed_seconds")) |
| 291 | if str(goal.get("status") or "active") == "active": |
| 292 | elapsed += _seconds_between(str(goal.get("active_since") or ""), _now()) |
| 293 | return elapsed |
| 294 | |
| 295 | |
| 296 | def _clean_elapsed_seconds(value: Any) -> int: |
| 297 | try: |
| 298 | return max(0, int(value)) |
| 299 | except (TypeError, ValueError): |
| 300 | return 0 |
| 301 | |
| 302 | |
| 303 | def _seconds_between(start: str, end: str) -> int: |
| 304 | start_dt, end_dt = _parse_time(start), _parse_time(end) |
| 305 | return max(0, int((end_dt - start_dt).total_seconds())) if start_dt and end_dt else 0 |
| 306 | |
| 307 | |
| 308 | def _parse_time(value: str) -> datetime | None: |
| 309 | value = str(value or "").strip() |
| 310 | if not value: |
| 311 | return None |
| 312 | if value.endswith("Z"): |
| 313 | value = f"{value[:-1]}+00:00" |
| 314 | try: |
| 315 | parsed = datetime.fromisoformat(value) |
| 316 | except ValueError: |
| 317 | return None |
| 318 | return parsed.replace(tzinfo=timezone.utc) if parsed.tzinfo is None else parsed.astimezone(timezone.utc) |
| 319 | |
| 320 | |
| 321 | def _format_elapsed(seconds: int) -> str: |
| 322 | hours, remainder = divmod(max(0, int(seconds)), 3600) |
| 323 | minutes, seconds = divmod(remainder, 60) |
| 324 | return f"{hours}h {minutes}m" if hours else f"{minutes}m {seconds}s" if minutes else f"{seconds}s" |
| 325 | |
| 326 | |
| 327 | def _now() -> str: |
| 328 | return datetime.now(timezone.utc).isoformat(timespec="seconds").replace("+00:00", "Z") |