first commit

backend/internal/rpc/client.go

@@ -6,12 +6,13 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"os"
 	"os/exec"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
 
-	"github.com/google/uuid"
 	"sproutclaw-web/internal/models"
 )
 
@@ -20,10 +21,29 @@ const (
 	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{}
+	Ch   chan []byte
+	Done <-chan struct{}
 }
 
 // RunSnapshot holds a snapshot of the current pi agent run state.
@@ -36,27 +56,33 @@ type RunSnapshot struct {
 // 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
+	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.
-// onExit is called with the exit code when the subprocess terminates.
-func NewClient(piCmd string, piArgs []string, agentDir string, onExit func(int)) (*Client, error) {
+// 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)")
 	}
 
-	args := append(piArgs, "--rpc")
+	// pi enters RPC mode via "--mode rpc"
+	args := append(append([]string{}, piArgs...), "--mode", "rpc")
 	cmd := exec.Command(piCmd, args...)
-	cmd.Dir = agentDir
-	cmd.Env = append(cmd.Environ(), fmt.Sprintf("PI_CODING_AGENT_DIR=%s", agentDir))
+	cmd.Dir = repoRoot
+	cmd.Env = append(os.Environ(), fmt.Sprintf("PI_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 {
@@ -66,12 +92,15 @@ func NewClient(piCmd string, piArgs []string, agentDir string, onExit func(int))
 	if err != nil {
 		return nil, fmt.Errorf("stdout pipe: %w", err)
 	}
-	cmd.Stderr = log.Writer()
+	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,
@@ -86,82 +115,103 @@ func NewClient(piCmd string, piArgs []string, agentDir string, onExit func(int))
 		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), 4*1024*1024)
+	scanner.Buffer(make([]byte, 4*1024*1024), 16*1024*1024)
 	for scanner.Scan() {
 		line := scanner.Bytes()
 		if len(line) == 0 {
 			continue
 		}
-		c.dispatch(line)
+		c.dispatch(append([]byte(nil), line...))
 	}
 }
 
 // dispatch parses one JSON line from stdout and routes it.
 func (c *Client) dispatch(raw []byte) {
-	var msg struct {
+	var head struct {
 		Type string `json:"type"`
 		ID   string `json:"id"`
 	}
-	if err := json.Unmarshal(raw, &msg); err != nil {
-		return
+	if err := json.Unmarshal(raw, &head); err != nil {
+		return // ignore non-JSON lines
 	}
 
-	c.mu.Lock()
-	defer c.mu.Unlock()
-
-	switch msg.Type {
-	case "response":
+	// Response frame: resolve the pending request.
+	if head.Type == "response" && head.ID != "" {
 		var resp models.RPCResponse
 		if err := json.Unmarshal(raw, &resp); err == nil {
-			if ch, ok := c.pending[resp.ID]; ok {
+			// 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 = 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)
-
+		c.replay = []json.RawMessage{raw}
 	case "agent_end":
+		c.replay = append(c.replay, raw)
 		c.isRunning = false
-		rawCopy := append([]byte(nil), raw...)
-		c.replay = append(c.replay, rawCopy)
-		c.broadcastLocked(rawCopy)
-
+		c.replay = nil
 	default:
-		rawCopy := append([]byte(nil), raw...)
-		if c.isRunning {
-			c.replay = append(c.replay, rawCopy)
+		if c.isRunning && bufferedEventTypes[typ] {
+			c.replay = append(c.replay, raw)
 		}
-		c.broadcastLocked(rawCopy)
 	}
 }
 
 // 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 <- raw:
+		case client.Ch <- frame:
 		default:
 			// slow client: drop
 		}
@@ -176,15 +226,13 @@ func (c *Client) send(cmd models.RPCCommand) error {
 	}
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	_, err = fmt.Fprintf(c.stdin, "%s\n", b)
+	_, 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) {
-	if cmd.ID == "" {
-		cmd.ID = uuid.New().String()
-	}
+	cmd.ID = c.nextReqID()
 	ch := make(chan models.RPCResponse, 1)
 
 	c.mu.Lock()
@@ -205,20 +253,50 @@ func (c *Client) SendCmd(cmd models.RPCCommand) (models.RPCResponse, error) {
 		c.mu.Lock()
 		delete(c.pending, cmd.ID)
 		c.mu.Unlock()
-		return models.RPCResponse{}, fmt.Errorf("rpc timeout: %s", cmd.Type)
+		return models.RPCResponse{}, fmt.Errorf("命令超时: %s", cmd.Type)
 	}
 }
 
-// SubmitPrompt sends a prompt without waiting for a response.
+// 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:                uuid.New().String(),
-		Text:              req.Message,
+		ID:                id,
+		Message:           req.Message,
 		Images:            req.Images,
 		StreamingBehavior: req.StreamingBehavior,
 	}
-	return c.send(cmd)
+
+	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.
@@ -234,25 +312,26 @@ func (c *Client) GetSnapshot() RunSnapshot {
 	}
 }
 
-// RegisterSSEClient adds a client to the broadcast set and replays buffered events.
+// 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 current run events
 	replay := make([]json.RawMessage, len(c.replay))
 	copy(replay, c.replay)
+	isStreaming := c.isRunning
 	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
-			}
-		}
-	}()
+	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.
@@ -263,12 +342,23 @@ func (c *Client) UnregisterSSEClient(client *SSEClient) {
 }
 
 // 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 {
-	return c.send(models.RPCCommand{
-		Type:     "extension_ui_response",
-		ID:       id,
-		Response: response,
-	})
+	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.
@@ -278,17 +368,6 @@ func (c *Client) Stop() {
 	}
 }
 
-// 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()
@@ -296,9 +375,8 @@ func (c *Client) IsRunning() bool {
 	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")
+// 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")
 }