picoclaw/pkg/heartbeat/service.go

package heartbeat

import (
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/sipeed/picoclaw/pkg/logger"
)

// ToolResult represents a structured result from tool execution.
// This is a minimal local definition to avoid circular dependencies.
type ToolResult struct {
	ForLLM  string `json:"for_llm"`
	ForUser string `json:"for_user,omitempty"`
	Silent  bool   `json:"silent"`
	IsError bool   `json:"is_error"`
	Async   bool   `json:"async"`
	Err     error  `json:"-"`
}

// HeartbeatHandler is the function type for handling heartbeat with tool support.
// It returns a ToolResult that can indicate async operations.
type HeartbeatHandler func(prompt string) *ToolResult

type HeartbeatService struct {
	workspace   string
	onHeartbeat func(string) (string, error)
	// onHeartbeatWithTools is the new handler that supports ToolResult returns
	onHeartbeatWithTools HeartbeatHandler
	interval             time.Duration
	enabled              bool
	mu                   sync.RWMutex
	stopChan             chan struct{}
}

func NewHeartbeatService(workspace string, onHeartbeat func(string) (string, error), intervalS int, enabled bool) *HeartbeatService {
	return &HeartbeatService{
		workspace:   workspace,
		onHeartbeat: onHeartbeat,
		interval:    time.Duration(intervalS) * time.Second,
		enabled:     enabled,
		stopChan:    make(chan struct{}),
	}
}

// SetOnHeartbeatWithTools sets the tool-supporting heartbeat handler.
// This handler returns a ToolResult that can indicate async operations.
// When set, this handler takes precedence over the legacy onHeartbeat callback.
func (hs *HeartbeatService) SetOnHeartbeatWithTools(handler HeartbeatHandler) {
	hs.mu.Lock()
	defer hs.mu.Unlock()
	hs.onHeartbeatWithTools = handler
}

func (hs *HeartbeatService) Start() error {
	hs.mu.Lock()
	defer hs.mu.Unlock()

	if hs.running() {
		return nil
	}

	if !hs.enabled {
		return fmt.Errorf("heartbeat service is disabled")
	}

	go hs.runLoop()

	return nil
}

func (hs *HeartbeatService) Stop() {
	hs.mu.Lock()
	defer hs.mu.Unlock()

	if !hs.running() {
		return
	}

	close(hs.stopChan)
}

func (hs *HeartbeatService) running() bool {
	select {
	case <-hs.stopChan:
		return false
	default:
		return true
	}
}

func (hs *HeartbeatService) runLoop() {
	ticker := time.NewTicker(hs.interval)
	defer ticker.Stop()

	for {
		select {
		case <-hs.stopChan:
			return
		case <-ticker.C:
			hs.checkHeartbeat()
		}
	}
}

func (hs *HeartbeatService) checkHeartbeat() {
	hs.mu.RLock()
	if !hs.enabled || !hs.running() {
		hs.mu.RUnlock()
		return
	}
	hs.mu.RUnlock()

	prompt := hs.buildPrompt()

	// Prefer the new tool-supporting handler
	if hs.onHeartbeatWithTools != nil {
		hs.executeHeartbeatWithTools(prompt)
	} else if hs.onHeartbeat != nil {
		_, err := hs.onHeartbeat(prompt)
		if err != nil {
			hs.log(fmt.Sprintf("Heartbeat error: %v", err))
		}
	}
}

// ExecuteHeartbeatWithTools executes a heartbeat using the tool-supporting handler.
// This method processes ToolResult returns and handles async tasks appropriately.
// If the result is async, it logs that the task started in background.
// If the result is an error, it logs the error message.
// This method is designed to be called from checkHeartbeat or directly by external code.
func (hs *HeartbeatService) ExecuteHeartbeatWithTools(prompt string) {
	hs.executeHeartbeatWithTools(prompt)
}

// executeHeartbeatWithTools is the internal implementation of tool-supporting heartbeat.
func (hs *HeartbeatService) executeHeartbeatWithTools(prompt string) {
	result := hs.onHeartbeatWithTools(prompt)

	if result == nil {
		hs.log("Heartbeat handler returned nil result")
		return
	}

	// Handle different result types
	if result.IsError {
		hs.log(fmt.Sprintf("Heartbeat error: %s", result.ForLLM))
		return
	}

	if result.Async {
		// Async task started - log and return immediately
		hs.log(fmt.Sprintf("Async task started: %s", result.ForLLM))
		logger.InfoCF("heartbeat", "Async heartbeat task started",
			map[string]interface{}{
				"message": result.ForLLM,
			})
		return
	}

	// Normal completion - log result
	hs.log(fmt.Sprintf("Heartbeat completed: %s", result.ForLLM))
}

func (hs *HeartbeatService) buildPrompt() string {
	notesDir := filepath.Join(hs.workspace, "memory")
	notesFile := filepath.Join(notesDir, "HEARTBEAT.md")

	var notes string
	if data, err := os.ReadFile(notesFile); err == nil {
		notes = string(data)
	}

	now := time.Now().Format("2006-01-02 15:04")

	prompt := fmt.Sprintf(`# Heartbeat Check

Current time: %s

Check if there are any tasks I should be aware of or actions I should take.
Review the memory file for any important updates or changes.
Be proactive in identifying potential issues or improvements.

%s
`, now, notes)

	return prompt
}

func (hs *HeartbeatService) log(message string) {
	logFile := filepath.Join(hs.workspace, "memory", "heartbeat.log")
	f, err := os.OpenFile(logFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return
	}
	defer f.Close()

	timestamp := time.Now().Format("2006-01-02 15:04:05")
	fmt.Fprintf(f, "[%s] %s\n", timestamp, message)
}