pi-dashboard/state.go

package main

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

// AgentState matches the TypeScript interface from dashboard.ts
type AgentState struct {
	Session          string  `json:"session"`
	Machine          string  `json:"machine"`
	Cwd              string  `json:"cwd"`
	AgentType        string  `json:"agent_type"`
	Model            string  `json:"model"`
	AgentStatus      string  `json:"agent_status"`
	Status           string  `json:"status"` // legacy alias of agent_status
	SessionState     string  `json:"session_state"`
	SessionStartedAt *string `json:"session_started_at"`
	ClosedAt         *string `json:"closed_at"`
	LastSeenAt       *string `json:"last_seen_at"`
	StartedAt        *string `json:"started_at"`
	CompletedAt      *string `json:"completed_at"`
	DurationSeconds  int     `json:"duration_seconds"`
	ToolCalls        int     `json:"tool_calls"`
	CurrentTask      string  `json:"current_task"`
	LastTool         *string `json:"last_tool"`
	LastToolAt       *string `json:"last_tool_at"`
	TerminalType     string  `json:"terminal_type"`
	ZellijSession    string  `json:"zellij_session"`
	TmuxSession      string  `json:"tmux_session"`
	PID              int     `json:"pid"`
	IsSubagent       bool    `json:"is_subagent"`
	ConnectionType   string  `json:"connection_type"`
	ConnectionName   string  `json:"connection_name"`
	Blocked          bool    `json:"blocked"`
	BlockedPrompt    *string `json:"blocked_prompt"`
	CostEstimate     float64 `json:"cost_estimate"`
}

// agentStatus returns the agent-turn status, falling back to the legacy
// "status" field for state files written by older extension versions.
func (a AgentState) agentStatus() string {
	if a.AgentStatus != "" {
		return a.AgentStatus
	}
	return a.Status
}

// legacyStale is the staleness window applied to legacy state files (written
// by the pre-heartbeat extension). Those files are only rewritten on events,
// so an idle-but-alive session can go minutes without a write — be generous.
const legacyStale = 30 * time.Minute

// isLegacy reports whether the state file predates the heartbeat extension
// (new files always carry pid + session_state; legacy ones have neither).
func (a AgentState) isLegacy() bool {
	return a.PID == 0 && a.SessionState == ""
}

// IsClosed reports whether the session is no longer alive:
//   - the extension marked it closed (session_state == "closed"), or
//   - the heartbeat (last_seen_at) is older than staleAfter (killed
//     without session_shutdown, e.g. closed terminal, kill -9), or
//   - there is no heartbeat at all (pre-heartbeat state files)
func (a AgentState) IsClosed(staleAfter time.Duration) bool {
	if a.SessionState == "closed" {
		return true
	}
	ls := a.LastSeenAt
	if ls == nil || *ls == "" {
		return true
	}
	t, err := time.Parse(time.RFC3339, *ls)
	if err != nil {
		return true
	}
	window := staleAfter
	if a.isLegacy() {
		window = legacyStale
	}
	return time.Since(t) > window
}

// WasKilled is true when the session died without a graceful shutdown
// (closed_at never written by the extension).
func (a AgentState) WasKilled(staleAfter time.Duration) bool {
	if a.SessionState == "closed" || a.ClosedAt != nil {
		return false
	}
	return a.IsClosed(staleAfter)
}

func (a AgentState) StatusStr(staleAfter time.Duration) string {
	if a.IsClosed(staleAfter) {
		if a.WasKilled(staleAfter) {
			return "✗ lost"
		}
		return "✗ closed"
	}
	switch a.agentStatus() {
	case "running":
		return "● running"
	case "blocked":
		return "⚠ blocked"
	case "error":
		return "✕ error"
	default:
		return "○ idle"
	}
}

// Folder returns the basename of the working directory the session ran in.
func (a AgentState) Folder() string {
	if a.Cwd == "" {
		return a.Session
	}
	trimmed := strings.TrimRight(a.Cwd, "/")
	base := filepath.Base(trimmed)
	if base == "/" || base == "." || base == "" {
		return a.Cwd
	}
	return base
}

func (a AgentState) ShortModel() string {
	parts := strings.Split(a.Model, "/")
	if len(parts) == 2 {
		s := parts[1]
		if len(s) > 14 {
			s = s[:12] + ".."
		}
		return s
	}
	if len(a.Model) > 14 {
		return a.Model[:12] + ".."
	}
	return a.Model
}

