fix(mcp): harden dashboard dispatch lifecycle

Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>

cmd/dashboard/controller/mcp_transfer_lifecycle.go

@@ -0,0 +1,84 @@
+package controller
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"io"
+	"sync"
+	"time"
+
+	"github.com/hashicorp/go-uuid"
+
+	"github.com/nezhahq/nezha/model"
+	pb "github.com/nezhahq/nezha/proto"
+	"github.com/nezhahq/nezha/service/rpc"
+	"github.com/nezhahq/nezha/service/singleton"
+)
+
+// openFsTransferStream owns the task-to-agent IOStream lifecycle. The returned
+// cleanup is safe for concurrent callers and shares ownership with cancellation.
+func openFsTransferStream(ctx context.Context, serverID uint64, req *model.FsTransferRequest) (io.ReadWriteCloser, func(), error) {
+	if singleton.Conf == nil || !singleton.Conf.MCPEnabled() {
+		return nil, func() {}, errors.New("MCP is disabled by the dashboard administrator")
+	}
+	server, _ := singleton.ServerShared.Get(serverID)
+	if server == nil || server.GetTaskStream() == nil {
+		return nil, func() {}, errors.New("server offline")
+	}
+	handler := rpc.NezhaHandlerSingleton
+
+	streamID, err := uuid.GenerateUUID()
+	if err != nil {
+		return nil, func() {}, err
+	}
+	req.StreamID = streamID
+	if err := handler.CreateStreamWithPurpose(streamID, 0, serverID, rpc.PurposeMCPTransfer); err != nil {
+		return nil, func() {}, err
+	}
+	var cleanupOnce sync.Once
+	cleanup := func() { cleanupOnce.Do(func() { _ = handler.CloseStream(streamID) }) }
+
+	body, err := json.Marshal(req)
+	if err != nil {
+		cleanup()
+		return nil, func() {}, err
+	}
+	// The stream is owned by cleanup until the caller receives it; every failure path releases it.
+	if singleton.Conf == nil || !singleton.Conf.MCPEnabled() {
+		cleanup()
+		return nil, func() {}, errors.New("MCP is disabled by the dashboard administrator")
+	}
+	if err := ctx.Err(); err != nil {
+		cleanup()
+		return nil, func() {}, err
+	}
+	if err := server.SendTask(&pb.Task{Type: model.TaskTypeFsTransfer, Data: string(body)}); err != nil {
+		cleanup()
+		if errors.Is(err, model.ErrTaskStreamOffline) {
+			return nil, func() {}, errors.New("server offline")
+		}
+		return nil, func() {}, err
+	}
+
+	agentStream, ok := handler.WaitForAgent(ctx, streamID, 30*time.Second)
+	if !ok {
+		cleanup()
+		return nil, func() {}, errors.New("agent did not attach within 30s")
+	}
+
+	watcherDone := make(chan struct{})
+	var watcherOnce sync.Once
+	go func() {
+		select {
+		case <-ctx.Done():
+			cleanup()
+		case <-watcherDone:
+		}
+	}()
+	wrappedCleanup := func() {
+		watcherOnce.Do(func() { close(watcherDone) })
+		cleanup()
+	}
+	return agentStream, wrappedCleanup, nil
+}