|
| 1 | +/** |
| 2 | + * Heartbeat Extension |
| 3 | + * |
| 4 | + * Periodically injects a heartbeat prompt into the agent's conversation so it |
| 5 | + * can perform health checks, clean up stale resources, and act proactively |
| 6 | + * without waiting for external events. |
| 7 | + * |
| 8 | + * The heartbeat reads a configurable checklist file (HEARTBEAT.md) and sends |
| 9 | + * it as a follow-up message. If the file is empty or missing, no heartbeat |
| 10 | + * fires (saves tokens). |
| 11 | + * |
| 12 | + * Configuration (env vars): |
| 13 | + * HEARTBEAT_INTERVAL_MS — interval between heartbeats (default: 600000 = 10 min) |
| 14 | + * HEARTBEAT_FILE — path to checklist file (default: ~/.pi/agent/HEARTBEAT.md) |
| 15 | + * HEARTBEAT_ENABLED — set to "0" or "false" to disable (default: enabled) |
| 16 | + * |
| 17 | + * Inspired by OpenClaw's HEARTBEAT.md pattern — a user-configurable Markdown |
| 18 | + * checklist that the agent evaluates on each tick. |
| 19 | + */ |
| 20 | + |
| 21 | +import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent"; |
| 22 | +import { Type } from "@sinclair/typebox"; |
| 23 | +import { StringEnum } from "@mariozechner/pi-ai"; |
| 24 | +import { existsSync, readFileSync } from "node:fs"; |
| 25 | +import { homedir } from "node:os"; |
| 26 | +import { join } from "node:path"; |
| 27 | + |
| 28 | +const DEFAULT_INTERVAL_MS = 10 * 60 * 1000; // 10 minutes |
| 29 | +const DEFAULT_HEARTBEAT_FILE = join(homedir(), ".pi", "agent", "HEARTBEAT.md"); |
| 30 | + |
| 31 | +// Minimum interval to prevent accidental token burn (2 minutes) |
| 32 | +const MIN_INTERVAL_MS = 2 * 60 * 1000; |
| 33 | + |
| 34 | +// Maximum consecutive errors before backing off |
| 35 | +const MAX_CONSECUTIVE_ERRORS = 5; |
| 36 | +const BACKOFF_MULTIPLIER = 2; |
| 37 | +const MAX_BACKOFF_MS = 60 * 60 * 1000; // 1 hour |
| 38 | + |
| 39 | +type HeartbeatState = { |
| 40 | + enabled: boolean; |
| 41 | + intervalMs: number; |
| 42 | + heartbeatFile: string; |
| 43 | + lastRunAt: number | null; |
| 44 | + consecutiveErrors: number; |
| 45 | + totalRuns: number; |
| 46 | +}; |
| 47 | + |
| 48 | +const HEARTBEAT_STATE_ENTRY = "heartbeat-state"; |
| 49 | + |
| 50 | +function isDisabledByEnv(): boolean { |
| 51 | + const val = process.env.HEARTBEAT_ENABLED?.trim().toLowerCase(); |
| 52 | + return val === "0" || val === "false" || val === "no"; |
| 53 | +} |
| 54 | + |
| 55 | +function resolveConfig(): { intervalMs: number; heartbeatFile: string; enabled: boolean } { |
| 56 | + const envInterval = parseInt(process.env.HEARTBEAT_INTERVAL_MS || "", 10); |
| 57 | + const intervalMs = Math.max( |
| 58 | + MIN_INTERVAL_MS, |
| 59 | + Number.isFinite(envInterval) ? envInterval : DEFAULT_INTERVAL_MS |
| 60 | + ); |
| 61 | + const heartbeatFile = process.env.HEARTBEAT_FILE?.trim() || DEFAULT_HEARTBEAT_FILE; |
| 62 | + const enabled = !isDisabledByEnv(); |
| 63 | + return { intervalMs, heartbeatFile, enabled }; |
| 64 | +} |
| 65 | + |
| 66 | +function readHeartbeatFile(filepath: string): string | null { |
| 67 | + try { |
| 68 | + if (!existsSync(filepath)) return null; |
| 69 | + const content = readFileSync(filepath, "utf-8").trim(); |
| 70 | + // Skip if empty or only comments/whitespace |
| 71 | + const meaningful = content |
| 72 | + .split("\n") |
| 73 | + .filter((line) => { |
| 74 | + const trimmed = line.trim(); |
| 75 | + return trimmed.length > 0 && !trimmed.startsWith("#"); |
| 76 | + }) |
| 77 | + .join("\n") |
| 78 | + .trim(); |
| 79 | + return meaningful.length > 0 ? content : null; |
| 80 | + } catch { |
| 81 | + return null; |
| 82 | + } |
| 83 | +} |
| 84 | + |
| 85 | +function computeBackoffMs(consecutiveErrors: number, baseInterval: number): number { |
| 86 | + if (consecutiveErrors <= 0) return baseInterval; |
| 87 | + const backoff = baseInterval * Math.pow(BACKOFF_MULTIPLIER, consecutiveErrors); |
| 88 | + return Math.min(backoff, MAX_BACKOFF_MS); |
| 89 | +} |
| 90 | + |
| 91 | +export default function heartbeatExtension(pi: ExtensionAPI): void { |
| 92 | + let timer: ReturnType<typeof setTimeout> | null = null; |
| 93 | + let state: HeartbeatState = { |
| 94 | + enabled: true, |
| 95 | + intervalMs: DEFAULT_INTERVAL_MS, |
| 96 | + heartbeatFile: DEFAULT_HEARTBEAT_FILE, |
| 97 | + lastRunAt: null, |
| 98 | + consecutiveErrors: 0, |
| 99 | + totalRuns: 0, |
| 100 | + }; |
| 101 | + |
| 102 | + function saveState() { |
| 103 | + pi.appendEntry(HEARTBEAT_STATE_ENTRY, { |
| 104 | + lastRunAt: state.lastRunAt, |
| 105 | + consecutiveErrors: state.consecutiveErrors, |
| 106 | + totalRuns: state.totalRuns, |
| 107 | + }); |
| 108 | + } |
| 109 | + |
| 110 | + function armTimer() { |
| 111 | + if (timer) clearTimeout(timer); |
| 112 | + timer = null; |
| 113 | + |
| 114 | + if (!state.enabled) return; |
| 115 | + |
| 116 | + const delay = computeBackoffMs(state.consecutiveErrors, state.intervalMs); |
| 117 | + timer = setTimeout(() => { |
| 118 | + fireHeartbeat(); |
| 119 | + }, delay); |
| 120 | + } |
| 121 | + |
| 122 | + function fireHeartbeat() { |
| 123 | + const content = readHeartbeatFile(state.heartbeatFile); |
| 124 | + if (!content) { |
| 125 | + // No checklist — skip silently, re-arm for next interval |
| 126 | + armTimer(); |
| 127 | + return; |
| 128 | + } |
| 129 | + |
| 130 | + const now = Date.now(); |
| 131 | + state.lastRunAt = now; |
| 132 | + state.totalRuns += 1; |
| 133 | + |
| 134 | + const prompt = [ |
| 135 | + `🫀 **Heartbeat** (run #${state.totalRuns}, ${new Date(now).toISOString()})`, |
| 136 | + ``, |
| 137 | + `Review the following checklist and take action on any items that need attention.`, |
| 138 | + `If everything is healthy, respond briefly with what you checked. Do NOT take action unless something is wrong.`, |
| 139 | + ``, |
| 140 | + `---`, |
| 141 | + content, |
| 142 | + `---`, |
| 143 | + ``, |
| 144 | + `If you find issues, fix them. If everything looks good, say so briefly and move on.`, |
| 145 | + ].join("\n"); |
| 146 | + |
| 147 | + pi.sendMessage( |
| 148 | + { |
| 149 | + customType: "heartbeat", |
| 150 | + content: prompt, |
| 151 | + display: true, |
| 152 | + }, |
| 153 | + { |
| 154 | + deliverAs: "followUp", |
| 155 | + triggerTurn: true, |
| 156 | + } |
| 157 | + ); |
| 158 | + |
| 159 | + saveState(); |
| 160 | + // Re-arm after firing (the agent_end handler will also re-arm on error) |
| 161 | + armTimer(); |
| 162 | + } |
| 163 | + |
| 164 | + function stopTimer() { |
| 165 | + if (timer) { |
| 166 | + clearTimeout(timer); |
| 167 | + timer = null; |
| 168 | + } |
| 169 | + } |
| 170 | + |
| 171 | + // ── Tool: heartbeat control ─────────────────────────────────────────────── |
| 172 | + |
| 173 | + pi.registerTool({ |
| 174 | + name: "heartbeat", |
| 175 | + label: "Heartbeat", |
| 176 | + description: |
| 177 | + "Manage the periodic heartbeat loop. " + |
| 178 | + "Actions: status (check state), pause (stop heartbeats), resume (restart), " + |
| 179 | + "trigger (fire one now), config (show configuration).", |
| 180 | + parameters: Type.Object({ |
| 181 | + action: StringEnum(["status", "pause", "resume", "trigger", "config"] as const), |
| 182 | + }), |
| 183 | + async execute(_toolCallId, params, _signal, _onUpdate, _ctx) { |
| 184 | + switch (params.action) { |
| 185 | + case "status": { |
| 186 | + const nextIn = timer |
| 187 | + ? `~${Math.round(computeBackoffMs(state.consecutiveErrors, state.intervalMs) / 1000)}s` |
| 188 | + : "paused"; |
| 189 | + return { |
| 190 | + content: [ |
| 191 | + { |
| 192 | + type: "text" as const, |
| 193 | + text: [ |
| 194 | + `Heartbeat Status:`, |
| 195 | + ` Enabled: ${state.enabled ? "✅" : "⏹"}`, |
| 196 | + ` Interval: ${state.intervalMs / 1000}s`, |
| 197 | + ` Next fire: ${nextIn}`, |
| 198 | + ` Total runs: ${state.totalRuns}`, |
| 199 | + ` Consecutive errors: ${state.consecutiveErrors}`, |
| 200 | + ` Last run: ${state.lastRunAt ? new Date(state.lastRunAt).toISOString() : "never"}`, |
| 201 | + ` Checklist: ${state.heartbeatFile}`, |
| 202 | + ].join("\n"), |
| 203 | + }, |
| 204 | + ], |
| 205 | + }; |
| 206 | + } |
| 207 | + |
| 208 | + case "pause": { |
| 209 | + state.enabled = false; |
| 210 | + stopTimer(); |
| 211 | + saveState(); |
| 212 | + return { |
| 213 | + content: [{ type: "text" as const, text: "⏹ Heartbeat paused." }], |
| 214 | + }; |
| 215 | + } |
| 216 | + |
| 217 | + case "resume": { |
| 218 | + state.enabled = true; |
| 219 | + state.consecutiveErrors = 0; |
| 220 | + armTimer(); |
| 221 | + saveState(); |
| 222 | + return { |
| 223 | + content: [ |
| 224 | + { |
| 225 | + type: "text" as const, |
| 226 | + text: `✅ Heartbeat resumed (every ${state.intervalMs / 1000}s).`, |
| 227 | + }, |
| 228 | + ], |
| 229 | + }; |
| 230 | + } |
| 231 | + |
| 232 | + case "trigger": { |
| 233 | + fireHeartbeat(); |
| 234 | + return { |
| 235 | + content: [{ type: "text" as const, text: "🫀 Heartbeat triggered." }], |
| 236 | + }; |
| 237 | + } |
| 238 | + |
| 239 | + case "config": { |
| 240 | + const content = readHeartbeatFile(state.heartbeatFile); |
| 241 | + return { |
| 242 | + content: [ |
| 243 | + { |
| 244 | + type: "text" as const, |
| 245 | + text: [ |
| 246 | + `Heartbeat Configuration:`, |
| 247 | + ` File: ${state.heartbeatFile}`, |
| 248 | + ` File exists: ${existsSync(state.heartbeatFile) ? "yes" : "no"}`, |
| 249 | + ` Has content: ${content ? "yes" : "no (empty or comments only)"}`, |
| 250 | + ` Interval: ${state.intervalMs / 1000}s (env: HEARTBEAT_INTERVAL_MS)`, |
| 251 | + ` Min interval: ${MIN_INTERVAL_MS / 1000}s`, |
| 252 | + ` Backoff multiplier: ${BACKOFF_MULTIPLIER}x per error`, |
| 253 | + ` Max backoff: ${MAX_BACKOFF_MS / 1000}s`, |
| 254 | + ``, |
| 255 | + content ? `Current checklist:\n${content}` : `(no checklist loaded)`, |
| 256 | + ].join("\n"), |
| 257 | + }, |
| 258 | + ], |
| 259 | + }; |
| 260 | + } |
| 261 | + |
| 262 | + default: |
| 263 | + return { |
| 264 | + content: [{ type: "text" as const, text: `Unknown action: ${params.action}` }], |
| 265 | + }; |
| 266 | + } |
| 267 | + }, |
| 268 | + }); |
| 269 | + |
| 270 | + // ── Lifecycle ───────────────────────────────────────────────────────────── |
| 271 | + |
| 272 | + pi.on("session_start", async (_event, ctx) => { |
| 273 | + // Restore persisted state |
| 274 | + for (const entry of ctx.sessionManager.getEntries()) { |
| 275 | + const e = entry as { type: string; customType?: string; data?: any }; |
| 276 | + if (e.type === "custom" && e.customType === HEARTBEAT_STATE_ENTRY && e.data) { |
| 277 | + if (typeof e.data.consecutiveErrors === "number") |
| 278 | + state.consecutiveErrors = e.data.consecutiveErrors; |
| 279 | + if (typeof e.data.totalRuns === "number") state.totalRuns = e.data.totalRuns; |
| 280 | + if (typeof e.data.lastRunAt === "number") state.lastRunAt = e.data.lastRunAt; |
| 281 | + } |
| 282 | + } |
| 283 | + |
| 284 | + // Apply env config |
| 285 | + const config = resolveConfig(); |
| 286 | + state.intervalMs = config.intervalMs; |
| 287 | + state.heartbeatFile = config.heartbeatFile; |
| 288 | + state.enabled = config.enabled; |
| 289 | + |
| 290 | + if (state.enabled) { |
| 291 | + armTimer(); |
| 292 | + } |
| 293 | + }); |
| 294 | + |
| 295 | + pi.on("session_shutdown", async () => { |
| 296 | + stopTimer(); |
| 297 | + }); |
| 298 | +} |
0 commit comments