func (a AgentState) MachineShort() string {
	s := a.Machine
	if len(s) > 12 {
		s = s[:10] + ".."
	}
	return s
}

func (a AgentState) DurationStr() string {
	d := a.DurationSeconds
	if d == 0 {
		return ""
	}
	if d < 60 {
		return fmt.Sprintf("%ds", d)
	}
	if d < 3600 {
		return fmt.Sprintf("%dm%ds", d/60, d%60)
	}
	return fmt.Sprintf("%dh%dm", d/3600, (d%3600)/60)
}

// fmtTime renders an RFC3339 timestamp compactly: HH:MM for today,
// MM-DD HH:MM otherwise.
func fmtTime(s *string) string {
	if s == nil || *s == "" {
		return ""
	}
	t, err := time.Parse(time.RFC3339, *s)
	if err != nil {
		return ""
	}
	now := time.Now()
	if t.Year() == now.Year() && t.YearDay() == now.YearDay() {
		return t.Format("15:04")
	}
	return t.Format("01-02 15:04")
}

// lastActivity returns the most recent timestamp we know about for the session.
func (a AgentState) lastActivity() time.Time {
	for _, s := range []*string{a.LastSeenAt, a.CompletedAt, a.StartedAt, a.SessionStartedAt} {
		if s != nil {
			if t, err := time.Parse(time.RFC3339, *s); err == nil {
				return t
			}
		}
	}
	return time.Time{}
}

// sessionStart returns when the pi session began.
func (a AgentState) sessionStart() time.Time {
	if a.SessionStartedAt != nil {
		if t, err := time.Parse(time.RFC3339, *a.SessionStartedAt); err == nil {
			return t
		}
	}
	return a.lastActivity()
}

// IsLocal returns true if this agent is on the same machine
func (a AgentState) IsLocal() bool {
	return a.Machine == localHostname
}

var localHostname string

// readLocalStates reads all JSON files from the given directory
func readLocalStates(dir string) ([]AgentState, error) {
	entries, err := os.ReadDir(dir)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, fmt.Errorf("reading dir %s: %w", dir, err)
	}

	var states []AgentState
	for _, entry := range entries {
		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") {
			continue
		}
		s, err := readStateFile(filepath.Join(dir, entry.Name()))
		if err == nil {
			states = append(states, s)
		}
	}
	return states, nil
}

func readStateFile(path string) (AgentState, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return AgentState{}, err
	}
	var s AgentState
	if err := json.Unmarshal(data, &s); err != nil {
		return AgentState{}, err
	}

	// Legacy state files (pre-heartbeat extension) have no last_seen_at.
	// The old extension rewrote the file on every event, so the file mtime
	// is a good liveness proxy for them.
	if s.LastSeenAt == nil || *s.LastSeenAt == "" {
		if fi, err := os.Stat(path); err == nil {
			t := fi.ModTime().UTC().Format(time.RFC3339)
			s.LastSeenAt = &t
		}
	}
	return s, nil
}

// readRemoteStates SSHes into a machine and reads its state files.
// State files are pretty-printed multi-line JSON, so `cat *.json`
// concatenates several objects — decode them as a stream.
func readRemoteStates(cfg RemoteConfig) ([]AgentState, error) {
	hostPort := cfg.Host
	if cfg.Port > 0 && cfg.Port != 22 {
		hostPort = fmt.Sprintf("%s:%d", cfg.Host, cfg.Port)
	}

	user := cfg.User
	if user == "" {
		user = "sam"
	}

	remotePath := cfg.Path
	if remotePath == "" {
		home, _ := os.UserHomeDir()
		remotePath = filepath.Join(home, ".pi", "agent", "dashboard")
	}

	// SSH in, cat all JSON files
	cmd := exec.Command("ssh",
		"-o", "ConnectTimeout=5",
		"-o", "BatchMode=yes",
		fmt.Sprintf("%s@%s", user, hostPort),
		fmt.Sprintf("cat %s/*.json 2>/dev/null || true", remotePath),
	)
	output, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("ssh %s: %w", cfg.Host, err)
	}

	// Multiple pretty-printed JSON objects concatenated — stream-decode
	var states []AgentState
	dec := json.NewDecoder(strings.NewReader(string(output)))
	for {
		var s AgentState
		if err := dec.Decode(&s); err != nil {
			break
		}
		states = append(states, s)
	}
	return states, nil
}