fix: refine HEARTBEAT_OK handling

src/auto-reply/heartbeat.test.ts

@@ -5,55 +5,89 @@ import { HEARTBEAT_TOKEN } from "./tokens.js";
 
 describe("stripHeartbeatToken", () => {
   it("skips empty or token-only replies", () => {
-    expect(stripHeartbeatToken(undefined)).toEqual({
+    expect(stripHeartbeatToken(undefined, { mode: "heartbeat" })).toEqual({
       shouldSkip: true,
       text: "",
+      didStrip: false,
     });
-    expect(stripHeartbeatToken("  ")).toEqual({
+    expect(stripHeartbeatToken("  ", { mode: "heartbeat" })).toEqual({
       shouldSkip: true,
       text: "",
+      didStrip: false,
     });
-    expect(stripHeartbeatToken(HEARTBEAT_TOKEN)).toEqual({
+    expect(stripHeartbeatToken(HEARTBEAT_TOKEN, { mode: "heartbeat" })).toEqual(
+      {
+        shouldSkip: true,
+        text: "",
+        didStrip: true,
+      },
+    );
+  });
+
+  it("drops heartbeats with small junk in heartbeat mode", () => {
+    expect(
+      stripHeartbeatToken("HEARTBEAT_OK 🦞", { mode: "heartbeat" }),
+    ).toEqual({
       shouldSkip: true,
       text: "",
+      didStrip: true,
+    });
+    expect(
+      stripHeartbeatToken(`🦞 ${HEARTBEAT_TOKEN}`, { mode: "heartbeat" }),
+    ).toEqual({
+      shouldSkip: true,
+      text: "",
+      didStrip: true,
     });
   });
 
-  it("skips any reply that includes the heartbeat token", () => {
-    expect(stripHeartbeatToken(`ALERT ${HEARTBEAT_TOKEN}`)).toEqual({
-      shouldSkip: true,
-      text: "",
-    });
-    expect(stripHeartbeatToken("HEARTBEAT_OK 🦞")).toEqual({
-      shouldSkip: true,
-      text: "",
-    });
-    expect(stripHeartbeatToken("HEARTBEAT_OK_OK_OK")).toEqual({
-      shouldSkip: true,
-      text: "",
-    });
-    expect(stripHeartbeatToken("HEARTBEAT_OK_OK")).toEqual({
-      shouldSkip: true,
-      text: "",
-    });
-    expect(stripHeartbeatToken("HEARTBEAT_OK _OK")).toEqual({
-      shouldSkip: true,
-      text: "",
-    });
-    expect(stripHeartbeatToken("HEARTBEAT_OK OK")).toEqual({
-      shouldSkip: true,
-      text: "",
-    });
-    expect(stripHeartbeatToken("ALERT HEARTBEAT_OK_OK")).toEqual({
+  it("drops short remainder in heartbeat mode", () => {
+    expect(
+      stripHeartbeatToken(`ALERT ${HEARTBEAT_TOKEN}`, { mode: "heartbeat" }),
+    ).toEqual({
       shouldSkip: true,
       text: "",
+      didStrip: true,
     });
   });
 
-  it("keeps non-heartbeat content", () => {
-    expect(stripHeartbeatToken("hello")).toEqual({
+  it("keeps heartbeat replies when remaining content exceeds threshold", () => {
+    const long = "A".repeat(31);
+    expect(
+      stripHeartbeatToken(`${long} ${HEARTBEAT_TOKEN}`, { mode: "heartbeat" }),
+    ).toEqual({
+      shouldSkip: false,
+      text: long,
+      didStrip: true,
+    });
+  });
+
+  it("strips token at edges for normal messages", () => {
+    expect(
+      stripHeartbeatToken(`${HEARTBEAT_TOKEN} hello`, { mode: "message" }),
+    ).toEqual({
       shouldSkip: false,
       text: "hello",
+      didStrip: true,
+    });
+    expect(
+      stripHeartbeatToken(`hello ${HEARTBEAT_TOKEN}`, { mode: "message" }),
+    ).toEqual({
+      shouldSkip: false,
+      text: "hello",
+      didStrip: true,
+    });
+  });
+
+  it("does not touch token in the middle", () => {
+    expect(
+      stripHeartbeatToken(`hello ${HEARTBEAT_TOKEN} there`, {
+        mode: "message",
+      }),
+    ).toEqual({
+      shouldSkip: false,
+      text: `hello ${HEARTBEAT_TOKEN} there`,
+      didStrip: false,
     });
   });
 });

src/auto-reply/heartbeat.ts

@@ -2,12 +2,69 @@ import { HEARTBEAT_TOKEN } from "./tokens.js";
 
 export const HEARTBEAT_PROMPT = "HEARTBEAT";
 
-export function stripHeartbeatToken(raw?: string) {
-  if (!raw) return { shouldSkip: true, text: "" };
-  const trimmed = raw.trim();
-  if (!trimmed) return { shouldSkip: true, text: "" };
-  if (trimmed.includes(HEARTBEAT_TOKEN)) {
-    return { shouldSkip: true, text: "" };
+export type StripHeartbeatMode = "heartbeat" | "message";
+
+function stripTokenAtEdges(raw: string): { text: string; didStrip: boolean } {
+  let text = raw.trim();
+  if (!text) return { text: "", didStrip: false };
+
+  const token = HEARTBEAT_TOKEN;
+  if (!text.includes(token)) return { text, didStrip: false };
+
+  let didStrip = false;
+  let changed = true;
+  while (changed) {
+    changed = false;
+    const next = text.trim();
+    if (next.startsWith(token)) {
+      const after = next.slice(token.length).trimStart();
+      text = after;
+      didStrip = true;
+      changed = true;
+      continue;
+    }
+    if (next.endsWith(token)) {
+      const before = next.slice(0, Math.max(0, next.length - token.length));
+      text = before.trimEnd();
+      didStrip = true;
+      changed = true;
+    }
   }
-  return { shouldSkip: false, text: trimmed };
+
+  const collapsed = text.replace(/\s+/g, " ").trim();
+  return { text: collapsed, didStrip };
+}
+
+export function stripHeartbeatToken(
+  raw?: string,
+  opts: { mode?: StripHeartbeatMode; maxAckChars?: number } = {},
+) {
+  if (!raw) return { shouldSkip: true, text: "", didStrip: false };
+  const trimmed = raw.trim();
+  if (!trimmed) return { shouldSkip: true, text: "", didStrip: false };
+
+  const mode: StripHeartbeatMode = opts.mode ?? "message";
+  const maxAckChars = Math.max(0, opts.maxAckChars ?? 30);
+
+  if (!trimmed.includes(HEARTBEAT_TOKEN)) {
+    return { shouldSkip: false, text: trimmed, didStrip: false };
+  }
+
+  const stripped = stripTokenAtEdges(trimmed);
+  if (!stripped.didStrip) {
+    return { shouldSkip: false, text: trimmed, didStrip: false };
+  }
+
+  if (!stripped.text) {
+    return { shouldSkip: true, text: "", didStrip: true };
+  }
+
+  if (mode === "heartbeat") {
+    const rest = stripped.text.trim();
+    if (rest.length <= maxAckChars) {
+      return { shouldSkip: true, text: "", didStrip: true };
+    }
+  }
+
+  return { shouldSkip: false, text: stripped.text, didStrip: true };
 }

src/auto-reply/reply.triggers.test.ts

@@ -186,6 +186,31 @@ describe("trigger handling", () => {
     });
   });
 
+  it("strips HEARTBEAT_OK at edges outside heartbeat runs", async () => {
+    await withTempHome(async (home) => {
+      vi.mocked(runEmbeddedPiAgent).mockResolvedValue({
+        payloads: [{ text: `${HEARTBEAT_TOKEN} hello` }],
+        meta: {
+          durationMs: 1,
+          agentMeta: { sessionId: "s", provider: "p", model: "m" },
+        },
+      });
+
+      const res = await getReplyFromConfig(
+        {
+          Body: "hello",
+          From: "+1002",
+          To: "+2000",
+        },
+        {},
+        makeCfg(home),
+      );
+
+      const text = Array.isArray(res) ? res[0]?.text : res?.text;
+      expect(text).toBe("hello");
+    });
+  });
+
   it("updates group activation when the owner sends /activation", async () => {
     await withTempHome(async (home) => {
       const cfg = makeCfg(home);

src/auto-reply/reply.ts

@@ -53,6 +53,7 @@ import {
   normalizeGroupActivation,
   parseActivationCommand,
 } from "./group-activation.js";
+import { stripHeartbeatToken } from "./heartbeat.js";
 import { extractModelDirective } from "./model.js";
 import { buildStatusMessage } from "./status.js";
 import type { MsgContext, TemplateContext } from "./templating.js";
@@ -62,7 +63,7 @@ import {
   type ThinkLevel,
   type VerboseLevel,
 } from "./thinking.js";
-import { HEARTBEAT_TOKEN, SILENT_REPLY_TOKEN } from "./tokens.js";
+import { SILENT_REPLY_TOKEN } from "./tokens.js";
 import { isAudio, transcribeInboundAudio } from "./transcription.js";
 import type { GetReplyOptions, ReplyPayload } from "./types.js";
 
@@ -1191,7 +1192,7 @@ export async function getReplyFromConfig(
     return undefined;
   }
 
-  let suppressedByHeartbeatAck = false;
+  let didLogHeartbeatStrip = false;
   try {
     if (shouldEagerType) {
       await startTypingLoop();
@@ -1221,20 +1222,24 @@ export async function getReplyFromConfig(
         runId,
         onPartialReply: opts?.onPartialReply
           ? async (payload) => {
-              if (
-                !opts?.isHeartbeat &&
-                payload.text?.includes(HEARTBEAT_TOKEN)
-              ) {
-                suppressedByHeartbeatAck = true;
-                logVerbose(
-                  "Suppressing partial reply: detected HEARTBEAT_OK token",
-                );
-                return;
+              let text = payload.text;
+              if (!opts?.isHeartbeat && text?.includes("HEARTBEAT_OK")) {
+                const stripped = stripHeartbeatToken(text, { mode: "message" });
+                if (stripped.didStrip && !didLogHeartbeatStrip) {
+                  didLogHeartbeatStrip = true;
+                  logVerbose("Stripped stray HEARTBEAT_OK token from reply");
+                }
+                if (
+                  stripped.shouldSkip &&
+                  (payload.mediaUrls?.length ?? 0) === 0
+                ) {
+                  return;
+                }
+                text = stripped.text;
               }
-              if (suppressedByHeartbeatAck) return;
-              await startTypingOnText(payload.text);
+              await startTypingOnText(text);
               await opts.onPartialReply?.({
-                text: payload.text,
+                text,
                 mediaUrls: payload.mediaUrls,
               });
             }
@@ -1242,22 +1247,23 @@ export async function getReplyFromConfig(
         shouldEmitToolResult,
         onToolResult: opts?.onToolResult
           ? async (payload) => {
-              if (
-                !opts?.isHeartbeat &&
-                payload.text?.includes(HEARTBEAT_TOKEN)
-              ) {
-                suppressedByHeartbeatAck = true;
-                logVerbose(
-                  "Suppressing tool result: detected HEARTBEAT_OK token",
-                );
-                return;
+              let text = payload.text;
+              if (!opts?.isHeartbeat && text?.includes("HEARTBEAT_OK")) {
+                const stripped = stripHeartbeatToken(text, { mode: "message" });
+                if (stripped.didStrip && !didLogHeartbeatStrip) {
+                  didLogHeartbeatStrip = true;
+                  logVerbose("Stripped stray HEARTBEAT_OK token from reply");
+                }
+                if (
+                  stripped.shouldSkip &&
+                  (payload.mediaUrls?.length ?? 0) === 0
+                ) {
+                  return;
+                }
+                text = stripped.text;
               }
-              if (suppressedByHeartbeatAck) return;
-              await startTypingOnText(payload.text);
-              await opts.onToolResult?.({
-                text: payload.text,
-                mediaUrls: payload.mediaUrls,
-              });
+              await startTypingOnText(text);
+              await opts.onToolResult?.({ text, mediaUrls: payload.mediaUrls });
             }
           : undefined,
       });
@@ -1288,22 +1294,28 @@ export async function getReplyFromConfig(
 
     const payloadArray = runResult.payloads ?? [];
     if (payloadArray.length === 0) return undefined;
-    if (
-      suppressedByHeartbeatAck ||
-      (!opts?.isHeartbeat &&
-        payloadArray.some((payload) => payload.text?.includes(HEARTBEAT_TOKEN)))
-    ) {
-      logVerbose("Suppressing reply: detected HEARTBEAT_OK token");
-      return undefined;
-    }
-    const shouldSignalTyping = payloadArray.some((payload) => {
+
+    const sanitizedPayloads = opts?.isHeartbeat
+      ? payloadArray
+      : payloadArray.flatMap((payload) => {
+          const text = payload.text;
+          if (!text || !text.includes("HEARTBEAT_OK")) return [payload];
+          const stripped = stripHeartbeatToken(text, { mode: "message" });
+          if (stripped.didStrip && !didLogHeartbeatStrip) {
+            didLogHeartbeatStrip = true;
+            logVerbose("Stripped stray HEARTBEAT_OK token from reply");
+          }
+          const hasMedia =
+            Boolean(payload.mediaUrl) || (payload.mediaUrls?.length ?? 0) > 0;
+          if (stripped.shouldSkip && !hasMedia) return [];
+          return [{ ...payload, text: stripped.text }];
+        });
+
+    if (sanitizedPayloads.length === 0) return undefined;
+
+    const shouldSignalTyping = sanitizedPayloads.some((payload) => {
       const trimmed = payload.text?.trim();
-      if (
-        trimmed &&
-        trimmed !== SILENT_REPLY_TOKEN &&
-        !trimmed.includes(HEARTBEAT_TOKEN)
-      )
-        return true;
+      if (trimmed && trimmed !== SILENT_REPLY_TOKEN) return true;
       if (payload.mediaUrl) return true;
       if (payload.mediaUrls && payload.mediaUrls.length > 0) return true;
       return false;
@@ -1356,11 +1368,11 @@ export async function getReplyFromConfig(
     }
 
     // If verbose is enabled and this is a new session, prepend a session hint.
-    let finalPayloads = payloadArray;
+    let finalPayloads = sanitizedPayloads;
     if (resolvedVerboseLevel === "on" && isNewSession) {
       finalPayloads = [
         { text: `🧭 New session: ${sessionIdFinal}` },
-        ...payloadArray,
+        ...finalPayloads,
       ];
     }