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

package rpc

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

	"sproutclaw-web/internal/models"
)

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

// bufferedEventTypes are the event types replayed to newly-connected SSE
// clients while the agent is mid-run (mirrors the original pi-client.ts).
var bufferedEventTypes = map[string]bool{
	"agent_start":            true,
	"agent_end":              true,
	"message_start":          true,
	"message_update":         true,
	"message_end":            true,
	"tool_execution_start":   true,
	"tool_execution_update":  true,
	"tool_execution_end":     true,
	"compaction_start":       true,
	"compaction_end":         true,
	"queue_update":           true,
	"auto_retry_start":       true,
	"auto_retry_end":         true,
	"bash_update":            true,
}

// 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
	reqMu       sync.Mutex
	reqCounter  int
}

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

	// pi enters RPC mode via "--mode rpc"
	args := append(append([]string{}, piArgs...), "--mode", "rpc")
	cmd := exec.Command(piCmd, args...)
	cmd.Dir = repoRoot
	cmd.Env = append(os.Environ(), fmt.Sprintf("SPROUTCLAW_CODING_AGENT_DIR=%s", agentDir))

	log.Printf("[sproutclaw-web] launching pi RPC: %s %s (cwd=%s)", piCmd, strings.Join(args, " "), repoRoot)

	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 = os.Stderr

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

	// Keep stdin open; closing it on Windows can trigger RPC shutdown.
	_, _ = stdin.Write([]byte(""))

	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()
		}
		log.Printf("[sproutclaw-web] pi RPC exited, code=%d", code)
		onExit(code)
	}()

	return c, nil
}

func (c *Client) nextReqID() string {
	c.reqMu.Lock()
	c.reqCounter++
	id := c.reqCounter
	c.reqMu.Unlock()
	return "req_" + strconv.Itoa(id)
}

// 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), 16*1024*1024)
	for scanner.Scan() {
		line := scanner.Bytes()
		if len(line) == 0 {
			continue
		}
		c.dispatch(append([]byte(nil), line...))
	}
}

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

	// Response frame: resolve the pending request.
	if head.Type == "response" && head.ID != "" {
		var resp models.RPCResponse
		if err := json.Unmarshal(raw, &resp); err == nil {
			// capture sessionFile from get_state responses
			if resp.Command == "get_state" && resp.Success {
				var data struct {
					SessionFile string `json:"sessionFile"`
				}
				if json.Unmarshal(resp.Data, &data) == nil && data.SessionFile != "" {
					c.mu.Lock()
					c.sessionFile = data.SessionFile
					c.mu.Unlock()
				}
			}
			c.mu.Lock()
			ch, ok := c.pending[resp.ID]
			if ok {
				delete(c.pending, resp.ID)
			}
			c.mu.Unlock()
			if ok {
				ch <- resp
			}
			return
		}
	}

	// Event frame: track agent run state + broadcast to SSE clients.
	c.trackAgentEvent(head.Type, raw)
	c.mu.Lock()
	c.broadcastLocked(raw)
	c.mu.Unlock()
}

func (c *Client) trackAgentEvent(typ string, raw []byte) {
	c.mu.Lock()
	defer c.mu.Unlock()
	switch typ {
	case "agent_start":
		c.isRunning = true
		c.replay = []json.RawMessage{raw}
	case "agent_end":
		c.replay = append(c.replay, raw)
		c.isRunning = false
		c.replay = nil
	default:
		if c.isRunning && bufferedEventTypes[typ] {
			c.replay = append(c.replay, raw)
		}
	}
}

// broadcastLocked sends raw JSON to all SSE clients. Must hold c.mu.
func (c *Client) broadcastLocked(raw []byte) {
	frame := formatSSEData(raw)
	for client := range c.sseClients {
		select {
		case client.Ch <- frame:
		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 = c.stdin.Write(append(b, '\n'))
	return err
}

// SendCmd sends a command and waits for the matching response.
func (c *Client) SendCmd(cmd models.RPCCommand) (models.RPCResponse, error) {
	cmd.ID = c.nextReqID()
	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("命令超时: %s", cmd.Type)
	}
}

// SubmitPrompt sends a prompt. It registers a preflight pending entry but does
// not block the HTTP handler — agent output arrives via SSE.
func (c *Client) SubmitPrompt(req models.ChatRequest) error {
	id := c.nextReqID()
	cmd := models.RPCCommand{
		Type:              "prompt",
		ID:                id,
		Message:           req.Message,
		Images:            req.Images,
		StreamingBehavior: req.StreamingBehavior,
	}

	ch := make(chan models.RPCResponse, 1)
	c.mu.Lock()
	c.pending[id] = ch
	c.mu.Unlock()

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

	// Preflight: wait briefly for a rejection; success streams via SSE.
	go func() {
		select {
		case resp := <-ch:
			if !resp.Success {
				c.mu.Lock()
				c.broadcastLocked([]byte(fmt.Sprintf(`{"type":"prompt_rejected","error":%q}`, resp.Error)))
				c.mu.Unlock()
			}
		case <-time.After(promptTimeout):
			c.mu.Lock()
			delete(c.pending, id)
			c.mu.Unlock()
		}
	}()
	return nil
}

// 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 and sends the initial "connected" frame
// carrying the current run snapshot (matches the original pi-client.ts).
func (c *Client) RegisterSSEClient(client *SSEClient) {
	c.mu.Lock()
	c.sseClients[client] = struct{}{}
	replay := make([]json.RawMessage, len(c.replay))
	copy(replay, c.replay)
	isStreaming := c.isRunning
	c.mu.Unlock()

	connected := map[string]any{
		"type":        "connected",
		"isStreaming": isStreaming,
		"replay":      replay,
	}
	b, _ := json.Marshal(connected)
	select {
	case client.Ch <- formatSSEData(b):
	case <-client.Done:
	}
}

// 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.
// The response object is merged into the top-level frame: {type, id, ...response}.
func (c *Client) SendExtensionUIResponse(id string, response json.RawMessage) error {
	merged := map[string]any{"type": "extension_ui_response", "id": id}
	var extra map[string]any
	if json.Unmarshal(response, &extra) == nil {
		for k, v := range extra {
			merged[k] = v
		}
	}
	b, err := json.Marshal(merged)
	if err != nil {
		return err
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	_, err = c.stdin.Write(append(b, '\n'))
	return err
}

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

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

// formatSSEData wraps a raw JSON payload as an SSE "data:" frame.
func formatSSEData(raw []byte) []byte {
	trimmed := strings.TrimRight(string(raw), "\r\n")
	return []byte("data: " + trimmed + "\n\n")
}