sproutclaw-web/backend/internal/rpc/client.go

package rpc

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os/exec"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"sproutclaw-web/internal/models"
)

const (
	cmdTimeout    = 60 * time.Second
	promptTimeout = 30 * time.Second
)

// SSEClient represents a connected browser SSE client.
type SSEClient struct {
	Ch     chan []byte
	Done   <-chan struct{}
}

// RunSnapshot holds a snapshot of the current pi agent run state.
type RunSnapshot struct {
	IsStreaming  bool
	SessionFile  string
	ReplayEvents []json.RawMessage
}

// Client manages a pi CLI subprocess and bridges its JSON-RPC protocol
// to HTTP handlers and SSE clients.
type Client struct {
	mu         sync.Mutex
	cmd        *exec.Cmd
	stdin      io.WriteCloser
	pending    map[string]chan models.RPCResponse
	sseClients map[*SSEClient]struct{}
	replay     []json.RawMessage // events in current agent run
	isRunning  bool
	sessionFile string
}

// NewClient starts the pi CLI subprocess and begins reading its stdout.
// onExit is called with the exit code when the subprocess terminates.
func NewClient(piCmd string, piArgs []string, agentDir string, onExit func(int)) (*Client, error) {
	if piCmd == "" {
		return nil, fmt.Errorf("pi CLI command not specified (use --pi-cmd)")
	}

	args := append(piArgs, "--rpc")
	cmd := exec.Command(piCmd, args...)
	cmd.Dir = agentDir
	cmd.Env = append(cmd.Environ(), fmt.Sprintf("PI_CODING_AGENT_DIR=%s", agentDir))

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("stdin pipe: %w", err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("stdout pipe: %w", err)
	}
	cmd.Stderr = log.Writer()

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("start pi cli: %w", err)
	}

	c := &Client{
		cmd:        cmd,
		stdin:      stdin,
		pending:    make(map[string]chan models.RPCResponse),
		sseClients: make(map[*SSEClient]struct{}),
	}

	go c.readLoop(stdout)
	go func() {
		_ = cmd.Wait()
		code := 0
		if cmd.ProcessState != nil {
			code = cmd.ProcessState.ExitCode()
		}
		onExit(code)
	}()

	return c, nil
}

// readLoop reads stdout line-by-line and dispatches messages.
func (c *Client) readLoop(r io.Reader) {
	scanner := bufio.NewScanner(r)
	scanner.Buffer(make([]byte, 4*1024*1024), 4*1024*1024)
	for scanner.Scan() {
		line := scanner.Bytes()
		if len(line) == 0 {
			continue
		}
		c.dispatch(line)
	}
}

