Merge pull request #26 from dj-oyu/codex/change-heartbeat-routing-to-multiple-targets

Add Telegram heartbeat thread routing, persistence and control command
This commit is contained in:
dj-oyu 2026-03-04 18:11:01 +09:00 committed by GitHub
commit 9befdc6a55
7 changed files with 489 additions and 43 deletions

View file

@ -113,7 +113,12 @@ func gatewayCmd(debug bool, orchestration bool, enableStats bool) error {
cfg.Heartbeat.Interval,
cfg.Heartbeat.Enabled,
)
heartbeatService.SetHeartbeatThreadID(cfg.Channels.Telegram.HeartbeatThreadID)
heartbeatService.SetBus(msgBus)
agentLoop.SetHeartbeatThreadUpdater(heartbeatService.SetHeartbeatThreadID)
agentLoop.SetConfigSaver(func(c *config.Config) error {
return config.SaveConfig(internal.GetConfigPath(), c)
})
heartbeatService.SetHandler(func(prompt, channel, chatID string) *tools.ToolResult {
// Use cli:direct as fallback if no valid channel
if channel == "" || chatID == "" {

View file

@ -102,6 +102,8 @@ type AgentLoop struct {
promptDirty atomic.Bool // true = rebuild needed on next GetSystemPrompt read
OnStateChange func() // called on plan/session/skills mutations
OnUserMessage func() // called when a real user message is processed
saveConfig func(*config.Config) error
onHeartbeatThreadUpdate func(int)
orchBroadcaster *orch.Broadcaster // nil when --orchestration not set
orchReporter orch.AgentReporter // always non-nil (Noop when disabled)
}
@ -213,6 +215,16 @@ func (al *AgentLoop) notifyStateChange() {
}
}
// SetConfigSaver registers a callback used by slash commands that persist runtime config changes.
func (al *AgentLoop) SetConfigSaver(fn func(*config.Config) error) {
al.saveConfig = fn
}
// SetHeartbeatThreadUpdater registers a callback to apply runtime heartbeat thread updates.
func (al *AgentLoop) SetHeartbeatThreadUpdater(fn func(int)) {
al.onHeartbeatThreadUpdate = fn
}
// registerSharedTools registers tools that are shared across all agents (web, message, spawn).
func registerSharedTools(
cfg *config.Config,
@ -610,6 +622,16 @@ func (al *AgentLoop) RecordLastChatID(chatID string) error {
return al.state.SetLastChatID(chatID)
}
// RecordLastHeartbeatTarget records the latest heartbeat-safe destination.
// This is intentionally separate from LastChannel so heartbeat routing can be
// reasoned about and evolved without breaking generic last-activity tracking.
func (al *AgentLoop) RecordLastHeartbeatTarget(target string) error {
if al.state == nil {
return nil
}
return al.state.SetLastHeartbeatTarget(target)
}
func (al *AgentLoop) ProcessDirect(ctx context.Context, content, sessionKey string) (string, error) {
return al.ProcessDirectWithChannel(ctx, content, sessionKey, "cli", "direct")
}
@ -1162,6 +1184,9 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt
if err := al.RecordLastChannel(channelKey); err != nil {
logger.WarnCF("agent", "Failed to record last channel", map[string]any{"error": err.Error()})
}
if err := al.RecordLastHeartbeatTarget(channelKey); err != nil {
logger.WarnCF("agent", "Failed to record last heartbeat target", map[string]any{"error": err.Error()})
}
}
}
@ -3543,11 +3568,94 @@ func (al *AgentLoop) handleCommand(ctx context.Context, msg bus.InboundMessage)
al.notifyStateChange()
}
return resp, handled
case "/heartbeat":
resp, handled := al.handleHeartbeatCommand(args, msg)
if handled {
al.notifyStateChange()
}
return resp, handled
}
return "", false
}
func (al *AgentLoop) handleHeartbeatCommand(args []string, msg bus.InboundMessage) (string, bool) {
if len(args) == 0 {
return "Usage: /heartbeat thread [here|off|<thread_id>]", true
}
if args[0] != "thread" {
return "Usage: /heartbeat thread [here|off|<thread_id>]", true
}
if len(args) < 2 {
return "Usage: /heartbeat thread [here|off|<thread_id>]", true
}
if msg.Channel != "telegram" {
return "/heartbeat thread is only supported from Telegram chats.", true
}
baseChatID, currentThreadID := splitChatAndThread(msg.ChatID)
if baseChatID == "" {
return "Unable to detect Telegram chat ID for heartbeat routing.", true
}
arg := strings.ToLower(strings.TrimSpace(args[1]))
var threadID int
var err error
switch arg {
case "off", "disable", "clear":
threadID = 0
case "here", "this":
if currentThreadID <= 0 {
return "Current Telegram message is not in a thread. Usage: /heartbeat thread <thread_id>", true
}
threadID = currentThreadID
default:
threadID, err = strconv.Atoi(arg)
if err != nil || threadID < 0 {
return "Usage: /heartbeat thread [here|off|<thread_id>]", true
}
}
al.cfg.Channels.Telegram.HeartbeatThreadID = threadID
if al.state != nil {
_ = al.state.SetHeartbeatTarget(fmt.Sprintf("telegram:%s", baseChatID))
}
if al.onHeartbeatThreadUpdate != nil {
al.onHeartbeatThreadUpdate(threadID)
}
if al.saveConfig != nil {
if err := al.saveConfig(al.cfg); err != nil {
return fmt.Sprintf("Failed to persist config.json: %v", err), true
}
}
if threadID == 0 {
return fmt.Sprintf("Heartbeat thread routing disabled for chat %s and saved to config.json.", baseChatID), true
}
return fmt.Sprintf("Heartbeat thread set to %d for chat %s and saved to config.json.", threadID, baseChatID), true
}
func splitChatAndThread(chatID string) (baseChatID string, threadID int) {
baseChatID = strings.TrimSpace(chatID)
if baseChatID == "" {
return "", 0
}
if slash := strings.Index(baseChatID, "/"); slash >= 0 {
threadPart := strings.TrimSpace(baseChatID[slash+1:])
baseChatID = strings.TrimSpace(baseChatID[:slash])
if tid, err := strconv.Atoi(threadPart); err == nil && tid > 0 {
threadID = tid
}
}
return baseChatID, threadID
}
// handleSessionCommand returns usage statistics or resets them.
func (al *AgentLoop) handleSessionCommand(args []string) string {
if al.stats == nil {

View file

@ -120,6 +120,38 @@ func TestRecordLastChatID(t *testing.T) {
}
}
func TestRecordLastHeartbeatTarget(t *testing.T) {
tmpDir, err := os.MkdirTemp("", "agent-test-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
defer os.RemoveAll(tmpDir)
cfg := &config.Config{
Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{
Workspace: tmpDir,
Model: "test-model",
MaxTokens: 4096,
MaxToolIterations: 10,
},
},
}
msgBus := bus.NewMessageBus()
provider := &mockProvider{}
al := NewAgentLoop(cfg, msgBus, provider)
target := "telegram:-100123/42"
if err := al.RecordLastHeartbeatTarget(target); err != nil {
t.Fatalf("RecordLastHeartbeatTarget failed: %v", err)
}
if got := al.state.GetLastHeartbeatTarget(); got != target {
t.Fatalf("GetLastHeartbeatTarget = %q, want %q", got, target)
}
}
func TestNewAgentLoop_StateInitialized(t *testing.T) {
// Create temp workspace
tmpDir, err := os.MkdirTemp("", "agent-test-*")
@ -1019,6 +1051,98 @@ func TestPlanCommand_ShowNoPlan(t *testing.T) {
}
}
func TestSplitChatAndThread(t *testing.T) {
tests := []struct {
name string
chatID string
wantChatID string
wantThread int
}{
{name: "plain chat", chatID: "-100123", wantChatID: "-100123", wantThread: 0},
{name: "chat with thread", chatID: "-100123/77", wantChatID: "-100123", wantThread: 77},
{name: "invalid thread", chatID: "-100123/abc", wantChatID: "-100123", wantThread: 0},
{name: "empty", chatID: "", wantChatID: "", wantThread: 0},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
gotChatID, gotThread := splitChatAndThread(tt.chatID)
if gotChatID != tt.wantChatID || gotThread != tt.wantThread {
t.Fatalf(
"splitChatAndThread(%q) = (%q, %d), want (%q, %d)",
tt.chatID,
gotChatID,
gotThread,
tt.wantChatID,
tt.wantThread,
)
}
})
}
}
func TestHeartbeatCommandThreadHerePersistsConfig(t *testing.T) {
al, cleanup := newTestAgentLoop(t)
defer cleanup()
var saved bool
var updatedThread int
al.SetConfigSaver(func(cfg *config.Config) error {
saved = true
if cfg.Channels.Telegram.HeartbeatThreadID != 42 {
t.Fatalf("HeartbeatThreadID in saver = %d, want 42", cfg.Channels.Telegram.HeartbeatThreadID)
}
return nil
})
al.SetHeartbeatThreadUpdater(func(threadID int) { updatedThread = threadID })
msg := bus.InboundMessage{
Content: "/heartbeat thread here",
Channel: "telegram",
ChatID: "-100500/42",
}
resp, handled := al.handleCommand(context.Background(), msg)
if !handled {
t.Fatal("expected /heartbeat command to be handled")
}
if !strings.Contains(resp, "Heartbeat thread set to 42") {
t.Fatalf("unexpected response: %q", resp)
}
if !saved {
t.Fatal("expected config saver to be called")
}
if updatedThread != 42 {
t.Fatalf("updatedThread = %d, want 42", updatedThread)
}
if got := al.cfg.Channels.Telegram.HeartbeatThreadID; got != 42 {
t.Fatalf("cfg heartbeat thread = %d, want 42", got)
}
if got := al.state.GetHeartbeatTarget(); got != "telegram:-100500" {
t.Fatalf("state heartbeat target = %q, want %q", got, "telegram:-100500")
}
}
func TestHeartbeatCommandThreadOff(t *testing.T) {
al, cleanup := newTestAgentLoop(t)
defer cleanup()
al.cfg.Channels.Telegram.HeartbeatThreadID = 99
resp, handled := al.handleCommand(context.Background(), bus.InboundMessage{
Content: "/heartbeat thread off",
Channel: "telegram",
ChatID: "-100500/42",
})
if !handled {
t.Fatal("expected /heartbeat command to be handled")
}
if !strings.Contains(resp, "disabled") {
t.Fatalf("unexpected response: %q", resp)
}
if got := al.cfg.Channels.Telegram.HeartbeatThreadID; got != 0 {
t.Fatalf("cfg heartbeat thread = %d, want 0", got)
}
}
func TestPlanCommand_StartNewPlan(t *testing.T) {
al, cleanup := newTestAgentLoop(t)
defer cleanup()

View file

@ -44,6 +44,7 @@ type HeartbeatService struct {
mu sync.RWMutex
stopChan chan struct{}
lastNotifiedAt time.Time // when a non-silent result was last sent to user
heartbeatThreadID int
}
// NewHeartbeatService creates a new heartbeat service
@ -79,6 +80,13 @@ func (hs *HeartbeatService) SetHandler(handler HeartbeatHandler) {
hs.handler = handler
}
// SetHeartbeatThreadID configures Telegram thread routing for heartbeat messages.
func (hs *HeartbeatService) SetHeartbeatThreadID(threadID int) {
hs.mu.Lock()
defer hs.mu.Unlock()
hs.heartbeatThreadID = threadID
}
// ResetSuppression clears the notification suppression so the next
// non-silent heartbeat result will be delivered to the user again.
// Typically called when a user message arrives.
@ -182,12 +190,8 @@ func (hs *HeartbeatService) executeHeartbeat() {
return
}
// Get last channel info for context
lastChannel := hs.state.GetLastChannel()
channel, chatID := hs.parseLastChannel(lastChannel)
// Debug log for channel resolution
hs.logInfof("Resolved channel: %s, chatID: %s (from lastChannel: %s)", channel, chatID, lastChannel)
channel, chatID, reason := hs.resolveHeartbeatTarget()
hs.logInfof("Resolved channel: %s, chatID: %s (%s)", channel, chatID, reason)
result := handler(prompt, channel, chatID)
@ -236,6 +240,67 @@ func (hs *HeartbeatService) executeHeartbeat() {
hs.logInfof("Heartbeat completed: %s", result.ForLLM)
}
func (hs *HeartbeatService) resolveHeartbeatTarget() (channel, chatID, reason string) {
if explicit := hs.state.GetHeartbeatTarget(); explicit != "" {
if ch, cid := hs.parseTarget(explicit); ch != "" && cid != "" {
return ch, cid, fmt.Sprintf("explicit heartbeat target: %s", explicit)
}
hs.logErrorf("Invalid explicit heartbeat target: %s", explicit)
}
if threadID := hs.telegramHeartbeatThreadID(); threadID > 0 {
if ch, cid, src := hs.resolveTelegramThreadTarget(threadID); ch != "" && cid != "" {
return ch, cid, src
}
}
lastChannel := hs.state.GetLastChannel()
channel, chatID = hs.parseLastChannel(lastChannel)
return channel, chatID, fmt.Sprintf("fallback last channel: %s", lastChannel)
}
func (hs *HeartbeatService) resolveTelegramThreadTarget(threadID int) (channel, chatID, reason string) {
candidates := []struct {
value string
reason string
}{
{value: hs.state.GetLastHeartbeatTarget(), reason: "last heartbeat target"},
{value: hs.state.GetLastChannel(), reason: "last channel"},
}
for _, candidate := range candidates {
ch, cid := hs.parseTarget(candidate.value)
if ch != "telegram" || cid == "" {
continue
}
return ch,
withTelegramThread(cid, threadID),
fmt.Sprintf("telegram heartbeat_thread_id from %s", candidate.reason)
}
return "", "", ""
}
func (hs *HeartbeatService) telegramHeartbeatThreadID() int {
hs.mu.RLock()
defer hs.mu.RUnlock()
return hs.heartbeatThreadID
}
func withTelegramThread(chatID string, threadID int) string {
if threadID <= 0 || chatID == "" {
return chatID
}
baseChatID := chatID
if slash := strings.Index(baseChatID, "/"); slash >= 0 {
baseChatID = baseChatID[:slash]
}
if baseChatID == "" {
return chatID
}
return fmt.Sprintf("%s/%d", baseChatID, threadID)
}
// buildPrompt builds the heartbeat prompt from HEARTBEAT.md
func (hs *HeartbeatService) buildPrompt() string {
heartbeatPath := filepath.Join(hs.workspace, "HEARTBEAT.md")
@ -306,20 +371,21 @@ Add your heartbeat tasks below this line:
// parseLastChannel parses the last channel string into platform and userID.
// Returns empty strings for invalid or internal channels.
func (hs *HeartbeatService) parseLastChannel(lastChannel string) (platform, userID string) {
if lastChannel == "" {
return hs.parseTarget(lastChannel)
}
func (hs *HeartbeatService) parseTarget(target string) (platform, userID string) {
if target == "" {
return "", ""
}
// Parse channel format: "platform:user_id" (e.g., "telegram:123456")
parts := strings.SplitN(lastChannel, ":", 2)
parts := strings.SplitN(target, ":", 2)
if len(parts) != 2 || parts[0] == "" || parts[1] == "" {
hs.logErrorf("Invalid last channel format: %s", lastChannel)
hs.logErrorf("Invalid heartbeat target format: %s", target)
return "", ""
}
platform, userID = parts[0], parts[1]
// Skip internal channels
if constants.IsInternalChannel(platform) {
hs.logInfof("Skipping internal channel: %s", platform)
return "", ""

View file

@ -255,3 +255,62 @@ func TestHeartbeatFilePath(t *testing.T) {
t.Errorf("Expected HEARTBEAT.md at %s, but it doesn't exist", expectedPath)
}
}
func TestExecuteHeartbeat_TargetPriority_ExplicitTarget(t *testing.T) {
tmpDir, err := os.MkdirTemp("", "heartbeat-test-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
defer os.RemoveAll(tmpDir)
hs := NewHeartbeatService(tmpDir, 30, true)
hs.stopChan = make(chan struct{})
hs.SetHeartbeatThreadID(77)
if err := hs.state.SetHeartbeatTarget("slack:C12345/999"); err != nil {
t.Fatalf("SetHeartbeatTarget failed: %v", err)
}
if err := hs.state.SetLastHeartbeatTarget("telegram:-100500"); err != nil {
t.Fatalf("SetLastHeartbeatTarget failed: %v", err)
}
var gotChannel, gotChatID string
hs.SetHandler(func(prompt, channel, chatID string) *tools.ToolResult {
gotChannel, gotChatID = channel, chatID
return tools.SilentResult("ok")
})
os.WriteFile(filepath.Join(tmpDir, "HEARTBEAT.md"), []byte("Test task"), 0o644)
hs.executeHeartbeat()
if gotChannel != "slack" || gotChatID != "C12345/999" {
t.Fatalf("handler target = %s:%s, want slack:C12345/999", gotChannel, gotChatID)
}
}
func TestExecuteHeartbeat_TargetPriority_TelegramThread(t *testing.T) {
tmpDir, err := os.MkdirTemp("", "heartbeat-test-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
defer os.RemoveAll(tmpDir)
hs := NewHeartbeatService(tmpDir, 30, true)
hs.stopChan = make(chan struct{})
hs.SetHeartbeatThreadID(77)
if err := hs.state.SetLastHeartbeatTarget("telegram:-100500"); err != nil {
t.Fatalf("SetLastHeartbeatTarget failed: %v", err)
}
var gotChannel, gotChatID string
hs.SetHandler(func(prompt, channel, chatID string) *tools.ToolResult {
gotChannel, gotChatID = channel, chatID
return tools.SilentResult("ok")
})
os.WriteFile(filepath.Join(tmpDir, "HEARTBEAT.md"), []byte("Test task"), 0o644)
hs.executeHeartbeat()
if gotChannel != "telegram" || gotChatID != "-100500/77" {
t.Fatalf("handler target = %s:%s, want telegram:-100500/77", gotChannel, gotChatID)
}
}

View file

@ -18,6 +18,14 @@ type State struct {
// LastChannel is the last channel used for communication
LastChannel string `json:"last_channel,omitempty"`
// LastHeartbeatTarget is the last destination considered safe for heartbeat delivery.
// Format: "channel:chatID[/thread]".
LastHeartbeatTarget string `json:"last_heartbeat_target,omitempty"`
// HeartbeatTarget is an explicit heartbeat destination override.
// Format: "channel:chatID[/thread]".
HeartbeatTarget string `json:"heartbeat_target,omitempty"`
// LastChatID is the last chat ID used for communication
LastChatID string `json:"last_chat_id,omitempty"`
@ -85,6 +93,36 @@ func (sm *Manager) SetLastChannel(channel string) error {
return nil
}
// SetLastHeartbeatTarget atomically updates the last heartbeat target and saves the state.
func (sm *Manager) SetLastHeartbeatTarget(target string) error {
sm.mu.Lock()
defer sm.mu.Unlock()
sm.state.LastHeartbeatTarget = target
sm.state.Timestamp = time.Now()
if err := sm.saveAtomic(); err != nil {
return fmt.Errorf("failed to save state atomically: %w", err)
}
return nil
}
// SetHeartbeatTarget atomically updates the explicit heartbeat target and saves the state.
func (sm *Manager) SetHeartbeatTarget(target string) error {
sm.mu.Lock()
defer sm.mu.Unlock()
sm.state.HeartbeatTarget = target
sm.state.Timestamp = time.Now()
if err := sm.saveAtomic(); err != nil {
return fmt.Errorf("failed to save state atomically: %w", err)
}
return nil
}
// SetLastChatID atomically updates the last chat ID and saves the state.
func (sm *Manager) SetLastChatID(chatID string) error {
sm.mu.Lock()
@ -109,6 +147,20 @@ func (sm *Manager) GetLastChannel() string {
return sm.state.LastChannel
}
// GetLastHeartbeatTarget returns the last heartbeat target from the state.
func (sm *Manager) GetLastHeartbeatTarget() string {
sm.mu.RLock()
defer sm.mu.RUnlock()
return sm.state.LastHeartbeatTarget
}
// GetHeartbeatTarget returns the explicit heartbeat target from the state.
func (sm *Manager) GetHeartbeatTarget() string {
sm.mu.RLock()
defer sm.mu.RUnlock()
return sm.state.HeartbeatTarget
}
// GetLastChatID returns the last chat ID from the state.
func (sm *Manager) GetLastChatID() string {
sm.mu.RLock()

View file

@ -214,3 +214,35 @@ func TestNewManager_EmptyWorkspace(t *testing.T) {
t.Error("Expected zero timestamp for new state")
}
}
func TestHeartbeatTargetsPersistence(t *testing.T) {
tmpDir, err := os.MkdirTemp("", "state-test-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
defer os.RemoveAll(tmpDir)
sm := NewManager(tmpDir)
if err := sm.SetLastHeartbeatTarget("telegram:-100123"); err != nil {
t.Fatalf("SetLastHeartbeatTarget failed: %v", err)
}
if err := sm.SetHeartbeatTarget("telegram:-100123/42"); err != nil {
t.Fatalf("SetHeartbeatTarget failed: %v", err)
}
if got := sm.GetLastHeartbeatTarget(); got != "telegram:-100123" {
t.Fatalf("GetLastHeartbeatTarget = %q, want %q", got, "telegram:-100123")
}
if got := sm.GetHeartbeatTarget(); got != "telegram:-100123/42" {
t.Fatalf("GetHeartbeatTarget = %q, want %q", got, "telegram:-100123/42")
}
sm2 := NewManager(tmpDir)
if got := sm2.GetLastHeartbeatTarget(); got != "telegram:-100123" {
t.Fatalf("persistent GetLastHeartbeatTarget = %q, want %q", got, "telegram:-100123")
}
if got := sm2.GetHeartbeatTarget(); got != "telegram:-100123/42" {
t.Fatalf("persistent GetHeartbeatTarget = %q, want %q", got, "telegram:-100123/42")
}
}