package methods import ( "context" "database/sql" "encoding/json" "errors" "fmt" "log/slog" "time" "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/agent" "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/gateway" "github.com/nextlevelbuilder/goclaw/internal/i18n" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) // HeartbeatMethods handles heartbeat.get/set/toggle/test/logs/checklist RPC methods. type HeartbeatMethods struct { hbStore store.HeartbeatStore agentStore store.AgentStore agentRouter *agent.Router // cache-aware lookup for resolveAgentUUIDCached hot path providerStore store.ProviderStore eventBus bus.EventPublisher wakeFn func(uuid.UUID) // triggers immediate heartbeat run } func NewHeartbeatMethods(hb store.HeartbeatStore, eventBus bus.EventPublisher) *HeartbeatMethods { return &HeartbeatMethods{hbStore: hb, eventBus: eventBus} } // SetAgentStore sets the agent store for HEARTBEAT.md read/write via RPC. func (m *HeartbeatMethods) SetAgentStore(as store.AgentStore) { m.agentStore = as } // SetAgentRouter wires the agent router for cache-aware agent_key resolution. // Optional — when nil, resolveAgentUUIDCached falls back to a pure DB lookup. func (m *HeartbeatMethods) SetAgentRouter(r *agent.Router) { m.agentRouter = r } // SetProviderStore sets the provider store for resolving provider names to UUIDs. func (m *HeartbeatMethods) SetProviderStore(ps store.ProviderStore) { m.providerStore = ps } // SetWakeFn sets the function called when "heartbeat.test" triggers an immediate run. func (m *HeartbeatMethods) SetWakeFn(fn func(uuid.UUID)) { m.wakeFn = fn } func (m *HeartbeatMethods) Register(router *gateway.MethodRouter) { router.Register(protocol.MethodHeartbeatGet, m.handleGet) router.Register(protocol.MethodHeartbeatSet, m.handleSet) router.Register(protocol.MethodHeartbeatToggle, m.handleToggle) router.Register(protocol.MethodHeartbeatTest, m.handleTest) 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) { locale := store.LocaleFromContext(ctx) var params struct { AgentID string `json:"agentId"` } if req.Params != nil { json.Unmarshal(req.Params, ¶ms) } if params.AgentID == "" { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, i18n.T(locale, i18n.MsgRequired, "agentId"))) return } agentUUID, err := resolveAgentUUIDCached(ctx, m.agentRouter, m.agentStore, params.AgentID) if err != nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, "invalid agentId")) return } hb, err := m.hbStore.Get(ctx, agentUUID) if err != nil { if errors.Is(err, sql.ErrNoRows) { client.SendResponse(protocol.NewOKResponse(req.ID, map[string]any{"heartbeat": nil})) return } client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("get", err))) return } client.SendResponse(protocol.NewOKResponse(req.ID, map[string]any{"heartbeat": hb})) } func (m *HeartbeatMethods) handleSet(ctx context.Context, client *gateway.Client, req *protocol.RequestFrame) { locale := store.LocaleFromContext(ctx) var params struct { AgentID string `json:"agentId"` Enabled *bool `json:"enabled"` IntervalSec *int `json:"intervalSec"` Prompt *string `json:"prompt"` ProviderName *string `json:"providerName"` Model *string `json:"model"` IsolatedSession *bool `json:"isolatedSession"` LightContext *bool `json:"lightContext"` AckMaxChars *int `json:"ackMaxChars"` MaxRetries *int `json:"maxRetries"` ActiveHoursStart *string `json:"activeHoursStart"` ActiveHoursEnd *string `json:"activeHoursEnd"` Timezone *string `json:"timezone"` Channel *string `json:"channel"` ChatID *string `json:"chatId"` } if req.Params != nil { json.Unmarshal(req.Params, ¶ms) } if params.AgentID == "" { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, i18n.T(locale, i18n.MsgRequired, "agentId"))) return } agentUUID, err := resolveAgentUUIDCached(ctx, m.agentRouter, m.agentStore, params.AgentID) if err != nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, "invalid agentId")) return } // 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, heartbeatInternalErr("set.load", err))) return } if hb == nil { hb = &store.AgentHeartbeat{ AgentID: agentUUID, IntervalSec: 1800, IsolatedSession: true, AckMaxChars: 300, MaxRetries: 2, } } if params.Enabled != nil { hb.Enabled = *params.Enabled } if params.IntervalSec != nil { if *params.IntervalSec < 300 { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, "minimum interval is 300 seconds")) return } hb.IntervalSec = *params.IntervalSec } if params.Prompt != nil { hb.Prompt = params.Prompt } if params.ProviderName != nil { if *params.ProviderName == "" { hb.ProviderID = nil // clear override } else if m.providerStore != nil { prov, err := m.providerStore.GetProviderByName(ctx, *params.ProviderName) if err != nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, "provider not found: "+*params.ProviderName)) return } hb.ProviderID = &prov.ID } } if params.Model != nil { if *params.Model == "" { hb.Model = nil // clear override } else { hb.Model = params.Model } } if params.IsolatedSession != nil { hb.IsolatedSession = *params.IsolatedSession } if params.LightContext != nil { hb.LightContext = *params.LightContext } if params.AckMaxChars != nil { if *params.AckMaxChars < 0 { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, "ackMaxChars must be >= 0")) return } hb.AckMaxChars = *params.AckMaxChars } if params.MaxRetries != nil { if *params.MaxRetries < 0 || *params.MaxRetries > 10 { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, "maxRetries must be 0-10")) return } hb.MaxRetries = *params.MaxRetries } if params.ActiveHoursStart != nil { hb.ActiveHoursStart = params.ActiveHoursStart } if params.ActiveHoursEnd != nil { hb.ActiveHoursEnd = params.ActiveHoursEnd } if params.Timezone != nil { if *params.Timezone != "" { if _, err := time.LoadLocation(*params.Timezone); err != nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, fmt.Sprintf("invalid timezone: %s", *params.Timezone))) return } } hb.Timezone = params.Timezone } if params.Channel != nil { hb.Channel = params.Channel } if params.ChatID != nil { hb.ChatID = params.ChatID } if hb.Enabled && hb.NextRunAt == nil { nextRun := time.Now().Add(time.Duration(hb.IntervalSec)*time.Second + store.StaggerOffset(hb.AgentID, hb.IntervalSec)) hb.NextRunAt = &nextRun } if err := m.hbStore.Upsert(ctx, hb); err != nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("set.upsert", err))) return } client.SendResponse(protocol.NewOKResponse(req.ID, map[string]any{"heartbeat": hb})) m.emitCacheInvalidate(hb.AgentID.String()) emitAudit(m.eventBus, client, "heartbeat.set", "heartbeat", hb.AgentID.String()) } func (m *HeartbeatMethods) handleToggle(ctx context.Context, client *gateway.Client, req *protocol.RequestFrame) { locale := store.LocaleFromContext(ctx) var params struct { AgentID string `json:"agentId"` Enabled bool `json:"enabled"` } if req.Params != nil { json.Unmarshal(req.Params, ¶ms) } if params.AgentID == "" { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, i18n.T(locale, i18n.MsgRequired, "agentId"))) return } agentUUID, err := resolveAgentUUIDCached(ctx, m.agentRouter, m.agentStore, params.AgentID) if err != nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, "invalid agentId")) return } hb, err := m.hbStore.Get(ctx, agentUUID) if err != nil { if errors.Is(err, sql.ErrNoRows) { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrNotFound, "heartbeat not configured")) return } client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("op", err))) return } hb.Enabled = params.Enabled if params.Enabled && hb.NextRunAt == nil { nextRun := time.Now().Add(time.Duration(hb.IntervalSec) * time.Second) hb.NextRunAt = &nextRun } if err := m.hbStore.Upsert(ctx, hb); err != nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("op", err))) return } client.SendResponse(protocol.NewOKResponse(req.ID, map[string]any{ "agentId": params.AgentID, "enabled": params.Enabled, })) m.emitCacheInvalidate(params.AgentID) emitAudit(m.eventBus, client, "heartbeat.toggled", "heartbeat", params.AgentID) } func (m *HeartbeatMethods) handleTest(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, ¶ms) } if params.AgentID == "" { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, i18n.T(locale, i18n.MsgRequired, "agentId"))) return } agentUUID, err := resolveAgentUUIDCached(ctx, m.agentRouter, m.agentStore, params.AgentID) if err != nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, "invalid agentId")) return } if m.wakeFn == nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, "heartbeat ticker not available")) return } client.SendResponse(protocol.NewOKResponse(req.ID, map[string]any{"ok": true})) emitAudit(m.eventBus, client, "heartbeat.test", "heartbeat", params.AgentID) go func() { m.wakeFn(agentUUID) slog.Info("heartbeat.test triggered", "agent_id", params.AgentID) }() } func (m *HeartbeatMethods) handleLogs(ctx context.Context, client *gateway.Client, req *protocol.RequestFrame) { locale := store.LocaleFromContext(ctx) var params struct { AgentID string `json:"agentId"` Limit int `json:"limit"` Offset int `json:"offset"` } if req.Params != nil { json.Unmarshal(req.Params, ¶ms) } if params.AgentID == "" { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, i18n.T(locale, i18n.MsgRequired, "agentId"))) return } agentUUID, err := resolveAgentUUIDCached(ctx, m.agentRouter, m.agentStore, params.AgentID) if err != nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, "invalid agentId")) return } logs, total, err := m.hbStore.ListLogs(ctx, agentUUID, params.Limit, params.Offset) if err != nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("op", err))) return } client.SendResponse(protocol.NewOKResponse(req.ID, map[string]any{ "logs": logs, "total": total, })) } func (m *HeartbeatMethods) handleChecklistGet(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, ¶ms) } if params.AgentID == "" { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, i18n.T(locale, i18n.MsgRequired, "agentId"))) return } if m.agentStore == nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, "agent store not configured")) return } agentUUID, err := resolveAgentUUIDCached(ctx, m.agentRouter, m.agentStore, params.AgentID) if err != nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, "invalid agentId")) return } files, err := m.agentStore.GetAgentContextFiles(ctx, agentUUID) if err != nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("op", err))) return } var content string for _, f := range files { if f.FileName == "HEARTBEAT.md" { content = f.Content break } } client.SendResponse(protocol.NewOKResponse(req.ID, map[string]any{ "content": content, })) } func (m *HeartbeatMethods) handleChecklistSet(ctx context.Context, client *gateway.Client, req *protocol.RequestFrame) { locale := store.LocaleFromContext(ctx) var params struct { AgentID string `json:"agentId"` Content string `json:"content"` } if req.Params != nil { json.Unmarshal(req.Params, ¶ms) } if params.AgentID == "" { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, i18n.T(locale, i18n.MsgRequired, "agentId"))) return } if m.agentStore == nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, "agent store not configured")) return } agentUUID, err := resolveAgentUUIDCached(ctx, m.agentRouter, m.agentStore, params.AgentID) if err != nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, "invalid agentId")) return } if err := m.agentStore.SetAgentContextFile(ctx, agentUUID, "HEARTBEAT.md", params.Content); err != nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, heartbeatInternalErr("op", err))) return } client.SendResponse(protocol.NewOKResponse(req.ID, map[string]any{ "ok": true, "length": len([]rune(params.Content)), })) 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, ¶ms) } if params.AgentID == "" { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, i18n.T(locale, i18n.MsgRequired, "agentId"))) return } // agentId param kept for backward compat but not used — targets are tenant-scoped. tenantID := store.TenantIDFromContext(ctx) if tenantID == uuid.Nil { tenantID = store.MasterTenantID } targets, err := m.hbStore.ListDeliveryTargets(ctx, tenantID) 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, Payload: bus.CacheInvalidatePayload{ Kind: bus.CacheKindHeartbeat, Key: agentID, }, }) }