// dispatch parses one JSON line from stdout and routes it.
func (c *Client) dispatch(raw []byte) {
	var msg struct {
		Type string `json:"type"`
		ID   string `json:"id"`
	}
	if err := json.Unmarshal(raw, &msg); err != nil {
		return
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	switch msg.Type {
	case "response":
		var resp models.RPCResponse
		if err := json.Unmarshal(raw, &resp); err == nil {
			if ch, ok := c.pending[resp.ID]; ok {
				delete(c.pending, resp.ID)
				ch <- resp
			}
		}

	case "agent_start":
		c.isRunning = true
		c.replay = nil
		// try to extract session path
		var ev struct {
			SessionFile string `json:"sessionFile"`
		}
		if err := json.Unmarshal(raw, &ev); err == nil && ev.SessionFile != "" {
			c.sessionFile = ev.SessionFile
		}
		rawCopy := append([]byte(nil), raw...)
		c.replay = append(c.replay, rawCopy)
		c.broadcastLocked(rawCopy)

	case "agent_end":
		c.isRunning = false
		rawCopy := append([]byte(nil), raw...)
		c.replay = append(c.replay, rawCopy)
		c.broadcastLocked(rawCopy)

	default:
		rawCopy := append([]byte(nil), raw...)
		if c.isRunning {
			c.replay = append(c.replay, rawCopy)
		}
		c.broadcastLocked(rawCopy)
	}
}

// broadcastLocked sends raw JSON to all SSE clients. Must hold c.mu.
func (c *Client) broadcastLocked(raw []byte) {
	for client := range c.sseClients {
		select {
		case client.Ch <- raw:
		default:
			// slow client: drop
		}
	}
}

// send writes a JSON command to the pi CLI stdin.
func (c *Client) send(cmd models.RPCCommand) error {
	b, err := json.Marshal(cmd)
	if err != nil {
		return err
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	_, err = fmt.Fprintf(c.stdin, "%s\n", b)
	return err
}

// SendCmd sends a command and waits for the matching response.
func (c *Client) SendCmd(cmd models.RPCCommand) (models.RPCResponse, error) {
	if cmd.ID == "" {
		cmd.ID = uuid.New().String()
	}
	ch := make(chan models.RPCResponse, 1)

	c.mu.Lock()
	c.pending[cmd.ID] = ch
	c.mu.Unlock()

	if err := c.send(cmd); err != nil {
		c.mu.Lock()
		delete(c.pending, cmd.ID)
		c.mu.Unlock()
		return models.RPCResponse{}, err
	}

	select {
	case resp := <-ch:
		return resp, nil
	case <-time.After(cmdTimeout):
		c.mu.Lock()
		delete(c.pending, cmd.ID)
		c.mu.Unlock()
		return models.RPCResponse{}, fmt.Errorf("rpc timeout: %s", cmd.Type)
	}
}

// SubmitPrompt sends a prompt without waiting for a response.
func (c *Client) SubmitPrompt(req models.ChatRequest) error {
	cmd := models.RPCCommand{
		Type:              "prompt",
		ID:                uuid.New().String(),
		Text:              req.Message,
		Images:            req.Images,
		StreamingBehavior: req.StreamingBehavior,
	}
	return c.send(cmd)
}

// GetSnapshot returns the current run snapshot.
func (c *Client) GetSnapshot() RunSnapshot {
	c.mu.Lock()
	defer c.mu.Unlock()
	replay := make([]json.RawMessage, len(c.replay))
	copy(replay, c.replay)
	return RunSnapshot{
		IsStreaming:  c.isRunning,
		SessionFile:  c.sessionFile,
		ReplayEvents: replay,
	}
}

// RegisterSSEClient adds a client to the broadcast set and replays buffered events.
func (c *Client) RegisterSSEClient(client *SSEClient) {
	c.mu.Lock()
	c.sseClients[client] = struct{}{}
	// replay current run events
	replay := make([]json.RawMessage, len(c.replay))
	copy(replay, c.replay)
	c.mu.Unlock()

	// send replay in a goroutine so we don't hold the lock
	go func() {
		for _, ev := range replay {
			select {
			case client.Ch <- ev:
			case <-client.Done:
				return
			}
		}
	}()
}

// UnregisterSSEClient removes a client from the broadcast set.
func (c *Client) UnregisterSSEClient(client *SSEClient) {
	c.mu.Lock()
	delete(c.sseClients, client)
	c.mu.Unlock()
}

// SendExtensionUIResponse forwards a UI response back to the pi CLI.
func (c *Client) SendExtensionUIResponse(id string, response json.RawMessage) error {
	return c.send(models.RPCCommand{
		Type:     "extension_ui_response",
		ID:       id,
		Response: response,
	})
}

// Stop terminates the pi CLI subprocess.
func (c *Client) Stop() {
	if c.cmd != nil && c.cmd.Process != nil {
		_ = c.cmd.Process.Kill()
	}
}

// buildEnv returns os.Environ() with extra variables appended.
func (c *Client) Environ() []string {
	return c.cmd.Environ()
}

// helper
func mustMarshal(v any) []byte {
	b, _ := json.Marshal(v)
	return b
}

// IsRunning returns true when the agent is actively streaming.
func (c *Client) IsRunning() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.isRunning
}

// FormatSSEEvent formats a raw JSON payload as an SSE data line.
func FormatSSEEvent(raw []byte) []byte {
	// strip newlines so the SSE frame stays on one logical line
	trimmed := strings.TrimRight(string(raw), "\n\r")
	return []byte("data: " + trimmed + "\n\n")
}