ccrm/daemon/cmd/ccrmd/main.go

package main

import (
	"bufio"
	"encoding/json"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/fmartingr/ccrm/daemon/internal/claude"
	"github.com/fmartingr/ccrm/daemon/internal/config"
	"github.com/fmartingr/ccrm/daemon/internal/relay"
	"github.com/fmartingr/ccrm/daemon/internal/session"
)

func main() {
	cfg, err := config.Load()
	if err != nil {
		log.Fatalf("Failed to load config: %v", err)
	}

	log.Printf("Starting ccrmd daemon...")

	// Resolve data directory
	dataDir := config.ExpandPath(cfg.Daemon.DataDir)

	// Declare both so closures can capture them by pointer
	var mgr *session.Manager
	var relayClient *relay.Client

	// Session manager with message handler and process exit callback
	mgr = session.NewManager(cfg.Daemon.TmuxPrefix, cfg.Claude.Binary, dataDir,
		func(sessionName string, msg *claude.CLIMessage) {
			handleSubprocessMessage(mgr, relayClient, sessionName, msg)
		},
		func(sessionName string, willRespawn bool) {
			if willRespawn {
				log.Printf("Session %s: subprocess exited, will respawn on next prompt", sessionName)
			} else {
				log.Printf("Session %s: subprocess exited (stopped)", sessionName)
				if relayClient != nil {
					_ = relayClient.Send(&relay.SessionEndedMsg{
						Type:        "session.ended",
						SessionName: sessionName,
					})
				}
			}
		},
	)

	// Recover sessions from tmux (survives daemon restart)
	if err := mgr.RecoverSessions(); err != nil {
		log.Printf("Warning: failed to recover sessions: %v", err)
	} else {
		// Re-register recovered sessions with relay client (done after relay starts)
		recovered := mgr.List()
		if len(recovered) > 0 {
			log.Printf("Recovered %d session(s) from tmux", len(recovered))
		}
	}

	// WebSocket relay client (optional)
	if cfg.Server.URL != "" {
		relayClient = relay.NewClient(cfg.Server.URL, cfg.Server.APIKey, func(msgType string, raw json.RawMessage) {
			handleServerMessage(mgr, relayClient, msgType, raw)
		})
		relayClient.Start()
		log.Printf("Relay client started, connecting to %s", cfg.Server.URL)

		// Re-register recovered sessions with relay
		for _, s := range mgr.List() {
			_ = relayClient.Send(&relay.SessionStartedMsg{
				Type:    "session.started",
				Session: s,
			})
		}
	}

	// Control socket listener
	shutdownCh := make(chan struct{})
	go controlSocketListener(cfg.Daemon.SocketPath, mgr, relayClient, shutdownCh)

	log.Printf("Daemon ready. Control socket: %s", cfg.Daemon.SocketPath)

	// Start heartbeat
	if relayClient != nil {
		go heartbeatLoop(mgr, relayClient)
	}

	// Wait for shutdown signal or control socket stop command
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	select {
	case <-sigCh:
	case <-shutdownCh:
	}

	log.Printf("Shutting down...")
	if relayClient != nil {
		relayClient.Stop()
	}
	// Stop all sessions
	for _, s := range mgr.List() {
		_ = mgr.Kill(s.Name)
	}
	log.Printf("Daemon stopped")
}

// --- Control socket ---

type controlCommand struct {
	Type        string   `json:"type"`
	SessionName string   `json:"session_name,omitempty"`
	Name        string   `json:"name,omitempty"`
	Path        string   `json:"path,omitempty"`
	ClaudeArgs  []string `json:"claude_args,omitempty"`
	Prompt      string   `json:"prompt,omitempty"`
}

type controlResponse struct {
	Status string `json:"status"` // "ok" or "error"
	Error  string `json:"error,omitempty"`
	Data   any    `json:"data,omitempty"`
}

func controlSocketListener(socketPath string, mgr *session.Manager, relayClient *relay.Client, shutdownCh chan struct{}) {
	_ = os.Remove(socketPath)

	ln, err := net.Listen("unix", socketPath)
	if err != nil {
		log.Fatalf("Failed to listen on control socket: %v", err)
	}
	defer ln.Close()
	defer os.Remove(socketPath)

	if err := os.Chmod(socketPath, 0666); err != nil {
		log.Printf("warning: failed to chmod socket: %v", err)
	}

	var wg sync.WaitGroup
	for {
		conn, err := ln.Accept()
		if err != nil {
			select {
			case <-shutdownCh:
				wg.Wait()
				return
			default:
				log.Printf("accept error: %v", err)
				continue
			}
		}

		wg.Add(1)
		go func() {
			defer wg.Done()
			handleControlConnection(conn, mgr, relayClient, shutdownCh)
		}()
	}
}

