feat(heartbeat): improve prompting, suppression, delivery targets and session cleanup

- Rewrite heartbeat prompt to instruct agent to EXECUTE checklist tasks, not echo them
- Simplify suppression: HEARTBEAT_OK present = always suppress, absent = always deliver
- Add delivery targets RPC (heartbeat.targets) for channel/chatId picker
- Sanitize backend errors — never expose raw SQL to client
- Add session cleanup for isolated heartbeat sessions after run
- Cap StaggerOffset at 10% of interval to avoid user-visible delay
- Fix Upsert to persist next_run_at correctly
This commit is contained in:
viettranx
2026-03-18 16:37:36 +07:00
parent b6f6af01d2
commit 96cfd1bf08
6 changed files with 155 additions and 31 deletions
+1
View File
@@ -711,6 +711,7 @@ func runGateway() {
heartbeatTicker := heartbeat.NewTicker(heartbeat.TickerConfig{
Store: pgStores.Heartbeats,
Agents: pgStores.Agents,
Sessions: pgStores.Sessions,
MsgBus: msgBus,
Sched: sched,
RunAgent: makeHeartbeatRunFn(sched),
+45 -8
View File
@@ -47,6 +47,7 @@ func (m *HeartbeatMethods) Register(router *gateway.MethodRouter) {
router.Register(protocol.MethodHeartbeatLogs, m.handleLogs)
router.Register(protocol.MethodHeartbeatChecklistGet, m.handleChecklistGet)
router.Register(protocol.MethodHeartbeatChecklistSet, m.handleChecklistSet)
router.Register(protocol.MethodHeartbeatTargets, m.handleTargets)
}
func (m *HeartbeatMethods) handleGet(ctx context.Context, client *gateway.Client, req *protocol.RequestFrame) {
@@ -74,7 +75,7 @@ func (m *HeartbeatMethods) handleGet(ctx context.Context, client *gateway.Client
client.SendResponse(protocol.NewOKResponse(req.ID, map[string]any{"heartbeat": nil}))
return
}
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, err.Error()))
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("get", err)))
return
}
@@ -116,7 +117,7 @@ func (m *HeartbeatMethods) handleSet(ctx context.Context, client *gateway.Client
// Load existing or create new.
hb, err := m.hbStore.Get(ctx, agentUUID)
if err != nil && !errors.Is(err, sql.ErrNoRows) {
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, err.Error()))
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("set.load", err)))
return
}
if hb == nil {
@@ -187,7 +188,7 @@ func (m *HeartbeatMethods) handleSet(ctx context.Context, client *gateway.Client
}
if err := m.hbStore.Upsert(ctx, hb); err != nil {
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, err.Error()))
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("set.upsert", err)))
return
}
@@ -222,7 +223,7 @@ func (m *HeartbeatMethods) handleToggle(ctx context.Context, client *gateway.Cli
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrNotFound, "heartbeat not configured"))
return
}
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, err.Error()))
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("op", err)))
return
}
@@ -233,7 +234,7 @@ func (m *HeartbeatMethods) handleToggle(ctx context.Context, client *gateway.Cli
}
if err := m.hbStore.Upsert(ctx, hb); err != nil {
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, err.Error()))
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("op", err)))
return
}
@@ -301,7 +302,7 @@ func (m *HeartbeatMethods) handleLogs(ctx context.Context, client *gateway.Clien
logs, total, err := m.hbStore.ListLogs(ctx, agentUUID, params.Limit, params.Offset)
if err != nil {
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, err.Error()))
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("op", err)))
return
}
@@ -336,7 +337,7 @@ func (m *HeartbeatMethods) handleChecklistGet(ctx context.Context, client *gatew
files, err := m.agentStore.GetAgentContextFiles(ctx, agentUUID)
if err != nil {
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, err.Error()))
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("op", err)))
return
}
@@ -378,7 +379,7 @@ func (m *HeartbeatMethods) handleChecklistSet(ctx context.Context, client *gatew
}
if err := m.agentStore.SetAgentContextFile(ctx, agentUUID, "HEARTBEAT.md", params.Content); err != nil {
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, err.Error()))
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("op", err)))
return
}
@@ -389,6 +390,42 @@ func (m *HeartbeatMethods) handleChecklistSet(ctx context.Context, client *gatew
emitAudit(m.eventBus, client, "heartbeat.checklist.set", "heartbeat", params.AgentID)
}
func (m *HeartbeatMethods) handleTargets(ctx context.Context, client *gateway.Client, req *protocol.RequestFrame) {
locale := store.LocaleFromContext(ctx)
var params struct {
AgentID string `json:"agentId"`
}
if req.Params != nil {
json.Unmarshal(req.Params, &params)
}
if params.AgentID == "" {
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, i18n.T(locale, i18n.MsgRequired, "agentId")))
return
}
agentUUID, err := uuid.Parse(params.AgentID)
if err != nil {
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, "invalid agentId"))
return
}
targets, err := m.hbStore.ListDeliveryTargets(ctx, agentUUID)
if err != nil {
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("targets", err)))
return
}
client.SendResponse(protocol.NewOKResponse(req.ID, map[string]any{
"targets": targets,
}))
}
// heartbeatInternalErr logs the real error and returns a safe message for the client.
func heartbeatInternalErr(action string, err error) string {
slog.Error("heartbeat RPC error", "action", action, "error", err)
return "internal error"
}
func (m *HeartbeatMethods) emitCacheInvalidate(agentID string) {
m.eventBus.Broadcast(bus.Event{
Name: protocol.EventCacheInvalidate,
+29 -18
View File
@@ -30,6 +30,7 @@ const (
type TickerConfig struct {
Store store.HeartbeatStore
Agents store.AgentStore
Sessions store.SessionStore // optional: for cleaning up isolated heartbeat sessions
MsgBus *bus.MessageBus
Sched *scheduler.Scheduler
RunAgent func(ctx context.Context, req agent.RunRequest) <-chan scheduler.RunOutcome
@@ -39,6 +40,7 @@ type TickerConfig struct {
type Ticker struct {
store store.HeartbeatStore
agents store.AgentStore
sessions store.SessionStore
msgBus *bus.MessageBus
sched *scheduler.Scheduler
runAgent func(ctx context.Context, req agent.RunRequest) <-chan scheduler.RunOutcome
@@ -54,6 +56,7 @@ func NewTicker(cfg TickerConfig) *Ticker {
return &Ticker{
store: cfg.Store,
agents: cfg.Agents,
sessions: cfg.Sessions,
msgBus: cfg.MsgBus,
sched: cfg.Sched,
runAgent: cfg.RunAgent,
@@ -183,16 +186,21 @@ func (t *Ticker) runOne(ctx context.Context, hb store.AgentHeartbeat) {
})
// [4] Build prompt.
prompt := "Read HEARTBEAT.md if it exists. Follow it strictly.\nIf nothing needs attention, reply HEARTBEAT_OK."
prompt := "Execute your heartbeat checklist now."
if hb.Prompt != nil && *hb.Prompt != "" {
prompt = *hb.Prompt
}
extraSystem := fmt.Sprintf(
"[Heartbeat Check-in]\nThis is a periodic heartbeat run for agent %s.\n"+
"Your HEARTBEAT.md checklist:\n---\n%s\n---\n"+
"Follow the checklist. If everything is OK, reply with HEARTBEAT_OK.\n"+
"If something needs attention, describe it clearly — your response will be delivered to the configured channel.",
"Your checklist:\n---\n%s\n---\n"+
"RULES:\n"+
"- EXECUTE the tasks in the checklist using your tools. Do NOT just read or quote the checklist back.\n"+
"- Your response will be delivered to the configured channel as-is.\n"+
"- HEARTBEAT_OK suppression: If your response contains the token HEARTBEAT_OK anywhere, "+
"the ENTIRE response is suppressed and NOT delivered to the channel.\n"+
"- Use HEARTBEAT_OK ONLY when there is nothing to deliver (e.g. monitoring checks all passed, no news).\n"+
"- Do NOT include HEARTBEAT_OK if the checklist asks you to send content (jokes, greetings, reports, etc.).",
agentKey, checklistContent,
)
@@ -244,7 +252,7 @@ func (t *Ticker) runOne(ctx context.Context, hb store.AgentHeartbeat) {
// [6] Process result.
if lastErr != nil {
t.finishRun(ctx, hb, agentKey, "error", lastErr.Error(), "", durationMS, 0, 0)
t.finishRun(ctx, hb, sessionKey, agentKey, "error", lastErr.Error(), "", durationMS, 0, 0)
return
}
@@ -258,7 +266,7 @@ func (t *Ticker) runOne(ctx context.Context, hb store.AgentHeartbeat) {
}
if !deliver {
t.finishRun(ctx, hb, agentKey, "suppressed", "", truncate(result.Content, maxSummaryLen), durationMS, inputTokens, outputTokens)
t.finishRun(ctx, hb, sessionKey, agentKey, "suppressed", "", truncate(result.Content, maxSummaryLen), durationMS, inputTokens, outputTokens)
return
}
@@ -271,10 +279,10 @@ func (t *Ticker) runOne(ctx context.Context, hb store.AgentHeartbeat) {
})
}
t.finishRun(ctx, hb, agentKey, "ok", "", truncate(cleaned, maxSummaryLen), durationMS, inputTokens, outputTokens)
t.finishRun(ctx, hb, sessionKey, agentKey, "ok", "", truncate(cleaned, maxSummaryLen), durationMS, inputTokens, outputTokens)
}
func (t *Ticker) finishRun(ctx context.Context, hb store.AgentHeartbeat, agentKey, status, errMsg, summary string, durationMS, inputTokens, outputTokens int) {
func (t *Ticker) finishRun(ctx context.Context, hb store.AgentHeartbeat, sessionKey, agentKey, status, errMsg, summary string, durationMS, inputTokens, outputTokens int) {
agentIDStr := hb.AgentID.String()
now := time.Now()
@@ -319,6 +327,13 @@ func (t *Ticker) finishRun(ctx context.Context, hb store.AgentHeartbeat, agentKe
slog.Warn("heartbeat.update_state_failed", "agent_id", agentIDStr, "error", err)
}
// Cleanup isolated session — data is already in heartbeat_run_logs.
if hb.IsolatedSession && t.sessions != nil && sessionKey != "" {
if err := t.sessions.Delete(sessionKey); err != nil {
slog.Debug("heartbeat.session_cleanup_failed", "session_key", sessionKey, "error", err)
}
}
// Emit event.
t.emitEvent(store.HeartbeatEvent{
Action: status,
@@ -383,18 +398,14 @@ func (t *Ticker) readChecklist(ctx context.Context, agentID uuid.UUID) string {
}
// processResponse implements smart suppression.
// If response contains HEARTBEAT_OK and cleaned content is within threshold, suppress delivery.
func processResponse(response string, ackMaxChars int) (deliver bool, cleaned string) {
// If response contains HEARTBEAT_OK, agent confirms everything is fine — always suppress.
// Only deliver when HEARTBEAT_OK is absent (agent found something needing attention).
func processResponse(response string, _ int) (deliver bool, cleaned string) {
const ackToken = "HEARTBEAT_OK"
if !strings.Contains(response, ackToken) {
return true, response // has alert, deliver as-is
if strings.Contains(response, ackToken) {
return false, "" // agent says OK → suppress regardless of extra content
}
cleaned = strings.ReplaceAll(response, ackToken, "")
cleaned = strings.TrimSpace(cleaned)
if len([]rune(cleaned)) <= ackMaxChars {
return false, "" // suppress
}
return true, cleaned // over threshold, deliver cleaned version
return true, response // no OK token → something needs attention, deliver
}
// isWithinActiveHours checks if current time falls within the configured active hours.
+18 -2
View File
@@ -65,7 +65,8 @@ type HeartbeatRunLog struct {
}
// StaggerOffset returns a deterministic offset for spreading heartbeats evenly.
// Uses FNV-1a hash of agent ID to produce a value in [0, intervalSec).
// Uses FNV-1a hash of agent ID to produce a value in [0, 10% of intervalSec).
// Capped at 10% to avoid user-visible delay while still preventing thundering herd.
func StaggerOffset(agentID uuid.UUID, intervalSec int) time.Duration {
if intervalSec <= 0 {
return 0
@@ -75,7 +76,11 @@ func StaggerOffset(agentID uuid.UUID, intervalSec int) time.Duration {
h ^= uint32(b)
h *= 16777619 // FNV prime
}
offset := int(h) % intervalSec
maxOffset := intervalSec / 10 // 10% of interval
if maxOffset < 1 {
maxOffset = 1
}
offset := int(h) % maxOffset
if offset < 0 {
offset = -offset
}
@@ -92,6 +97,14 @@ type HeartbeatEvent struct {
Reason string `json:"reason,omitempty"` // skip reason
}
// DeliveryTarget represents a known channel+chatID pair from session history.
type DeliveryTarget struct {
Channel string `json:"channel"`
ChatID string `json:"chatId"`
Title string `json:"title,omitempty"` // chat/group title from session metadata
Kind string `json:"kind"` // "dm" or "group"
}
// HeartbeatStore manages agent heartbeat configurations and run logs.
type HeartbeatStore interface {
Get(ctx context.Context, agentID uuid.UUID) (*AgentHeartbeat, error)
@@ -104,6 +117,9 @@ type HeartbeatStore interface {
InsertLog(ctx context.Context, log *HeartbeatRunLog) error
ListLogs(ctx context.Context, agentID uuid.UUID, limit, offset int) ([]HeartbeatRunLog, int, error)
// Delivery targets — distinct (channel, chatID) from session history for an agent.
ListDeliveryTargets(ctx context.Context, agentID uuid.UUID) ([]DeliveryTarget, error)
// Events
SetOnEvent(fn func(HeartbeatEvent))
}
+61 -3
View File
@@ -89,8 +89,8 @@ func (s *PGHeartbeatStore) Upsert(ctx context.Context, hb *store.AgentHeartbeat)
`INSERT INTO agent_heartbeats (agent_id, enabled, interval_sec, prompt, provider_id, model,
isolated_session, light_context, ack_max_chars, max_retries,
active_hours_start, active_hours_end, timezone,
channel, chat_id, metadata, created_at, updated_at)
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$17)
channel, chat_id, next_run_at, metadata, created_at, updated_at)
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$18)
ON CONFLICT (agent_id) DO UPDATE SET
enabled = EXCLUDED.enabled,
interval_sec = EXCLUDED.interval_sec,
@@ -106,13 +106,14 @@ func (s *PGHeartbeatStore) Upsert(ctx context.Context, hb *store.AgentHeartbeat)
timezone = EXCLUDED.timezone,
channel = EXCLUDED.channel,
chat_id = EXCLUDED.chat_id,
next_run_at = EXCLUDED.next_run_at,
metadata = EXCLUDED.metadata,
updated_at = EXCLUDED.updated_at
RETURNING id, created_at, updated_at`,
hb.AgentID, hb.Enabled, hb.IntervalSec, hb.Prompt, hb.ProviderID, hb.Model,
hb.IsolatedSession, hb.LightContext, hb.AckMaxChars, hb.MaxRetries,
hb.ActiveHoursStart, hb.ActiveHoursEnd, hb.Timezone,
hb.Channel, hb.ChatID, meta, now,
hb.Channel, hb.ChatID, hb.NextRunAt, meta, now,
).Scan(&hb.ID, &hb.CreatedAt, &hb.UpdatedAt)
if err != nil {
return err
@@ -269,3 +270,60 @@ func (s *PGHeartbeatStore) ListLogs(ctx context.Context, agentID uuid.UUID, limi
}
return logs, total, nil
}
// ListDeliveryTargets returns distinct (channel, chatID, title, kind) pairs from sessions for an agent.
// Uses idx_sessions_agent (btree on agent_id) for fast lookup.
// Session key format: agent:{key}:{channel}:{kind}:{chatId}
func (s *PGHeartbeatStore) ListDeliveryTargets(ctx context.Context, agentID uuid.UUID) ([]store.DeliveryTarget, error) {
rows, err := s.db.QueryContext(ctx,
`SELECT DISTINCT ON (s.channel, chat_id)
s.channel,
split_part(s.session_key, ':', 5) AS chat_id,
COALESCE(
s.metadata->>'chat_title',
cc.display_name,
CASE WHEN cc.username != '' THEN '@' || cc.username ELSE '' END,
''
) AS title,
CASE
WHEN s.session_key LIKE '%:group:%' THEN 'group'
WHEN s.session_key LIKE '%:direct:%' THEN 'dm'
ELSE 'other'
END AS kind
FROM sessions s
LEFT JOIN channel_contacts cc
ON cc.sender_id = split_part(s.session_key, ':', 5)
AND cc.channel_type = s.channel
WHERE s.agent_id = $1
AND s.channel IS NOT NULL AND s.channel != ''
AND s.session_key NOT LIKE '%:heartbeat%'
AND s.session_key NOT LIKE '%:cron%'
AND s.session_key NOT LIKE '%:subagent%'
AND s.session_key NOT LIKE '%:team:%'
ORDER BY s.channel, chat_id, title`,
agentID,
)
if err != nil {
return nil, err
}
defer rows.Close()
var targets []store.DeliveryTarget
seen := make(map[string]bool)
for rows.Next() {
var t store.DeliveryTarget
if err := rows.Scan(&t.Channel, &t.ChatID, &t.Title, &t.Kind); err != nil {
return nil, err
}
// Deduplicate by channel+chatID (sessions can have multiple rows for same target).
key := t.Channel + ":" + t.ChatID
if t.ChatID != "" && !seen[key] {
seen[key] = true
targets = append(targets, t)
}
}
if err := rows.Err(); err != nil {
return nil, err
}
return targets, nil
}
+1
View File
@@ -92,6 +92,7 @@ const (
MethodHeartbeatLogs = "heartbeat.logs"
MethodHeartbeatChecklistGet = "heartbeat.checklist.get"
MethodHeartbeatChecklistSet = "heartbeat.checklist.set"
MethodHeartbeatTargets = "heartbeat.targets"
)
// Channel instances management