fix: stop heartbeat transcript truncation races (#60998) (thanks @nxmxbbd)

src/auto-reply/heartbeat-filter.test.ts

@@ -0,0 +1,141 @@
+import { describe, expect, it } from "vitest";
+import {
+  filterHeartbeatPairs,
+  isHeartbeatOkResponse,
+  isHeartbeatUserMessage,
+} from "./heartbeat-filter.js";
+import { HEARTBEAT_PROMPT } from "./heartbeat.js";
+
+describe("isHeartbeatUserMessage", () => {
+  it("matches heartbeat prompts", () => {
+    expect(
+      isHeartbeatUserMessage(
+        {
+          role: "user",
+          content: `${HEARTBEAT_PROMPT}\nWhen reading HEARTBEAT.md, use workspace file /tmp/HEARTBEAT.md (exact case). Do not read docs/heartbeat.md.`,
+        },
+        HEARTBEAT_PROMPT,
+      ),
+    ).toBe(true);
+
+    expect(
+      isHeartbeatUserMessage({
+        role: "user",
+        content:
+          "Run the following periodic tasks (only those due based on their intervals):\n\n- email-check: Check for urgent unread emails\n\nAfter completing all due tasks, reply HEARTBEAT_OK.",
+      }),
+    ).toBe(true);
+  });
+
+  it("ignores quoted or non-user token mentions", () => {
+    expect(
+      isHeartbeatUserMessage({
+        role: "user",
+        content: "Please reply HEARTBEAT_OK so I can test something.",
+      }),
+    ).toBe(false);
+
+    expect(
+      isHeartbeatUserMessage({
+        role: "assistant",
+        content: "HEARTBEAT_OK",
+      }),
+    ).toBe(false);
+  });
+});
+
+describe("isHeartbeatOkResponse", () => {
+  it("matches no-op heartbeat acknowledgements", () => {
+    expect(
+      isHeartbeatOkResponse({
+        role: "assistant",
+        content: "**HEARTBEAT_OK**",
+      }),
+    ).toBe(true);
+
+    expect(
+      isHeartbeatOkResponse({
+        role: "assistant",
+        content: "You have 3 unread urgent emails. HEARTBEAT_OK",
+      }),
+    ).toBe(true);
+  });
+
+  it("preserves meaningful or non-text responses", () => {
+    expect(
+      isHeartbeatOkResponse({
+        role: "assistant",
+        content: "Status HEARTBEAT_OK due to watchdog failure",
+      }),
+    ).toBe(false);
+
+    expect(
+      isHeartbeatOkResponse({
+        role: "assistant",
+        content: [{ type: "tool_use", id: "tool-1", name: "search", input: {} }],
+      }),
+    ).toBe(false);
+  });
+
+  it("respects ackMaxChars overrides", () => {
+    expect(
+      isHeartbeatOkResponse(
+        {
+          role: "assistant",
+          content: "HEARTBEAT_OK all good",
+        },
+        0,
+      ),
+    ).toBe(false);
+  });
+});
+
+describe("filterHeartbeatPairs", () => {
+  it("removes no-op heartbeat pairs", () => {
+    const messages = [
+      { role: "user", content: "Hello" },
+      { role: "assistant", content: "Hi there!" },
+      { role: "user", content: HEARTBEAT_PROMPT },
+      { role: "assistant", content: "HEARTBEAT_OK" },
+      { role: "user", content: "What time is it?" },
+      { role: "assistant", content: "It is 3pm." },
+    ];
+
+    expect(filterHeartbeatPairs(messages, undefined, HEARTBEAT_PROMPT)).toEqual([
+      { role: "user", content: "Hello" },
+      { role: "assistant", content: "Hi there!" },
+      { role: "user", content: "What time is it?" },
+      { role: "assistant", content: "It is 3pm." },
+    ]);
+  });
+
+  it("keeps meaningful heartbeat results and non-text assistant turns", () => {
+    const meaningfulMessages = [
+      { role: "user", content: HEARTBEAT_PROMPT },
+      { role: "assistant", content: "Status HEARTBEAT_OK due to watchdog failure" },
+    ];
+    expect(filterHeartbeatPairs(meaningfulMessages, undefined, HEARTBEAT_PROMPT)).toEqual(
+      meaningfulMessages,
+    );
+
+    const nonTextMessages = [
+      { role: "user", content: HEARTBEAT_PROMPT },
+      {
+        role: "assistant",
+        content: [{ type: "tool_use", id: "tool-1", name: "search", input: {} }],
+      },
+    ];
+    expect(filterHeartbeatPairs(nonTextMessages, undefined, HEARTBEAT_PROMPT)).toEqual(
+      nonTextMessages,
+    );
+  });
+
+  it("keeps ordinary chats that mention the token", () => {
+    const messages = [
+      { role: "user", content: "Please reply HEARTBEAT_OK so I can test something." },
+      { role: "assistant", content: "HEARTBEAT_OK" },
+    ];
+
+    expect(filterHeartbeatPairs(messages, undefined, HEARTBEAT_PROMPT)).toEqual(messages);
+  });
+});

src/auto-reply/heartbeat-filter.ts

@@ -0,0 +1,96 @@
+import { stripHeartbeatToken } from "./heartbeat.js";
+
+const HEARTBEAT_TASK_PROMPT_PREFIX =
+  "Run the following periodic tasks (only those due based on their intervals):";
+const HEARTBEAT_TASK_PROMPT_ACK = "After completing all due tasks, reply HEARTBEAT_OK.";
+
+function resolveMessageText(content: unknown): { text: string; hasNonTextContent: boolean } {
+  if (typeof content === "string") {
+    return { text: content, hasNonTextContent: false };
+  }
+  if (!Array.isArray(content)) {
+    return { text: "", hasNonTextContent: content != null };
+  }
+  let hasNonTextContent = false;
+  const text = content
+    .filter((block): block is { type: "text"; text: string } => {
+      if (typeof block !== "object" || block === null || !("type" in block)) {
+        hasNonTextContent = true;
+        return false;
+      }
+      if (block.type !== "text") {
+        hasNonTextContent = true;
+        return false;
+      }
+      if (typeof (block as { text?: unknown }).text !== "string") {
+        hasNonTextContent = true;
+        return false;
+      }
+      return true;
+    })
+    .map((block) => block.text)
+    .join("");
+  return { text, hasNonTextContent };
+}
+
+export function isHeartbeatUserMessage(
+  message: { role: string; content?: unknown },
+  heartbeatPrompt?: string,
+): boolean {
+  if (message.role !== "user") {
+    return false;
+  }
+  const { text } = resolveMessageText(message.content);
+  const trimmed = text.trim();
+  if (!trimmed) {
+    return false;
+  }
+  const normalizedHeartbeatPrompt = heartbeatPrompt?.trim();
+  if (normalizedHeartbeatPrompt && trimmed.startsWith(normalizedHeartbeatPrompt)) {
+    return true;
+  }
+  return (
+    trimmed.startsWith(HEARTBEAT_TASK_PROMPT_PREFIX) && trimmed.includes(HEARTBEAT_TASK_PROMPT_ACK)
+  );
+}
+
+export function isHeartbeatOkResponse(
+  message: { role: string; content?: unknown },
+  ackMaxChars?: number,
+): boolean {
+  if (message.role !== "assistant") {
+    return false;
+  }
+  const { text, hasNonTextContent } = resolveMessageText(message.content);
+  if (hasNonTextContent) {
+    return false;
+  }
+  return stripHeartbeatToken(text, { mode: "heartbeat", maxAckChars: ackMaxChars }).shouldSkip;
+}
+
+export function filterHeartbeatPairs<T extends { role: string; content?: unknown }>(
+  messages: T[],
+  ackMaxChars?: number,
+  heartbeatPrompt?: string,
+): T[] {
+  if (messages.length < 2) {
+    return messages;
+  }
+
+  const result: T[] = [];
+  let i = 0;
+  while (i < messages.length) {
+    if (
+      i + 1 < messages.length &&
+      isHeartbeatUserMessage(messages[i], heartbeatPrompt) &&
+      isHeartbeatOkResponse(messages[i + 1], ackMaxChars)
+    ) {
+      i += 2;
+      continue;
+    }
+    result.push(messages[i]);
+    i++;
+  }
+
+  return result;
+}