func handleControlConnection(conn net.Conn, mgr *session.Manager, relayClient *relay.Client, shutdownCh chan struct{}) {
	defer conn.Close()

	scanner := bufio.NewScanner(conn)
	scanner.Buffer(make([]byte, 1024*1024), 1024*1024)

	if !scanner.Scan() {
		return
	}

	var cmd controlCommand
	if err := json.Unmarshal(scanner.Bytes(), &cmd); err != nil {
		writeResponse(conn, controlResponse{Status: "error", Error: "invalid JSON"})
		return
	}

	var resp controlResponse

	switch cmd.Type {
	case "session.start":
		s, err := mgr.Start(cmd.Name, cmd.Path, cmd.ClaudeArgs)
		if err != nil {
			resp = controlResponse{Status: "error", Error: err.Error()}
		} else {
			// Notify server
			if relayClient != nil {
				_ = relayClient.Send(&relay.SessionStartedMsg{
					Type:    "session.started",
					Session: s,
				})
			}
			resp = controlResponse{Status: "ok", Data: s}
		}

	case "session.stop":
		if err := mgr.Stop(cmd.SessionName); err != nil {
			resp = controlResponse{Status: "error", Error: err.Error()}
		} else {
			if relayClient != nil {
				_ = relayClient.Send(&relay.SessionEndedMsg{
					Type:        "session.ended",
					SessionName: cmd.SessionName,
				})
			}
			resp = controlResponse{Status: "ok"}
		}

	case "session.kill":
		if err := mgr.Kill(cmd.SessionName); err != nil {
			resp = controlResponse{Status: "error", Error: err.Error()}
		} else {
			if relayClient != nil {
				_ = relayClient.Send(&relay.SessionEndedMsg{
					Type:        "session.ended",
					SessionName: cmd.SessionName,
				})
			}
			resp = controlResponse{Status: "ok"}
		}

	case "session.attach":
		tmuxName, err := mgr.Attach(cmd.SessionName)
		if err != nil {
			resp = controlResponse{Status: "error", Error: err.Error()}
		} else {
			resp = controlResponse{Status: "ok", Data: map[string]string{"tmux_session": tmuxName}}
		}

	case "sessions.list":
		sessions := mgr.List()
		resp = controlResponse{Status: "ok", Data: sessions}

	case "prompt.send":
		err := mgr.SendPrompt(cmd.SessionName, cmd.Prompt)
		if err != nil {
			resp = controlResponse{Status: "error", Error: err.Error()}
		} else {
			s, _ := mgr.Get(cmd.SessionName)
			status := "sent"
			if s != nil && len(s.PendingPrompts) > 0 {
				status = "queued"
			}
			resp = controlResponse{Status: "ok", Data: map[string]string{"prompt_status": status}}
		}

	case "daemon.stop":
		resp = controlResponse{Status: "ok"}
		writeResponse(conn, resp)
		close(shutdownCh)
		return

	default:
		resp = controlResponse{Status: "error", Error: "unknown command type: " + cmd.Type}
	}

	writeResponse(conn, resp)
}

func writeResponse(conn net.Conn, resp controlResponse) {
	data, _ := json.Marshal(resp)
	data = append(data, '\n')
	conn.Write(data)
}

// --- Subprocess message handling ---

