feat(agentcompat): add scoped IO stream capabilities

Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>
2026-09-19 17:50:12 +00:00 · 2026-07-20 04:27:33 +00:00
parent c756ef9385
commit 26e92da33b
31 changed files with 3540 additions and 0 deletions
@@ -0,0 +1,197 @@
+//go:build agentcompat
+
+package rpc
+
+import (
+	"bufio"
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"strconv"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/nezhahq/nezha/model"
+	pb "github.com/nezhahq/nezha/proto"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMCPReceiptGate_FormatsTaskAndResultWithGeneration(t *testing.T) {
+	// Given
+	serverConn, clientConn := net.Pipe()
+	defer clientConn.Close()
+	gate := installReceiptGateForTest(serverConn)
+	defer clearReceiptGateForTest()
+	reader := bufio.NewReader(clientConn)
+
+	// When
+	taskDone := make(chan struct{})
+	go func() {
+		notifyMCPTaskDispatched(7, 9, model.TaskTypeExec)
+		close(taskDone)
+	}()
+	taskLine := mustReadLine(t, reader)
+	<-taskDone
+	resultDone := make(chan struct{})
+	go func() {
+		notifyMCPTaskResultAccepted(7, 9, model.TaskTypeExec)
+		close(resultDone)
+	}()
+	resultLine := mustReadLine(t, reader)
+	<-resultDone
+
+	// Then
+	require.Equal(t, "task "+itoa(gate.generation)+" 7 9 "+itoa(model.TaskTypeExec)+"\n", taskLine)
+	require.Equal(t, "result "+itoa(gate.generation)+" 7 9 "+itoa(model.TaskTypeExec)+"\n", resultLine)
+}
+
+func TestCallAgent_EmitsOneTaskAndOneAcceptedResult(t *testing.T) {
+	// Given
+	serverConn, clientConn := net.Pipe()
+	defer clientConn.Close()
+	gate := installReceiptGateForTest(serverConn)
+	defer clearReceiptGateForTest()
+	stream := newFakeStream()
+	cleanup := installFakeServer(t, 801, stream)
+	defer cleanup()
+	reader := bufio.NewReader(clientConn)
+	lines := make(chan string, 2)
+	go func() {
+		line, err := reader.ReadString('\n')
+		if err != nil {
+			return
+		}
+		lines <- line
+		line, err = reader.ReadString('\n')
+		if err == nil {
+			lines <- line
+		}
+	}()
+
+	go func() {
+		sent := <-stream.sent
+		deliverMCPResult(&pb.TaskResult{Id: sent.GetId(), Type: sent.GetType(), Successful: true, Data: "{}"})
+		deliverMCPResult(&pb.TaskResult{Id: sent.GetId(), Type: sent.GetType(), Successful: true, Data: "{}"})
+	}()
+
+	// When
+	_, err := CallAgent(context.Background(), 801, model.TaskTypeExec, model.ExecRequest{Cmd: "x"}, time.Second)
+
+	// Then
+	require.NoError(t, err)
+	taskLine := <-lines
+	resultLine := <-lines
+	require.Equal(t, "task "+itoa(gate.generation)+" 801 "+itoa(parseReceiptTaskID(taskLine))+" "+itoa(model.TaskTypeExec)+"\n", taskLine)
+	require.Equal(t, "result "+itoa(gate.generation)+" 801 "+itoa(parseReceiptTaskID(resultLine))+" "+itoa(model.TaskTypeExec)+"\n", resultLine)
+	require.Equal(t, parseReceiptTaskID(taskLine), parseReceiptTaskID(resultLine))
+	clientConn.SetReadDeadline(time.Now().Add(20 * time.Millisecond))
+	_, readErr := reader.ReadString('\n')
+	require.Error(t, readErr)
+}
+
+func TestCallAgent_SendFailureEmitsNoTaskReceipt(t *testing.T) {
+	// Given
+	serverConn, clientConn := net.Pipe()
+	defer clientConn.Close()
+	installReceiptGateForTest(serverConn)
+	defer clearReceiptGateForTest()
+	stream := &fakeTaskStream{sent: make(chan *pb.Task, 1), err: errors.New("send failed")}
+	cleanup := installFakeServer(t, 802, stream)
+	defer cleanup()
+
+	// When
+	_, err := CallAgent(context.Background(), 802, model.TaskTypeExec, model.ExecRequest{Cmd: "x"}, time.Second)
+
+	// Then
+	require.EqualError(t, err, "send failed")
+	clientConn.SetReadDeadline(time.Now().Add(20 * time.Millisecond))
+	_, readErr := bufio.NewReader(clientConn).ReadString('\n')
+	require.Error(t, readErr)
+}
+
+func TestCallAgent_LateDuplicateAndCancelledResultsEmitNoAcceptedReceipt(t *testing.T) {
+	// Given
+	serverConn, clientConn := net.Pipe()
+	defer clientConn.Close()
+	installReceiptGateForTest(serverConn)
+	defer clearReceiptGateForTest()
+	stream := newFakeStream()
+	cleanup := installFakeServer(t, 803, stream)
+	defer cleanup()
+	reader := bufio.NewReader(clientConn)
+	taskLineCh := make(chan string, 1)
+
+	// When
+	taskIDCh := make(chan uint64, 1)
+	go func() {
+		sent := <-stream.sent
+		taskIDCh <- sent.GetId()
+		line, _ := reader.ReadString('\n')
+		taskLineCh <- line
+	}()
+	_, err := CallAgent(context.Background(), 803, model.TaskTypeFsRead, model.FsReadRequest{Path: "/x"}, 20*time.Millisecond)
+	require.ErrorIs(t, err, ErrAgentTimeout)
+	taskID := <-taskIDCh
+	deliverMCPResult(&pb.TaskResult{Id: taskID, Type: model.TaskTypeFsRead, Successful: true, Data: "{}"})
+	deliverMCPResult(&pb.TaskResult{Id: taskID, Type: model.TaskTypeFsRead, Successful: true, Data: "{}"})
+
+	// Then
+	require.Contains(t, <-taskLineCh, "task ")
+	clientConn.SetReadDeadline(time.Now().Add(20 * time.Millisecond))
+	_, readErr := reader.ReadString('\n')
+	require.Error(t, readErr)
+}
+
+func TestCallAgent_CancelledResultEmitsNoAcceptedReceipt(t *testing.T) {
+	// Given
+	serverConn, clientConn := net.Pipe()
+	defer clientConn.Close()
+	installReceiptGateForTest(serverConn)
+	defer clearReceiptGateForTest()
+	stream := newFakeStream()
+	cleanup := installFakeServer(t, 804, stream)
+	defer cleanup()
+	reader := bufio.NewReader(clientConn)
+	taskLineCh := make(chan string, 1)
+	go func() {
+		line, _ := reader.ReadString('\n')
+		taskLineCh <- line
+	}()
+	taskIDCh := make(chan uint64, 1)
+	go func() {
+		sent := <-stream.sent
+		taskIDCh <- sent.GetId()
+	}()
+
+	// When
+	errCh := make(chan error, 1)
+	go func() {
+		_, err := CallAgent(context.Background(), 804, model.TaskTypeExec, model.ExecRequest{Cmd: "x"}, time.Second)
+		errCh <- err
+	}()
+	taskID := <-taskIDCh
+	CancelAllMCPInflight()
+	deliverMCPResult(&pb.TaskResult{Id: taskID, Type: model.TaskTypeExec, Successful: true, Data: "{}"})
+
+	// Then
+	require.ErrorIs(t, <-errCh, ErrMCPDisabled)
+	require.Contains(t, <-taskLineCh, "task ")
+	clientConn.SetReadDeadline(time.Now().Add(20 * time.Millisecond))
+	_, readErr := reader.ReadString('\n')
+	require.Error(t, readErr)
+}
+
+func itoa(value uint64) string {
+	return strconv.FormatUint(value, 10)
+}
+
+func parseReceiptTaskID(line string) uint64 {
+	fields := strings.Fields(line)
+	value, err := strconv.ParseUint(fields[3], 10, 64)
+	if err != nil {
+		panic(fmt.Sprintf("invalid receipt line %q: %v", line, err))
+	}
+	return value
+}