fix: gate heartbeat prompt to default agent sessions (#630) (thanks @adam91holt)

src/agents/pi-embedded-runner.test.ts

@@ -2,10 +2,12 @@ import type { AgentMessage, AgentTool } from "@mariozechner/pi-agent-core";
 import { SessionManager } from "@mariozechner/pi-coding-agent";
 import { Type } from "@sinclair/typebox";
 import { describe, expect, it, vi } from "vitest";
+import type { ClawdbotConfig } from "../config/config.js";
 import {
   applyGoogleTurnOrderingFix,
   buildEmbeddedSandboxInfo,
   createSystemPromptOverride,
+  resolveSessionAgentIds,
   splitSdkTools,
 } from "./pi-embedded-runner.js";
 import type { SandboxContext } from "./sandbox.js";
@@ -57,6 +59,38 @@ describe("buildEmbeddedSandboxInfo", () => {
   });
 });
 
+describe("resolveSessionAgentIds", () => {
+  const cfg = {
+    agents: {
+      list: [{ id: "main" }, { id: "beta", default: true }],
+    },
+  } as ClawdbotConfig;
+
+  it("falls back to the configured default when sessionKey is missing", () => {
+    const { defaultAgentId, sessionAgentId } = resolveSessionAgentIds({
+      config: cfg,
+    });
+    expect(defaultAgentId).toBe("beta");
+    expect(sessionAgentId).toBe("beta");
+  });
+
+  it("falls back to the configured default when sessionKey is non-agent", () => {
+    const { sessionAgentId } = resolveSessionAgentIds({
+      sessionKey: "telegram:slash:123",
+      config: cfg,
+    });
+    expect(sessionAgentId).toBe("beta");
+  });
+
+  it("uses the agent id from agent session keys", () => {
+    const { sessionAgentId } = resolveSessionAgentIds({
+      sessionKey: "agent:main:main",
+      config: cfg,
+    });
+    expect(sessionAgentId).toBe("main");
+  });
+});
+
 function createStubTool(name: string): AgentTool {
   return {
     name,

src/agents/pi-embedded-runner.ts

@@ -33,7 +33,10 @@ import {
   type enqueueCommand,
   enqueueCommandInLane,
 } from "../process/command-queue.js";
-import { resolveAgentIdFromSessionKey } from "../routing/session-key.js";
+import {
+  normalizeAgentId,
+  parseAgentSessionKey,
+} from "../routing/session-key.js";
 import { normalizeMessageProvider } from "../utils/message-provider.js";
 import { resolveUserPath } from "../utils.js";
 import { resolveClawdbotAgentDir } from "./agent-paths.js";
@@ -557,6 +560,19 @@ export function buildEmbeddedSandboxInfo(
   };
 }
 
+export function resolveSessionAgentIds(params: {
+  sessionKey?: string;
+  config?: ClawdbotConfig;
+}): { defaultAgentId: string; sessionAgentId: string } {
+  const defaultAgentId = resolveDefaultAgentId(params.config ?? {});
+  const sessionKey = params.sessionKey?.trim();
+  const parsed = sessionKey ? parseAgentSessionKey(sessionKey) : null;
+  const sessionAgentId = parsed?.agentId
+    ? normalizeAgentId(parsed.agentId)
+    : defaultAgentId;
+  return { defaultAgentId, sessionAgentId };
+}
+
 function buildEmbeddedSystemPrompt(params: {
   workspaceDir: string;
   defaultThinkLevel?: ThinkLevel;
@@ -885,8 +901,10 @@ export async function compactEmbeddedPiSession(params: {
         );
         const userTime = formatUserTime(new Date(), userTimezone);
         // Only include heartbeat prompt for the default agent
-        const sessionAgentId = resolveAgentIdFromSessionKey(params.sessionKey);
-        const defaultAgentId = resolveDefaultAgentId(params.config ?? {});
+        const { defaultAgentId, sessionAgentId } = resolveSessionAgentIds({
+          sessionKey: params.sessionKey,
+          config: params.config,
+        });
         const isDefaultAgent = sessionAgentId === defaultAgentId;
         const appendPrompt = buildEmbeddedSystemPrompt({
           workspaceDir: effectiveWorkspace,
@@ -1254,10 +1272,10 @@ export async function runEmbeddedPiAgent(params: {
           );
           const userTime = formatUserTime(new Date(), userTimezone);
           // Only include heartbeat prompt for the default agent
-          const sessionAgentId = resolveAgentIdFromSessionKey(
-            params.sessionKey,
-          );
-          const defaultAgentId = resolveDefaultAgentId(params.config ?? {});
+          const { defaultAgentId, sessionAgentId } = resolveSessionAgentIds({
+            sessionKey: params.sessionKey,
+            config: params.config,
+          });
           const isDefaultAgent = sessionAgentId === defaultAgentId;
           const appendPrompt = buildEmbeddedSystemPrompt({
             workspaceDir: effectiveWorkspace,