func handleSubprocessMessage(mgr *session.Manager, relayClient *relay.Client, sessionName string, msg *claude.CLIMessage) {
	log.Printf("Subprocess message [%s]: type=%s", sessionName, msg.Type)

	switch msg.Type {
	case "system":
		var sysMsg claude.SystemMessage
		if err := json.Unmarshal(msg.Raw, &sysMsg); err != nil {
			log.Printf("Failed to parse system message: %v", err)
			return
		}
		mgr.UpdateClaudeSession(sessionName, sysMsg.SessionID, "")
		mgr.UpdateStatus(sessionName, session.StatusIdle)
		log.Printf("Session %s: Claude session ID = %s, model = %s", sessionName, sysMsg.SessionID, sysMsg.Model)

		if relayClient != nil {
			s, _ := mgr.Get(sessionName)
			if s != nil {
				_ = relayClient.Send(&relay.SessionStartedMsg{
					Type:    "session.started",
					Session: s,
				})
			}
		}

	case "assistant":
		mgr.UpdateStatus(sessionName, session.StatusActive)
		if relayClient != nil {
			_ = relayClient.Send(&relay.HookEventMsg{
				Type:        "hook.event",
				SessionName: sessionName,
				Event: map[string]any{
					"hook_event_name": "AssistantMessage",
					"session_id":     getClaudeSessionID(mgr, sessionName),
					"raw":            msg.Raw,
				},
			})
		}

	case "result":
		var resultMsg claude.ResultMessage
		if err := json.Unmarshal(msg.Raw, &resultMsg); err != nil {
			log.Printf("Failed to parse result message: %v", err)
			return
		}
		// Fallback: extract session_id from result if system message was missed
		if resultMsg.SessionID != "" {
			mgr.UpdateClaudeSession(sessionName, resultMsg.SessionID, "")
		}
		mgr.UpdateStatus(sessionName, session.StatusIdle)
		mgr.DrainPendingPrompts(sessionName)
		log.Printf("Session %s: turn complete (cost=$%.4f, total=$%.4f)", sessionName, resultMsg.CostUSD, resultMsg.TotalCostUSD)

		if relayClient != nil {
			_ = relayClient.Send(&relay.HookEventMsg{
				Type:        "hook.event",
				SessionName: sessionName,
				Event: map[string]any{
					"hook_event_name": "Stop",
					"session_id":     getClaudeSessionID(mgr, sessionName),
				},
			})
		}

	case "tool_use_permission":
		var ctrlReq claude.ControlRequest
		if err := json.Unmarshal(msg.Raw, &ctrlReq); err != nil {
			log.Printf("Failed to parse control request: %v", err)
			return
		}
		mgr.UpdateStatus(sessionName, session.StatusWaitingPermission)
		log.Printf("Session %s: permission request for tool %s (request_id=%s) — auto-allowing",
			sessionName, ctrlReq.Tool.ToolName, ctrlReq.Tool.RequestID)

		// Auto-allow in PoC
		if err := mgr.RespondToPermission(sessionName, ctrlReq.Tool.RequestID, true, ""); err != nil {
			log.Printf("Failed to auto-allow permission: %v", err)
		}
		mgr.UpdateStatus(sessionName, session.StatusActive)

		if relayClient != nil {
			_ = relayClient.Send(&relay.HookEventMsg{
				Type:        "hook.event",
				SessionName: sessionName,
				Event: map[string]any{
					"hook_event_name": "PermissionRequest",
					"session_id":     getClaudeSessionID(mgr, sessionName),
					"tool_name":      ctrlReq.Tool.ToolName,
					"tool_input":     ctrlReq.Tool.Input,
				},
			})
		}

	default:
		log.Printf("Session %s: unhandled message type %q", sessionName, msg.Type)
	}
}

func getClaudeSessionID(mgr *session.Manager, name string) string {
	s, ok := mgr.Get(name)
	if !ok {
		return ""
	}
	return s.ClaudeSessionID
}


// --- Server message handling (for remote commands) ---

func handleServerMessage(mgr *session.Manager, relayClient *relay.Client, msgType string, raw json.RawMessage) {
	switch msgType {
	case "prompt.send":
		var msg relay.PromptSendMsg
		if err := json.Unmarshal(raw, &msg); err != nil {
			log.Printf("Failed to parse prompt.send: %v", err)
			return
		}

		err := mgr.SendPrompt(msg.SessionName, msg.Prompt)
		ack := &relay.PromptAckMsg{
			Type:        "prompt.ack",
			SessionName: msg.SessionName,
		}
		if err != nil {
			ack.Status = "rejected"
			ack.Reason = err.Error()
		} else {
			s, _ := mgr.Get(msg.SessionName)
			if s != nil && len(s.PendingPrompts) > 0 {
				ack.Status = "queued"
			} else {
				ack.Status = "sent"
			}
		}
		_ = relayClient.Send(ack)

	case "sessions.list":
		sessions := mgr.List()
		for _, s := range sessions {
			_ = relayClient.Send(&relay.SessionStartedMsg{
				Type:    "session.started",
				Session: s,
			})
		}

	case "session.kill":
		var msg relay.SessionKillMsg
		if err := json.Unmarshal(raw, &msg); err != nil {
			return
		}
		if err := mgr.Kill(msg.SessionName); err != nil {
			log.Printf("Failed to kill session %s: %v", msg.SessionName, err)
		} else {
			_ = relayClient.Send(&relay.SessionEndedMsg{
				Type:        "session.ended",
				SessionName: msg.SessionName,
			})
		}

	case "session.start":
		var msg relay.SessionStartMsg
		if err := json.Unmarshal(raw, &msg); err != nil {
			return
		}
		s, err := mgr.Start(msg.Name, msg.Path, msg.ClaudeArgs)
		if err != nil {
			log.Printf("Failed to start session %s: %v", msg.Name, err)
		} else {
			_ = relayClient.Send(&relay.SessionStartedMsg{
				Type:    "session.started",
				Session: s,
			})
		}
	}
}

func sendHeartbeat(mgr *session.Manager, relayClient *relay.Client) {
	sessions := mgr.List()
	names := make([]string, 0, len(sessions))
	for _, s := range sessions {
		if !s.ExplicitStop {
			names = append(names, s.Name)
		}
	}
	_ = relayClient.Send(&relay.HeartbeatMsg{
		Type:      "heartbeat",
		Sessions:  names,
		Timestamp: time.Now().UTC().Format(time.RFC3339),
	})
}

func heartbeatLoop(mgr *session.Manager, relayClient *relay.Client) {
	// Send immediately on start to reconcile stale sessions
	sendHeartbeat(mgr, relayClient)

	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for range ticker.C {
		sendHeartbeat(mgr, relayClient)
	}
}