fix(rpc): make IO stream lifecycle race-safe

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

pkg/grpcx/io_stream_wrapper.go

@@ -120,6 +120,9 @@ func (iw *IOStreamWrapper) Write(p []byte) (n int, err error) {
 func (iw *IOStreamWrapper) Close() error {
 	if iw.closed.CompareAndSwap(false, true) {
 		close(iw.closeCh)
+		if closer, ok := iw.IOStream.(interface{ Close() error }); ok {
+			return closer.Close()
+		}
 	}
 	return nil
 }

pkg/utils/request_wrapper.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"sync"
 )
 
 var _ io.ReadWriteCloser = (*RequestWrapper)(nil)
@@ -14,6 +15,11 @@ type RequestWrapper struct {
 	req    *http.Request
 	reader *bytes.Buffer
 	writer net.Conn
+
+	closeOnce sync.Once
+	closeInit sync.Once
+	closeDone chan struct{}
+	closeErr  error
 }
 
 func NewRequestWrapper(req *http.Request, writer http.ResponseWriter) (*RequestWrapper, error) {
@@ -27,12 +33,17 @@ func NewRequestWrapper(req *http.Request, writer http.ResponseWriter) (*RequestW
 	}
 	buf := bytes.NewBuffer(nil)
 	if err = req.Write(buf); err != nil {
-		return nil, err
+		var bodyErr error
+		if req.Body != nil {
+			bodyErr = req.Body.Close()
+		}
+		return nil, errors.Join(err, bodyErr, conn.Close())
 	}
 	return &RequestWrapper{
-		req:    req,
-		reader: buf,
-		writer: conn,
+		req:       req,
+		reader:    buf,
+		writer:    conn,
+		closeDone: make(chan struct{}),
 	}, nil
 }
 
@@ -53,7 +64,17 @@ func (rw *RequestWrapper) Write(p []byte) (int, error) {
 }
 
 func (rw *RequestWrapper) Close() error {
-	rw.req.Body.Close()
-	rw.writer.Close()
-	return nil
+	rw.closeInit.Do(func() {
+		rw.closeDone = make(chan struct{})
+	})
+	rw.closeOnce.Do(func() {
+		var bodyErr error
+		if rw.req.Body != nil {
+			bodyErr = rw.req.Body.Close()
+		}
+		rw.closeErr = errors.Join(bodyErr, rw.writer.Close())
+		close(rw.closeDone)
+	})
+	<-rw.closeDone
+	return rw.closeErr
 }

pkg/utils/request_wrapper_test.go

@@ -0,0 +1,232 @@
+package utils
+
+import (
+	"bufio"
+	"bytes"
+	"errors"
+	"io"
+	"net"
+	"net/http"
+	"net/url"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+)
+
+var (
+	errRequestWrite = errors.New("request write failed")
+	errBodyClose    = errors.New("body close failed")
+	errConnClose    = errors.New("connection close failed")
+)
+
+func TestNewRequestWrapper_closesHijackedConnectionWhenRequestWriteFails(t *testing.T) {
+	// Given
+	conn := newRequestWrapperTestConn(errConnClose)
+	req := &http.Request{
+		Method:        "POST",
+		URL:           &url.URL{Scheme: "http", Host: "example.test", Path: "/nat"},
+		Body:          &requestWrapperTestBody{readErr: errRequestWrite, closeErr: errBodyClose},
+		ContentLength: 1,
+	}
+	writer := &requestWrapperTestResponseWriter{conn: conn}
+
+	// When
+	_, err := NewRequestWrapper(req, writer)
+
+	// Then
+	if err == nil || !strings.Contains(err.Error(), errRequestWrite.Error()) {
+		t.Fatalf("expected request write error, got %v", err)
+	}
+	if !errors.Is(err, errBodyClose) {
+		t.Fatalf("expected body close error, got %v", err)
+	}
+	if !errors.Is(err, errConnClose) {
+		t.Fatalf("expected connection close error, got %v", err)
+	}
+	if got := conn.closeCount.Load(); got != 1 {
+		t.Fatalf("expected one connection close, got %d", got)
+	}
+}
+
+func TestRequestWrapper_Close_joinsBodyAndConnectionErrors(t *testing.T) {
+	// Given
+	body := &requestWrapperTestBody{closeErr: errBodyClose}
+	conn := newRequestWrapperTestConn(errConnClose)
+	rw := &RequestWrapper{
+		req:       &http.Request{Body: body},
+		reader:    bytes.NewBuffer(nil),
+		writer:    conn,
+		closeDone: make(chan struct{}),
+	}
+
+	// When
+	err := rw.Close()
+
+	// Then
+	if !errors.Is(err, errBodyClose) {
+		t.Fatalf("expected body close error, got %v", err)
+	}
+	if !errors.Is(err, errConnClose) {
+		t.Fatalf("expected connection close error, got %v", err)
+	}
+}
+
+func TestRequestWrapper_Close_repeatedCallersReceiveRetainedErrorAndCloseOnce(t *testing.T) {
+	// Given
+	body := &requestWrapperTestBody{closeErr: errBodyClose}
+	conn := newRequestWrapperTestConn(errConnClose)
+	rw := &RequestWrapper{
+		req:       &http.Request{Body: body},
+		reader:    bytes.NewBuffer(nil),
+		writer:    conn,
+		closeDone: make(chan struct{}),
+	}
+
+	// When
+	firstErr := rw.Close()
+	secondErr := rw.Close()
+
+	// Then
+	if firstErr != secondErr {
+		t.Fatalf("expected identical retained error, got distinct values %p and %p", firstErr, secondErr)
+	}
+	if got := body.closeCount.Load(); got != 1 {
+		t.Fatalf("expected one body close, got %d", got)
+	}
+	if got := conn.closeCount.Load(); got != 1 {
+		t.Fatalf("expected one connection close, got %d", got)
+	}
+}
+
+func TestRequestWrapper_Close_concurrentCallersWaitForCopyToUnblock(t *testing.T) {
+	// Given
+	body := &requestWrapperTestBody{closeErr: errBodyClose}
+	conn := newRequestWrapperTestConn(errConnClose)
+	rw := &RequestWrapper{
+		req:       &http.Request{Body: body},
+		reader:    bytes.NewBuffer(nil),
+		writer:    conn,
+		closeDone: make(chan struct{}),
+	}
+	readDone := make(chan error, 1)
+	go func() {
+		_, err := rw.Read(make([]byte, 1))
+		readDone <- err
+	}()
+	<-conn.readStarted
+
+	// When
+	const callerCount = 8
+	results := make(chan error, callerCount)
+	var callers sync.WaitGroup
+	callers.Add(callerCount)
+	for range callerCount {
+		go func() {
+			defer callers.Done()
+			results <- rw.Close()
+		}()
+	}
+	callers.Wait()
+	close(results)
+
+	// Then
+	var retainedErr error
+	for err := range results {
+		if !errors.Is(err, errConnClose) {
+			t.Fatalf("expected retained connection close error, got %v", err)
+		}
+		if !errors.Is(err, errBodyClose) {
+			t.Fatalf("expected retained body close error, got %v", err)
+		}
+		if retainedErr == nil {
+			retainedErr = err
+			continue
+		}
+		if err != retainedErr {
+			t.Fatalf("expected identical retained error, got distinct values %p and %p", retainedErr, err)
+		}
+	}
+	<-readDone
+	if got := body.closeCount.Load(); got != 1 {
+		t.Fatalf("expected one body close, got %d", got)
+	}
+	if got := conn.closeCount.Load(); got != 1 {
+		t.Fatalf("expected one connection close, got %d", got)
+	}
+}
+
+type requestWrapperTestBody struct {
+	readErr    error
+	closeErr   error
+	closeCount atomic.Int32
+}
+
+func (b *requestWrapperTestBody) Read([]byte) (int, error) {
+	if b.readErr != nil {
+		return 0, b.readErr
+	}
+	return 0, io.EOF
+}
+
+func (b *requestWrapperTestBody) Close() error {
+	b.closeCount.Add(1)
+	return b.closeErr
+}
+
+type requestWrapperTestConn struct {
+	net.Conn
+	closeErr    error
+	closeCount  atomic.Int32
+	closeOnce   sync.Once
+	readStarted chan struct{}
+	readDone    chan struct{}
+}
+
+func newRequestWrapperTestConn(closeErr error) *requestWrapperTestConn {
+	return &requestWrapperTestConn{
+		closeErr:    closeErr,
+		readStarted: make(chan struct{}),
+		readDone:    make(chan struct{}),
+	}
+}
+
+func (c *requestWrapperTestConn) Read([]byte) (int, error) {
+	select {
+	case <-c.readStarted:
+	default:
+		close(c.readStarted)
+	}
+	<-c.readDone
+	return 0, io.ErrClosedPipe
+}
+
+func (c *requestWrapperTestConn) Write(p []byte) (int, error) {
+	return len(p), nil
+}
+
+func (c *requestWrapperTestConn) Close() error {
+	c.closeOnce.Do(func() {
+		c.closeCount.Add(1)
+		close(c.readDone)
+	})
+	return c.closeErr
+}
+
+type requestWrapperTestResponseWriter struct {
+	conn net.Conn
+}
+
+func (w *requestWrapperTestResponseWriter) Header() http.Header {
+	return make(http.Header)
+}
+
+func (w *requestWrapperTestResponseWriter) Write([]byte) (int, error) {
+	return 0, nil
+}
+
+func (w *requestWrapperTestResponseWriter) WriteHeader(int) {}
+
+func (w *requestWrapperTestResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return w.conn, bufio.NewReadWriter(bufio.NewReader(bytes.NewReader(nil)), bufio.NewWriter(io.Discard)), nil
+}

service/rpc/geoip_rpc.go

@@ -0,0 +1,57 @@
+package rpc
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"net"
+
+	"github.com/nezhahq/nezha/model"
+	geoipx "github.com/nezhahq/nezha/pkg/geoip"
+	pb "github.com/nezhahq/nezha/proto"
+	"github.com/nezhahq/nezha/service/singleton"
+)
+
+func (s *NezhaHandler) ReportGeoIP(ctx context.Context, report *pb.GeoIP) (*pb.GeoIP, error) {
+	clientID, err := s.Auth.Check(ctx)
+	if err != nil {
+		return nil, err
+	}
+	geoIP := model.PB2GeoIP(report)
+	if geoIP.IP.IPv4Addr == "" && geoIP.IP.IPv6Addr == "" {
+		ip, _ := ctx.Value(model.CtxKeyRealIP{}).(string)
+		if ip == "" {
+			ip, _ = ctx.Value(model.CtxKeyConnectingIP{}).(string)
+		}
+		geoIP.IP.IPv4Addr = ip
+	}
+	joinedIP := geoIP.IP.Join()
+	server, ok := singleton.ServerShared.Get(clientID)
+	if !ok || server == nil {
+		return nil, fmt.Errorf("server not found")
+	}
+	if server.EnableDDNS && joinedIP != "" && (server.GeoIP == nil || server.GeoIP.IP != geoIP.IP) {
+		if err := singleton.ServerShared.UpdateDDNS(server, &model.IP{IPv4Addr: geoIP.IP.IPv4Addr, IPv6Addr: geoIP.IP.IPv6Addr}); err != nil {
+			log.Printf("NEZHA>> Failed to update DDNS for server %d: %v", server.ID, err)
+		}
+	}
+	if server.GeoIP != nil && singleton.Conf.EnableIPChangeNotification &&
+		((singleton.Conf.Cover == model.ConfigCoverAll && !singleton.Conf.IgnoredIPNotificationServerIDs[clientID]) ||
+			(singleton.Conf.Cover == model.ConfigCoverIgnoreAll && singleton.Conf.IgnoredIPNotificationServerIDs[clientID])) &&
+		server.GeoIP.IP.Join() != "" && joinedIP != "" && server.GeoIP.IP != geoIP.IP {
+		singleton.NotificationShared.SendNotification(singleton.Conf.IPChangeNotificationGroupID,
+			fmt.Sprintf("[%s] %s, %s => %s", singleton.Localizer.T("IP Changed"), server.Name,
+				singleton.IPDesensitize(server.GeoIP.IP.Join()), singleton.IPDesensitize(joinedIP)), "")
+	}
+	ip := geoIP.IP.IPv4Addr
+	if geoIP.IP.IPv6Addr != "" && (report.GetUse6() || ip == "") {
+		ip = geoIP.IP.IPv6Addr
+	}
+	location, err := geoipx.Lookup(net.ParseIP(ip))
+	if err != nil {
+		log.Printf("NEZHA>> geoip.Lookup: %v", err)
+	}
+	geoIP.CountryCode = location
+	server.GeoIP = &geoIP
+	return &pb.GeoIP{Ip: nil, CountryCode: location, DashboardBootTime: singleton.DashboardBootTime}, nil
+}

service/rpc/io_stream.go

@@ -1,7 +1,6 @@
 package rpc
 
 import (
-	"context"
 	"errors"
 	"io"
 	"sync"
@@ -10,16 +9,14 @@ import (
 	"github.com/nezhahq/nezha/service/singleton"
 )
 
-// StreamPurpose tags every IOStream with the feature that opened it so
-// admin actions can drop only the relevant subset. Existing call sites
-// (terminal / fm / NAT / server-transfer) keep PurposeLegacy and the
-// previous semantics; only the new MCP fs.transfer path uses
-// PurposeMCPTransfer, which is what EnableMCP=false revokes.
 type StreamPurpose uint8
 
 const (
 	PurposeLegacy StreamPurpose = iota
 	PurposeMCPTransfer
+	PurposeTerminal
+	PurposeFileManager
+	PurposeNAT
 )
 
 type ioStreamContext struct {
@@ -34,315 +31,30 @@ type ioStreamContext struct {
 	agentIoChOnce    sync.Once
 	revokedCh        chan struct{}
 	revokedOnce      sync.Once
+	waitStartedCh    chan struct{}
+	waitStartedOnce  sync.Once
+	startCaptureCh   chan struct{}
+	startCaptureOnce sync.Once
 }
 
-type bp struct {
-	buf []byte
-}
-
-var bufPool = sync.Pool{
-	New: func() any {
-		return &bp{
-			buf: make([]byte, 1024*1024),
-		}
-	},
-}
-
-const (
-	maxStreamsPerUser   = 20
-	maxStreamsPerServer = 40
-)
-
-var (
-	ErrTooManyStreamsForUser   = errors.New("too many concurrent streams for this user")
-	ErrTooManyStreamsForServer = errors.New("too many concurrent streams for this server")
-)
-
-func (s *NezhaHandler) CreateStream(streamId string, creatorUserID uint64, targetServerID uint64) error {
-	return s.CreateStreamWithPurpose(streamId, creatorUserID, targetServerID, PurposeLegacy)
-}
-
-func (s *NezhaHandler) CreateStreamWithPurpose(streamId string, creatorUserID uint64, targetServerID uint64, purpose StreamPurpose) error {
-	s.ioStreamMutex.Lock()
-	defer s.ioStreamMutex.Unlock()
-
-	var perUser, perServer int
-	for _, ctx := range s.ioStreams {
-		if creatorUserID != 0 && ctx.creatorUserID == creatorUserID {
-			perUser++
-		}
-		if ctx.targetServerID == targetServerID {
-			perServer++
-		}
-	}
-	// creatorUserID==0 is a dashboard-internal stream (NAT, server transfer,
-	// MCP transfer); only end-user-initiated streams are capped per user, but
-	// every stream counts toward the per-server cap so one server cannot be
-	// flooded regardless of who opened the streams.
-	if creatorUserID != 0 && perUser >= maxStreamsPerUser {
-		return ErrTooManyStreamsForUser
-	}
-	if perServer >= maxStreamsPerServer {
-		return ErrTooManyStreamsForServer
-	}
-
-	s.ioStreams[streamId] = &ioStreamContext{
-		creatorUserID:    creatorUserID,
-		targetServerID:   targetServerID,
-		purpose:          purpose,
-		userIoConnectCh:  make(chan struct{}),
-		agentIoConnectCh: make(chan struct{}),
-		revokedCh:        make(chan struct{}),
-	}
-	return nil
-}
-
-// IsStreamAuthorizedForAgent reports whether the connecting agent is the
-// server the dashboard selected when CreateStream was called. Without this
-// check any authenticated agent that learns an active streamId — via
-// task-stream observation, leaked logs, or a shared global agent secret —
-// can race in via IOStream() and serve a terminal / fm / NAT session that
-// was addressed to a different server, turning the channel into a
-// session-hijack RCE primitive. This is the agent-side dual of
-// IsStreamAuthorizedForUser.
-func (s *NezhaHandler) IsStreamAuthorizedForAgent(streamId string, agentServerID uint64) bool {
-	s.ioStreamMutex.RLock()
-	defer s.ioStreamMutex.RUnlock()
-
-	ctx, ok := s.ioStreams[streamId]
-	if !ok {
-		return false
-	}
-	return ctx.targetServerID != 0 && ctx.targetServerID == agentServerID
-}
-
-// WaitForAgent 阻塞等待 agent 端通过 IOStream 接入并完成 AgentConnected。
-// dashboard 把 MCP 大文件传输的 task 派给 agent 后，需要等 agent dial 回来
-// 才能开始 Read/Write，这里以 timeout 内的轻量轮询暴露给 controller。
-//
-// 同时返回 agent 端流（io.ReadWriteCloser）以便 controller 调 io.CopyN 转发
-// HTTP body；ok=false 表示超时或流已被关闭。
-func (s *NezhaHandler) WaitForAgent(ctx context.Context, streamId string, timeout time.Duration) (io.ReadWriteCloser, bool) {
-	deadline := time.NewTimer(timeout)
-	defer deadline.Stop()
-	for {
-		s.ioStreamMutex.RLock()
-		sc, ok := s.ioStreams[streamId]
-		if ok && sc.agentIo != nil {
-			s.ioStreamMutex.RUnlock()
-			return sc.agentIo, true
-		}
-		s.ioStreamMutex.RUnlock()
-		if !ok {
-			return nil, false
-		}
-		select {
-		case <-ctx.Done():
-			return nil, false
-		case <-deadline.C:
-			return nil, false
-		case <-sc.revokedCh:
-			return nil, false
-		case <-sc.agentIoConnectCh:
-			s.ioStreamMutex.RLock()
-			ag := sc.agentIo
-			s.ioStreamMutex.RUnlock()
-			return ag, ag != nil
-		}
+func newIOStreamContext(creatorUserID, targetServerID uint64, purpose StreamPurpose) *ioStreamContext {
+	return &ioStreamContext{
+		creatorUserID: creatorUserID, targetServerID: targetServerID, purpose: purpose,
+		userIoConnectCh: make(chan struct{}), agentIoConnectCh: make(chan struct{}),
+		revokedCh: make(chan struct{}), waitStartedCh: make(chan struct{}), startCaptureCh: make(chan struct{}),
 	}
 }
 
-// IsStreamAuthorizedForUser checks whether the requesting user may attach to
-// the stream. A stream is reachable only by its creator or by an admin; any
-// other authenticated user must be rejected. Unknown streams are always
-// rejected.
-func (s *NezhaHandler) IsStreamAuthorizedForUser(streamId string, userID uint64, isAdmin bool) bool {
-	creator, found := s.StreamOwnership(streamId)
-	if !found {
-		return false
-	}
-	if isAdmin {
-		return true
-	}
-	return creator == userID
+func (stream *ioStreamContext) revoke() {
+	stream.revokedOnce.Do(func() { close(stream.revokedCh) })
 }
 
-// isValidIOStreamMagic reports whether the first four bytes of an IOStream
-// init message carry the ff05ff05 marker. Previously this was inlined as
-// `byte0 != 0xff && byte1 != 0x05 && byte2 != 0xff && byte3 == 0x05` to
-// detect *invalid* payloads — but && short-circuited so any payload whose
-// byte0 happened to be 0xff slipped through. Centralising the check here and
-// stating the contract positively (all four bytes must match) eliminates the
-// short-circuit class of mistakes.
+type bp struct{ buf []byte }
+
+var bufPool = sync.Pool{New: func() any { return &bp{buf: make([]byte, 1024*1024)} }}
+
 func isValidIOStreamMagic(data []byte) bool {
-	if len(data) < 4 {
-		return false
-	}
-	return data[0] == 0xff && data[1] == 0x05 && data[2] == 0xff && data[3] == 0x05
-}
-
-// StreamOwnership returns the user ID that created the stream and whether the
-// stream is still tracked. Callers must compare the returned creator against
-// the requesting user before attaching to the stream — without this the
-// channel becomes a session-hijack primitive (terminal/file manager RCE).
-func (s *NezhaHandler) StreamOwnership(streamId string) (uint64, bool) {
-	s.ioStreamMutex.RLock()
-	defer s.ioStreamMutex.RUnlock()
-
-	ctx, ok := s.ioStreams[streamId]
-	if !ok {
-		return 0, false
-	}
-	return ctx.creatorUserID, true
-}
-
-// StreamTarget returns the server ID the stream was opened against and
-// whether the stream is still tracked. Callers MUST pass this through the
-// requesting PAT's CanAccessServer check before allowing attachment —
-// IsStreamAuthorizedForUser only knows about creator/admin, so without this
-// dual gate an admin's server-limited PAT can hijack any stream by knowing
-// the streamId.
-func (s *NezhaHandler) StreamTarget(streamId string) (uint64, bool) {
-	s.ioStreamMutex.RLock()
-	defer s.ioStreamMutex.RUnlock()
-
-	ctx, ok := s.ioStreams[streamId]
-	if !ok {
-		return 0, false
-	}
-	return ctx.targetServerID, true
-}
-
-func (s *NezhaHandler) GetStream(streamId string) (*ioStreamContext, error) {
-	s.ioStreamMutex.RLock()
-	defer s.ioStreamMutex.RUnlock()
-
-	if ctx, ok := s.ioStreams[streamId]; ok {
-		return ctx, nil
-	}
-
-	return nil, errors.New("stream not found")
-}
-
-// RevokeStreamsForServer tears down every IOStream whose targetServerID
-// matches serverID. Called by the singleton package via the
-// ServerTransferStreamRevocationHook on every transfer ownership
-// transition — a stream the previous owner had open against this server
-// must not survive into the new tenant, otherwise terminal/file-manager/NAT
-// sessions become post-transfer hijack channels (effectively RCE).
-//
-// Underlying IO pipes are closed inline so the dashboard websocket loop
-// sees EOF immediately rather than at the next idle-timeout.
-func (s *NezhaHandler) RevokeStreamsForServer(serverID uint64) {
-	if serverID == 0 {
-		return
-	}
-	s.ioStreamMutex.Lock()
-	defer s.ioStreamMutex.Unlock()
-	for streamId, ctx := range s.ioStreams {
-		if ctx.targetServerID != serverID {
-			continue
-		}
-		if ctx.userIo != nil {
-			ctx.userIo.Close()
-		}
-		if ctx.agentIo != nil {
-			ctx.agentIo.Close()
-		}
-		delete(s.ioStreams, streamId)
-	}
-}
-
-// RevokeStreamsForPurpose tears down every IOStream tagged with the given
-// purpose. Used as the IOStream half of the MCP kill switch: when the
-// admin flips EnableMCP=false, any in-flight fs.transfer / fs.upload /
-// fs.download must drop immediately rather than wait out the 5min IO
-// timeout. Returns the number of streams revoked so the caller can log
-// the blast radius.
-func (s *NezhaHandler) RevokeStreamsForPurpose(purpose StreamPurpose) int {
-	s.ioStreamMutex.Lock()
-	defer s.ioStreamMutex.Unlock()
-	revoked := 0
-	for streamId, ctx := range s.ioStreams {
-		if ctx.purpose != purpose {
-			continue
-		}
-		ctx.revokedOnce.Do(func() {
-			if ctx.revokedCh != nil {
-				close(ctx.revokedCh)
-			}
-		})
-		if ctx.userIo != nil {
-			ctx.userIo.Close()
-		}
-		if ctx.agentIo != nil {
-			ctx.agentIo.Close()
-		}
-		delete(s.ioStreams, streamId)
-		revoked++
-	}
-	return revoked
-}
-
-func (s *NezhaHandler) CloseStream(streamId string) error {
-	s.ioStreamMutex.Lock()
-	defer s.ioStreamMutex.Unlock()
-
-	if ctx, ok := s.ioStreams[streamId]; ok {
-		if ctx.userIo != nil {
-			ctx.userIo.Close()
-		}
-		if ctx.agentIo != nil {
-			ctx.agentIo.Close()
-		}
-		delete(s.ioStreams, streamId)
-	}
-
-	return nil
-}
-
-// UserConnected publishes the user-side IO under ioStreamMutex so concurrent
-// Revoke* / WaitForAgent / StartStream see a consistent stream view.
-// Without the lock, the bare assignment to stream.userIo races with the
-// revoker's lock-protected read and triggers go-race.
-func (s *NezhaHandler) UserConnected(streamId string, userIo io.ReadWriteCloser) error {
-	s.ioStreamMutex.Lock()
-	stream, ok := s.ioStreams[streamId]
-	if !ok {
-		s.ioStreamMutex.Unlock()
-		return errors.New("stream not found")
-	}
-	stream.userIo = userIo
-	s.ioStreamMutex.Unlock()
-	stream.userIoChOnce.Do(func() {
-		close(stream.userIoConnectCh)
-	})
-	return nil
-}
-
-// AgentConnected is the agent-side dual of UserConnected. Same locking
-// rationale.
-func (s *NezhaHandler) AgentConnected(streamId string, agentIo io.ReadWriteCloser) error {
-	s.ioStreamMutex.Lock()
-	stream, ok := s.ioStreams[streamId]
-	if !ok {
-		s.ioStreamMutex.Unlock()
-		return errors.New("stream not found")
-	}
-	stream.agentIo = agentIo
-	s.ioStreamMutex.Unlock()
-	stream.agentIoChOnce.Do(func() {
-		close(stream.agentIoConnectCh)
-	})
-	return nil
-}
-
-// streamEndpoints returns the user/agent IO under ioStreamMutex so callers
-// never read the interface fields while UserConnected/AgentConnected write them.
-func (s *NezhaHandler) streamEndpoints(stream *ioStreamContext) (userIo, agentIo io.ReadWriteCloser) {
-	s.ioStreamMutex.RLock()
-	defer s.ioStreamMutex.RUnlock()
-	return stream.userIo, stream.agentIo
+	return len(data) >= 4 && data[0] == 0xff && data[1] == 0x05 && data[2] == 0xff && data[3] == 0x05
 }
 
 func (s *NezhaHandler) StartStream(streamId string, timeout time.Duration) error {
@@ -350,52 +62,62 @@ func (s *NezhaHandler) StartStream(streamId string, timeout time.Duration) error
 	if err != nil {
 		return err
 	}
+	return s.startStreamContext(streamId, stream, timeout)
+}
 
+func (s *NezhaHandler) startStreamContext(streamId string, stream *ioStreamContext, timeout time.Duration) error {
 	timeoutTimer := time.NewTimer(timeout)
 	defer timeoutTimer.Stop()
-
-LOOP:
+	userConnected := stream.userIoConnectCh
+	agentConnected := stream.agentIoConnectCh
 	for {
+		s.ioStreamMutex.RLock()
+		if current, exists := s.ioStreams[streamId]; !exists || current != stream {
+			s.ioStreamMutex.RUnlock()
+			return errors.New("stream revoked")
+		}
+		userIo, agentIo := stream.userIo, stream.agentIo
+		s.ioStreamMutex.RUnlock()
+		stream.startCaptureOnce.Do(func() { close(stream.startCaptureCh) })
+		if userIo != nil {
+			userConnected = nil
+		}
+		if agentIo != nil {
+			agentConnected = nil
+		}
+		if userIo != nil && agentIo != nil {
+			break
+		}
 		select {
-		case <-stream.userIoConnectCh:
-			if _, agentIo := s.streamEndpoints(stream); agentIo != nil {
-				break LOOP
-			}
-		case <-stream.agentIoConnectCh:
-			if userIo, _ := s.streamEndpoints(stream); userIo != nil {
-				break LOOP
-			}
+		case <-userConnected:
+			userConnected = nil
+		case <-agentConnected:
+			agentConnected = nil
+		case <-stream.revokedCh:
+			return errors.New("stream revoked")
 		case <-timeoutTimer.C:
-			break LOOP
+			return singleton.Localizer.ErrorT("timeout: stream endpoints not established")
 		}
-		time.Sleep(time.Millisecond * 500)
 	}
-
-	userIo, agentIo := s.streamEndpoints(stream)
-	if userIo == nil && agentIo == nil {
-		return singleton.Localizer.ErrorT("timeout: no connection established")
+	s.ioStreamMutex.RLock()
+	if current, exists := s.ioStreams[streamId]; !exists || current != stream {
+		s.ioStreamMutex.RUnlock()
+		return errors.New("stream revoked")
 	}
-	if userIo == nil {
-		return singleton.Localizer.ErrorT("timeout: user connection not established")
-	}
-	if agentIo == nil {
-		return singleton.Localizer.ErrorT("timeout: agent connection not established")
-	}
-
+	userIo, agentIo := stream.userIo, stream.agentIo
+	s.ioStreamMutex.RUnlock()
 	errCh := make(chan error, 2)
-
 	go func() {
 		bp := bufPool.Get().(*bp)
 		defer bufPool.Put(bp)
-		_, innerErr := io.CopyBuffer(userIo, agentIo, bp.buf)
-		errCh <- innerErr
+		_, copyErr := io.CopyBuffer(userIo, agentIo, bp.buf)
+		errCh <- copyErr
 	}()
 	go func() {
 		bp := bufPool.Get().(*bp)
 		defer bufPool.Put(bp)
-		_, innerErr := io.CopyBuffer(agentIo, userIo, bp.buf)
-		errCh <- innerErr
+		_, copyErr := io.CopyBuffer(agentIo, userIo, bp.buf)
+		errCh <- copyErr
 	}()
-
 	return <-errCh
 }

service/rpc/io_stream_lifecycle.go

@@ -0,0 +1,128 @@
+package rpc
+
+import (
+	"context"
+	"errors"
+	"io"
+	"time"
+)
+
+type ioStreamDetach struct {
+	stream    *ioStreamContext
+	endpoints []io.ReadWriteCloser
+}
+
+func detachStreamLocked(streamID string, retainedStream *ioStreamContext, streams map[string]*ioStreamContext) (ioStreamDetach, bool) {
+	current, live := streams[streamID]
+	if streamID == "" || !live || current != retainedStream {
+		return ioStreamDetach{}, false
+	}
+	retainedStream.revoke()
+	endpoints := make([]io.ReadWriteCloser, 0, 2)
+	if retainedStream.userIo != nil {
+		endpoints = append(endpoints, retainedStream.userIo)
+	}
+	if retainedStream.agentIo != nil {
+		endpoints = append(endpoints, retainedStream.agentIo)
+	}
+	delete(streams, streamID)
+	return ioStreamDetach{stream: retainedStream, endpoints: endpoints}, true
+}
+
+func (s *NezhaHandler) detachExactStream(streamID string, retainedStream *ioStreamContext) error {
+	s.ioStreamMutex.Lock()
+	detached, ok := detachStreamLocked(streamID, retainedStream, s.ioStreams)
+	if !ok {
+		s.ioStreamMutex.Unlock()
+		return nil
+	}
+	s.publishIOStreamStateChangeLocked()
+	s.ioStreamMutex.Unlock()
+
+	var closeErrors []error
+	for _, endpoint := range detached.endpoints {
+		if err := endpoint.Close(); err != nil {
+			closeErrors = append(closeErrors, err)
+		}
+	}
+	return errors.Join(closeErrors...)
+}
+
+func (s *NezhaHandler) detachStreams(shouldDetach func(*ioStreamContext) bool) (int, error) {
+	s.ioStreamMutex.Lock()
+	detached := make([]ioStreamDetach, 0)
+	for streamID, stream := range s.ioStreams {
+		if !shouldDetach(stream) {
+			continue
+		}
+		item, ok := detachStreamLocked(streamID, stream, s.ioStreams)
+		if ok {
+			detached = append(detached, item)
+		}
+	}
+	if len(detached) > 0 {
+		s.publishIOStreamStateChangeLocked()
+	}
+	s.ioStreamMutex.Unlock()
+
+	var closeErrors []error
+	for _, item := range detached {
+		// Registry publication must precede endpoint Close so Close implementations may reenter safely.
+		for _, endpoint := range item.endpoints {
+			if err := endpoint.Close(); err != nil {
+				closeErrors = append(closeErrors, err)
+			}
+		}
+	}
+	return len(detached), errors.Join(closeErrors...)
+}
+
+func (s *NezhaHandler) CloseStream(streamID string) error {
+	_, err := s.detachStreams(func(stream *ioStreamContext) bool {
+		return stream != nil && streamID != "" && stream == s.ioStreams[streamID]
+	})
+	return err
+}
+
+func (s *NezhaHandler) WaitForAgent(ctx context.Context, streamID string, timeout time.Duration) (io.ReadWriteCloser, bool) {
+	deadline := time.NewTimer(timeout)
+	defer deadline.Stop()
+	for {
+		s.ioStreamMutex.RLock()
+		stream, ok := s.ioStreams[streamID]
+		if ok && stream.agentIo != nil {
+			agentIo := stream.agentIo
+			s.ioStreamMutex.RUnlock()
+			return agentIo, true
+		}
+		if !ok {
+			s.ioStreamMutex.RUnlock()
+			return nil, false
+		}
+		revokedCh := stream.revokedCh
+		agentIoConnectCh := stream.agentIoConnectCh
+		stream.waitStartedOnce.Do(func() { close(stream.waitStartedCh) })
+		s.ioStreamMutex.RUnlock()
+		select {
+		case <-ctx.Done():
+			return nil, false
+		case <-deadline.C:
+			return nil, false
+		case <-revokedCh:
+			return nil, false
+		case <-agentIoConnectCh:
+		}
+	}
+}
+
+func (s *NezhaHandler) RevokeStreamsForServer(serverID uint64) {
+	if serverID == 0 {
+		return
+	}
+	_, _ = s.detachStreams(func(stream *ioStreamContext) bool { return stream.targetServerID == serverID })
+}
+
+func (s *NezhaHandler) RevokeStreamsForPurpose(purpose StreamPurpose) int {
+	revoked, _ := s.detachStreams(func(stream *ioStreamContext) bool { return stream.purpose == purpose })
+	return revoked
+}

service/rpc/io_stream_lifecycle_test.go

@@ -0,0 +1,259 @@
+package rpc
+
+import (
+	"context"
+	"errors"
+	"io"
+	"testing"
+	"time"
+)
+
+type lifecycleRWC struct {
+	closed chan struct{}
+}
+
+type reenteringErrorRWC struct {
+	handler  *NezhaHandler
+	streamID string
+	err      error
+}
+
+func (stream *reenteringErrorRWC) Read([]byte) (int, error)       { return 0, io.EOF }
+func (stream *reenteringErrorRWC) Write(data []byte) (int, error) { return len(data), nil }
+func (stream *reenteringErrorRWC) Close() error {
+	if _, ok := stream.handler.StreamOwnership(stream.streamID); ok {
+		return errors.Join(stream.err, errors.New("stream remained registered during endpoint close"))
+	}
+	return stream.err
+}
+
+func newLifecycleRWC() *lifecycleRWC {
+	return &lifecycleRWC{closed: make(chan struct{})}
+}
+
+func (stream *lifecycleRWC) Read([]byte) (int, error)       { return 0, io.EOF }
+func (stream *lifecycleRWC) Write(data []byte) (int, error) { return len(data), nil }
+func (stream *lifecycleRWC) Close() error {
+	select {
+	case <-stream.closed:
+	default:
+		close(stream.closed)
+	}
+	return nil
+}
+
+func TestIOStreamValidCreateAttachCloseLifecycle(t *testing.T) {
+	handler := NewNezhaHandler()
+	user := newLifecycleRWC()
+	agent := newLifecycleRWC()
+	if err := handler.CreateStream("valid-lifecycle", 11, 22); err != nil {
+		t.Fatalf("Given a new stream, CreateStream failed: %v", err)
+	}
+	if err := handler.UserConnected("valid-lifecycle", user); err != nil {
+		t.Fatalf("Given a tracked stream, UserConnected failed: %v", err)
+	}
+	if err := handler.AgentConnected("valid-lifecycle", agent); err != nil {
+		t.Fatalf("Given a tracked stream, AgentConnected failed: %v", err)
+	}
+	if _, ok := handler.StreamOwnership("valid-lifecycle"); !ok {
+		t.Fatal("Then a valid attached stream must remain tracked")
+	}
+	if err := handler.CloseStream("valid-lifecycle"); err != nil {
+		t.Fatalf("When closing the valid stream, CloseStream failed: %v", err)
+	}
+	if _, ok := handler.StreamOwnership("valid-lifecycle"); ok {
+		t.Fatal("Then CloseStream must remove the tracked stream")
+	}
+}
+
+func TestCreateStreamKeepsExistingStreamWhenIDIsDuplicated(t *testing.T) {
+	handler := NewNezhaHandler()
+	original := newLifecycleRWC()
+	if err := handler.CreateStream("duplicate-id", 11, 22); err != nil {
+		t.Fatalf("Given a new stream ID, CreateStream failed: %v", err)
+	}
+	if err := handler.AgentConnected("duplicate-id", original); err != nil {
+		t.Fatalf("Given a live stream, AgentConnected failed: %v", err)
+	}
+
+	err := handler.CreateStream("duplicate-id", 33, 44)
+	if !errors.Is(err, ErrStreamAlreadyExists) {
+		t.Fatalf("When reusing a live ID, expected ErrStreamAlreadyExists, got %v", err)
+	}
+	owner, found := handler.StreamOwnership("duplicate-id")
+	if !found || owner != 11 {
+		t.Fatalf("Then the original stream ownership must remain, found=%v owner=%d", found, owner)
+	}
+	select {
+	case <-original.closed:
+		t.Fatal("Then duplicate creation must not close the original endpoint")
+	default:
+	}
+}
+
+func TestAgentConnectedRejectsDuplicateEndpointWithoutReplacingLiveRelay(t *testing.T) {
+	handler := NewNezhaHandler()
+	first := newLifecycleRWC()
+	second := newLifecycleRWC()
+	if err := handler.CreateStream("agent-once", 11, 22); err != nil {
+		t.Fatalf("Given a new stream, CreateStream failed: %v", err)
+	}
+	if err := handler.AgentConnected("agent-once", first); err != nil {
+		t.Fatalf("Given no agent endpoint, AgentConnected failed: %v", err)
+	}
+	if err := handler.AgentConnected("agent-once", second); !errors.Is(err, ErrAgentStreamAlreadyConnected) {
+		t.Fatalf("When attaching a second agent endpoint, expected ErrAgentStreamAlreadyConnected, got %v", err)
+	}
+	endpoints, err := handler.GetStream("agent-once")
+	if err != nil || endpoints.agentIo != first {
+		t.Fatalf("Then the first endpoint must remain attached, err=%v", err)
+	}
+	select {
+	case <-second.closed:
+	default:
+		t.Fatal("Then the rejected duplicate endpoint must be closed")
+	}
+}
+
+func TestCloseStreamWakesWaitForAgentAndAllowsSlotReuse(t *testing.T) {
+	handler := NewNezhaHandler()
+	if err := handler.CreateStream("wait-close", 11, 22); err != nil {
+		t.Fatalf("Given a pending stream, CreateStream failed: %v", err)
+	}
+	stream, err := handler.GetStream("wait-close")
+	if err != nil {
+		t.Fatalf("Given a created stream, GetStream failed: %v", err)
+	}
+	result := make(chan bool, 1)
+	go func() {
+		_, ok := handler.WaitForAgent(context.Background(), "wait-close", time.Minute)
+		result <- ok
+	}()
+	select {
+	case <-stream.waitStartedCh:
+	case <-time.After(time.Second):
+		t.Fatal("WaitForAgent did not enter its blocking select")
+	}
+
+	if err := handler.CloseStream("wait-close"); err != nil {
+		t.Fatalf("When closing a pending stream, CloseStream failed: %v", err)
+	}
+	select {
+	case ok := <-result:
+		if ok {
+			t.Fatal("Then WaitForAgent must report no attached agent")
+		}
+	case <-time.After(time.Second):
+		t.Fatal("Then CloseStream must wake WaitForAgent")
+	}
+	if err := handler.CreateStream("wait-close-reused", 11, 22); err != nil {
+		t.Fatalf("Then the released user/server slot must be reusable: %v", err)
+	}
+}
+
+func TestRevokeStreamsForPurposeWakesWaitForAgentAndIsRepeatable(t *testing.T) {
+	handler := NewNezhaHandler()
+	if err := handler.CreateStreamWithPurpose("revoke-wait", 0, 22, PurposeMCPTransfer); err != nil {
+		t.Fatalf("Given a pending MCP stream, CreateStream failed: %v", err)
+	}
+	stream, err := handler.GetStream("revoke-wait")
+	if err != nil {
+		t.Fatalf("Given a created stream, GetStream failed: %v", err)
+	}
+	result := make(chan bool, 1)
+	go func() {
+		_, ok := handler.WaitForAgent(context.Background(), "revoke-wait", time.Minute)
+		result <- ok
+	}()
+	select {
+	case <-stream.waitStartedCh:
+	case <-time.After(time.Second):
+		t.Fatal("WaitForAgent did not enter its blocking select")
+	}
+
+	if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 1 {
+		t.Fatalf("When revoking the purpose, expected one stream, got %d", revoked)
+	}
+	if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 0 {
+		t.Fatalf("When repeating revocation, expected zero streams, got %d", revoked)
+	}
+	select {
+	case ok := <-result:
+		if ok {
+			t.Fatal("Then WaitForAgent must report no attached agent")
+		}
+	case <-time.After(time.Second):
+		t.Fatal("Then revocation must wake WaitForAgent")
+	}
+}
+
+func TestCloseStreamDetachesBeforeReenteringEndpointCloseAndJoinsErrors(t *testing.T) {
+	handler := NewNezhaHandler()
+	if err := handler.CreateStream("close-errors", 1, 1); err != nil {
+		t.Fatal(err)
+	}
+	firstErr := errors.New("first close error")
+	secondErr := errors.New("second close error")
+	if err := handler.UserConnected("close-errors", &reenteringErrorRWC{handler: handler, streamID: "close-errors", err: firstErr}); err != nil {
+		t.Fatal(err)
+	}
+	if err := handler.AgentConnected("close-errors", &reenteringErrorRWC{handler: handler, streamID: "close-errors", err: secondErr}); err != nil {
+		t.Fatal(err)
+	}
+	err := handler.CloseStream("close-errors")
+	if !errors.Is(err, firstErr) || !errors.Is(err, secondErr) {
+		t.Fatalf("close errors were not joined: %v", err)
+	}
+}
+
+func TestStartStreamReturnsImmediatelyWhenRevoked(t *testing.T) {
+	handler := NewNezhaHandler()
+	if err := handler.CreateStream("start-revoked", 1, 1); err != nil {
+		t.Fatal(err)
+	}
+	result := make(chan error, 1)
+	go func() { result <- handler.StartStream("start-revoked", time.Minute) }()
+	if revoked := handler.RevokeStreamsForPurpose(PurposeLegacy); revoked != 1 {
+		t.Fatalf("revoked streams: %d", revoked)
+	}
+	select {
+	case err := <-result:
+		if err == nil {
+			t.Fatal("revoked StartStream must return an error")
+		}
+	case <-time.After(time.Second):
+		t.Fatal("StartStream did not wake on revoke")
+	}
+}
+
+func TestConcurrentCloseAndRevokePublishOneGeneration(t *testing.T) {
+	handler := NewNezhaHandler()
+	if err := handler.CreateStream("single-generation", 1, 1); err != nil {
+		t.Fatal(err)
+	}
+	start := handler.SnapshotIOStreamState()
+	closeDone := make(chan struct{})
+	revokeDone := make(chan struct{})
+	go func() {
+		_ = handler.CloseStream("single-generation")
+		close(closeDone)
+	}()
+	go func() {
+		handler.RevokeStreamsForPurpose(PurposeLegacy)
+		close(revokeDone)
+	}()
+	select {
+	case <-closeDone:
+	case <-time.After(time.Second):
+		t.Fatal("CloseStream did not complete")
+	}
+	select {
+	case <-revokeDone:
+	case <-time.After(time.Second):
+		t.Fatal("RevokeStreamsForPurpose did not complete")
+	}
+	state := handler.SnapshotIOStreamState()
+	if state.Count != 0 || state.Generation != start.Generation+1 {
+		t.Fatalf("single detach publication: start=%+v final=%+v", start, state)
+	}
+}

service/rpc/io_stream_quota.go

@@ -0,0 +1,62 @@
+package rpc
+
+import "errors"
+
+const (
+	maxStreamsPerUser   = 20
+	maxStreamsPerServer = 40
+)
+
+var (
+	ErrTooManyStreamsForUser   = errors.New("too many concurrent streams for this user")
+	ErrTooManyStreamsForServer = errors.New("too many concurrent streams for this server")
+	ErrStreamAlreadyExists     = errors.New("stream already exists")
+)
+
+func (s *NezhaHandler) CreateStream(streamId string, creatorUserID uint64, targetServerID uint64) error {
+	return s.CreateStreamWithPurpose(streamId, creatorUserID, targetServerID, PurposeLegacy)
+}
+
+func (s *NezhaHandler) CreateStreamWithPurpose(streamId string, creatorUserID uint64, targetServerID uint64, purpose StreamPurpose) error {
+	s.ioStreamMutex.Lock()
+	defer s.ioStreamMutex.Unlock()
+	return s.createStreamLocked(streamId, creatorUserID, targetServerID, purpose)
+}
+
+func (s *NezhaHandler) createStreamLocked(streamId string, creatorUserID uint64, targetServerID uint64, purpose StreamPurpose) error {
+	if _, exists := s.ioStreams[streamId]; exists {
+		// Stream IDs identify live relay ownership; never overwrite one or orphan its endpoint.
+		return ErrStreamAlreadyExists
+	}
+
+	var perUser, perServer int
+	for _, ctx := range s.ioStreams {
+		if creatorUserID != 0 && ctx.creatorUserID == creatorUserID {
+			perUser++
+		}
+		if ctx.targetServerID == targetServerID {
+			perServer++
+		}
+	}
+	// creatorUserID==0 is a dashboard-internal stream (NAT, server transfer,
+	// MCP transfer); only end-user-initiated streams are capped per user, but
+	// every stream counts toward the per-server cap so one server cannot be
+	// flooded regardless of who opened the streams.
+	if creatorUserID != 0 && perUser >= maxStreamsPerUser {
+		return ErrTooManyStreamsForUser
+	}
+	if perServer >= maxStreamsPerServer {
+		return ErrTooManyStreamsForServer
+	}
+
+	s.ioStreams[streamId] = newIOStreamContext(creatorUserID, targetServerID, purpose)
+	s.publishIOStreamStateChangeLocked()
+	return nil
+}
+
+// StreamCount reports the registry size under the same lock used by lifecycle mutations.
+func (s *NezhaHandler) StreamCount() int {
+	s.ioStreamMutex.RLock()
+	defer s.ioStreamMutex.RUnlock()
+	return len(s.ioStreams)
+}

service/rpc/io_stream_quota_test.go

@@ -0,0 +1,106 @@
+package rpc
+
+import (
+	"errors"
+	"fmt"
+	"testing"
+)
+
+func TestCreateStreamExactUserBoundary(t *testing.T) {
+	h := NewNezhaHandler()
+	for i := 0; i < maxStreamsPerUser; i++ {
+		if err := h.CreateStream(fmt.Sprintf("quota-user-%d", i), 1, uint64(i+1)); err != nil {
+			t.Fatalf("20th user stream must succeed: %v", err)
+		}
+	}
+	if err := h.CreateStream("quota-user-21", 1, 100); !errors.Is(err, ErrTooManyStreamsForUser) {
+		t.Fatalf("21st user stream must be rejected: %v", err)
+	}
+}
+
+func TestCreateStreamNormalUserEverydayUseSucceeds(t *testing.T) {
+	h := NewNezhaHandler()
+	if err := h.CreateStream("term", 7, 1); err != nil {
+		t.Fatal(err)
+	}
+	if err := h.CreateStream("fm", 7, 1); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestCreateStreamNormalUsersAreIndependent(t *testing.T) {
+	h := NewNezhaHandler()
+	for userID := uint64(1); userID <= 5; userID++ {
+		for i := 0; i < maxStreamsPerUser; i++ {
+			if err := h.CreateStream(fmt.Sprintf("independent-%d-%d", userID, i), userID, 100+userID); err != nil {
+				t.Fatalf("user %d stream %d: %v", userID, i, err)
+			}
+		}
+	}
+}
+
+func TestCreateStreamExemptsInternalStreamsFromPerUserCap(t *testing.T) {
+	h := NewNezhaHandler()
+	for i := 0; i < maxStreamsPerUser*3; i++ {
+		if err := h.CreateStream(fmt.Sprintf("internal-user-%d", i), 0, uint64(i+1)); err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+func TestCreateStreamInternalStreamsStillCountTowardPerServerCap(t *testing.T) {
+	h := NewNezhaHandler()
+	for i := 0; i < maxStreamsPerServer; i++ {
+		if err := h.CreateStream(fmt.Sprintf("internal-server-%d", i), 0, 9); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if err := h.CreateStream("internal-server-over", 0, 9); !errors.Is(err, ErrTooManyStreamsForServer) {
+		t.Fatal(err)
+	}
+}
+
+func TestCreateStreamExactServerBoundary(t *testing.T) {
+	h := NewNezhaHandler()
+	for i := 0; i < maxStreamsPerServer; i++ {
+		if err := h.CreateStream(fmt.Sprintf("quota-server-%d", i), uint64(i+1), 2); err != nil {
+			t.Fatalf("40th server stream must succeed: %v", err)
+		}
+	}
+	if err := h.CreateStream("quota-server-41", 100, 2); !errors.Is(err, ErrTooManyStreamsForServer) {
+		t.Fatalf("41st server stream must be rejected: %v", err)
+	}
+}
+
+func TestCreateStreamReleasesUserAndServerSlots(t *testing.T) {
+	h := NewNezhaHandler()
+	for i := 0; i < maxStreamsPerUser; i++ {
+		if err := h.CreateStream(fmt.Sprintf("reuse-user-%d", i), 1, uint64(i+10)); err != nil {
+			t.Fatalf("user setup stream %d failed: %v", i, err)
+		}
+	}
+	if !errors.Is(h.CreateStream("reuse-user-over", 1, 100), ErrTooManyStreamsForUser) {
+		t.Fatal("user cap was not enforced")
+	}
+	if err := h.CloseStream("reuse-user-0"); err != nil {
+		t.Fatal(err)
+	}
+	if err := h.CreateStream("reuse-user-new", 1, 101); err != nil {
+		t.Fatalf("closed user slot must be reusable: %v", err)
+	}
+
+	for i := 0; i < maxStreamsPerServer; i++ {
+		if err := h.CreateStream(fmt.Sprintf("reuse-server-%d", i), uint64(i+2), 2); err != nil {
+			t.Fatalf("server setup stream %d failed: %v", i, err)
+		}
+	}
+	if !errors.Is(h.CreateStream("reuse-server-over", 100, 2), ErrTooManyStreamsForServer) {
+		t.Fatal("server cap was not enforced")
+	}
+	if err := h.CloseStream("reuse-server-0"); err != nil {
+		t.Fatal(err)
+	}
+	if err := h.CreateStream("reuse-server-new", 100, 2); err != nil {
+		t.Fatalf("closed server slot must be reusable: %v", err)
+	}
+}

service/rpc/io_stream_registry.go

@@ -0,0 +1,85 @@
+package rpc
+
+import (
+	"errors"
+	"io"
+)
+
+var ErrAgentStreamAlreadyConnected = errors.New("agent stream already connected")
+
+func (s *NezhaHandler) IsStreamAuthorizedForAgent(streamId string, agentServerID uint64) bool {
+	s.ioStreamMutex.RLock()
+	defer s.ioStreamMutex.RUnlock()
+	ctx, ok := s.ioStreams[streamId]
+	return ok && ctx.targetServerID != 0 && ctx.targetServerID == agentServerID
+}
+
+func (s *NezhaHandler) IsStreamAuthorizedForUser(streamId string, userID uint64, isAdmin bool) bool {
+	creator, found := s.StreamOwnership(streamId)
+	return found && (isAdmin || creator == userID)
+}
+
+func (s *NezhaHandler) StreamOwnership(streamId string) (uint64, bool) {
+	s.ioStreamMutex.RLock()
+	defer s.ioStreamMutex.RUnlock()
+	ctx, ok := s.ioStreams[streamId]
+	if !ok {
+		return 0, false
+	}
+	return ctx.creatorUserID, true
+}
+
+func (s *NezhaHandler) StreamTarget(streamId string) (uint64, bool) {
+	s.ioStreamMutex.RLock()
+	defer s.ioStreamMutex.RUnlock()
+	ctx, ok := s.ioStreams[streamId]
+	if !ok {
+		return 0, false
+	}
+	return ctx.targetServerID, true
+}
+
+func (s *NezhaHandler) GetStream(streamId string) (*ioStreamContext, error) {
+	s.ioStreamMutex.RLock()
+	defer s.ioStreamMutex.RUnlock()
+	if ctx, ok := s.ioStreams[streamId]; ok {
+		return ctx, nil
+	}
+	return nil, errors.New("stream not found")
+}
+
+func (s *NezhaHandler) UserConnected(streamId string, userIo io.ReadWriteCloser) error {
+	s.ioStreamMutex.Lock()
+	stream, ok := s.ioStreams[streamId]
+	if !ok {
+		s.ioStreamMutex.Unlock()
+		return errors.New("stream not found")
+	}
+	stream.userIo = userIo
+	s.ioStreamMutex.Unlock()
+	stream.userIoChOnce.Do(func() { close(stream.userIoConnectCh) })
+	return nil
+}
+
+func (s *NezhaHandler) AgentConnected(streamId string, agentIo io.ReadWriteCloser) error {
+	s.ioStreamMutex.Lock()
+	stream, ok := s.ioStreams[streamId]
+	if !ok {
+		s.ioStreamMutex.Unlock()
+		return errors.Join(errors.New("stream not found"), agentIo.Close())
+	}
+	if stream.agentIo != nil {
+		s.ioStreamMutex.Unlock()
+		return errors.Join(ErrAgentStreamAlreadyConnected, agentIo.Close())
+	}
+	stream.agentIo = agentIo
+	s.ioStreamMutex.Unlock()
+	stream.agentIoChOnce.Do(func() { close(stream.agentIoConnectCh) })
+	return nil
+}
+
+func (s *NezhaHandler) streamEndpoints(stream *ioStreamContext) (io.ReadWriteCloser, io.ReadWriteCloser) {
+	s.ioStreamMutex.RLock()
+	defer s.ioStreamMutex.RUnlock()
+	return stream.userIo, stream.agentIo
+}

service/rpc/io_stream_rpc.go

@@ -0,0 +1,58 @@
+package rpc
+
+import (
+	"fmt"
+	"log"
+	"time"
+
+	"github.com/nezhahq/nezha/pkg/grpcx"
+	pb "github.com/nezhahq/nezha/proto"
+)
+
+func (s *NezhaHandler) IOStream(stream pb.NezhaService_IOStreamServer) error {
+	clientID, err := s.Auth.Check(stream.Context())
+	if err != nil {
+		return err
+	}
+	id, err := stream.Recv()
+	if err != nil {
+		return err
+	}
+	if id == nil || !isValidIOStreamMagic(id.Data) {
+		return fmt.Errorf("invalid stream id")
+	}
+	streamID := string(id.Data[4:])
+	if !s.IsStreamAuthorizedForAgent(streamID, clientID) {
+		return fmt.Errorf("stream not authorized for agent")
+	}
+	if _, err := s.GetStream(streamID); err != nil {
+		return err
+	}
+	wrapper := grpcx.NewIOStreamWrapper(stream)
+	keepaliveDone := make(chan struct{})
+	go func() {
+		defer close(keepaliveDone)
+		ticker := time.NewTicker(30 * time.Second)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-wrapper.Context().Done():
+				return
+			case <-wrapper.Done():
+				return
+			case <-ticker.C:
+				if err := wrapper.SendKeepalive(); err != nil {
+					log.Printf("NEZHA>> IOStream keepAlive error: %v\n", err)
+					return
+				}
+			}
+		}
+	}()
+	if err := s.AgentConnected(streamID, wrapper); err != nil {
+		_ = wrapper.Close()
+		return err
+	}
+	wrapper.Wait()
+	<-keepaliveDone
+	return nil
+}

service/rpc/io_stream_state.go

@@ -0,0 +1,95 @@
+package rpc
+
+import (
+	"context"
+	"errors"
+)
+
+var ErrInvalidIOStreamStateExpectation = errors.New("invalid IOStream state expectation")
+
+type IOStreamState struct {
+	Count      int    `json:"count"`
+	Generation uint64 `json:"generation"`
+}
+
+type IOStreamStateExpectation struct {
+	// A pointer distinguishes an omitted count from an explicit zero count.
+	ExpectedCount   *int   `json:"expected_count,omitempty"`
+	PresentStreamID string `json:"present_stream_id,omitempty"`
+	AbsentStreamID  string `json:"absent_stream_id,omitempty"`
+}
+
+func ExpectedIOStreamCount(count int) *int {
+	return &count
+}
+
+func (s IOStreamStateExpectation) validate() error {
+	if s.ExpectedCount == nil && s.PresentStreamID == "" && s.AbsentStreamID == "" {
+		return ErrInvalidIOStreamStateExpectation
+	}
+	if s.ExpectedCount != nil && *s.ExpectedCount < 0 {
+		return ErrInvalidIOStreamStateExpectation
+	}
+	if s.PresentStreamID != "" && s.PresentStreamID == s.AbsentStreamID {
+		return ErrInvalidIOStreamStateExpectation
+	}
+	return nil
+}
+
+func (s *NezhaHandler) SnapshotIOStreamState() IOStreamState {
+	s.ioStreamMutex.RLock()
+	defer s.ioStreamMutex.RUnlock()
+	return s.snapshotIOStreamStateLocked()
+}
+
+func (s *NezhaHandler) snapshotIOStreamStateLocked() IOStreamState {
+	return IOStreamState{Count: len(s.ioStreams), Generation: s.ioStreamGeneration}
+}
+
+func (s *NezhaHandler) ioStreamStateExpectationSatisfiedLocked(expectation IOStreamStateExpectation) bool {
+	if expectation.ExpectedCount != nil && len(s.ioStreams) != *expectation.ExpectedCount {
+		return false
+	}
+	if expectation.PresentStreamID != "" {
+		if _, exists := s.ioStreams[expectation.PresentStreamID]; !exists {
+			return false
+		}
+	}
+	if expectation.AbsentStreamID != "" {
+		if _, exists := s.ioStreams[expectation.AbsentStreamID]; exists {
+			return false
+		}
+	}
+	return true
+}
+
+func (s *NezhaHandler) WaitForIOStreamState(ctx context.Context, expectation IOStreamStateExpectation) (IOStreamState, error) {
+	if err := expectation.validate(); err != nil {
+		return IOStreamState{}, err
+	}
+	for {
+		s.ioStreamMutex.RLock()
+		notify := s.ioStreamNotify
+		state := s.snapshotIOStreamStateLocked()
+		satisfied := s.ioStreamStateExpectationSatisfiedLocked(expectation)
+		observer := s.ioStreamWaitLockedHook
+		s.ioStreamMutex.RUnlock()
+		if observer != nil {
+			observer()
+		}
+		if satisfied {
+			return state, nil
+		}
+		select {
+		case <-ctx.Done():
+			return IOStreamState{}, ctx.Err()
+		case <-notify:
+		}
+	}
+}
+
+func (s *NezhaHandler) publishIOStreamStateChangeLocked() {
+	s.ioStreamGeneration++
+	close(s.ioStreamNotify)
+	s.ioStreamNotify = make(chan struct{})
+}

service/rpc/io_stream_state_expectation_test.go

@@ -0,0 +1,142 @@
+package rpc
+
+import (
+	"context"
+	"errors"
+	"strings"
+	"testing"
+)
+
+func TestWaitForIOStreamStateRejectsZeroValueExpectation(t *testing.T) {
+	handler := NewNezhaHandler()
+	if _, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{}); !errors.Is(err, ErrInvalidIOStreamStateExpectation) {
+		t.Fatalf("zero-value expectation error: %v", err)
+	}
+}
+
+func TestWaitForIOStreamStateAcceptsExplicitZeroCount(t *testing.T) {
+	handler := NewNezhaHandler()
+	state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0)})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if state != (IOStreamState{}) {
+		t.Fatalf("explicit zero state: %+v", state)
+	}
+}
+
+func TestWaitForIOStreamStateAcceptsPresentOnlyExpectation(t *testing.T) {
+	handler := NewNezhaHandler()
+	if err := handler.CreateStream("present-only", 1, 1); err != nil {
+		t.Fatal(err)
+	}
+	state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{PresentStreamID: "present-only"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if state.Count != 1 || state.Generation != 1 {
+		t.Fatalf("present-only state: %+v", state)
+	}
+}
+
+func TestWaitForIOStreamStateRejectsSamePresentAndAbsentID(t *testing.T) {
+	handler := NewNezhaHandler()
+	_, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{PresentStreamID: "same", AbsentStreamID: "same"})
+	if !errors.Is(err, ErrInvalidIOStreamStateExpectation) {
+		t.Fatalf("same identity expectation error: %v", err)
+	}
+}
+
+func TestWaitForIOStreamStateRequiresAllSpecifiedConditions(t *testing.T) {
+	handler := NewNezhaHandler()
+	if err := handler.CreateStream("present", 1, 1); err != nil {
+		t.Fatal(err)
+	}
+	if err := handler.CreateStream("other", 1, 2); err != nil {
+		t.Fatal(err)
+	}
+	if err := handler.CreateStream("absent", 1, 3); err != nil {
+		t.Fatal(err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	_, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{
+		ExpectedCount:   ExpectedIOStreamCount(2),
+		PresentStreamID: "present",
+		AbsentStreamID:  "absent",
+	})
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("combined expectation cancellation: %v", err)
+	}
+}
+
+func TestWaitForIOStreamStateAbsenceOnlyIgnoresUnrelatedStreams(t *testing.T) {
+	handler := NewNezhaHandler()
+	if err := handler.CreateStream("unrelated", 1, 1); err != nil {
+		t.Fatal(err)
+	}
+	state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{AbsentStreamID: "absent"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if state.Count != 1 || state.Generation != 1 {
+		t.Fatalf("absence-only state: %+v", state)
+	}
+}
+
+func TestWaitForIOStreamStateRejectsNegativeCountWithoutPrivateID(t *testing.T) {
+	handler := NewNezhaHandler()
+	_, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(-1), AbsentStreamID: "private-stream-id"})
+	if !errors.Is(err, ErrInvalidIOStreamStateExpectation) {
+		t.Fatalf("negative expectation error: %v", err)
+	}
+	if err != nil && strings.Contains(err.Error(), "private-stream-id") {
+		t.Fatalf("private stream ID leaked: %v", err)
+	}
+}
+
+func TestWaitForIOStreamStateRequiresCombinedCountAndAbsence(t *testing.T) {
+	handler := NewNezhaHandler()
+	if err := handler.CreateStream("present", 1, 1); err != nil {
+		t.Fatal(err)
+	}
+	state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1), AbsentStreamID: "absent"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if state.Count != 1 || state.Generation != 1 {
+		t.Fatalf("combined expectation state: %+v", state)
+	}
+}
+
+func TestWaitForIOStreamStateCancellationRemainsValidForUnsatisfiedExpectation(t *testing.T) {
+	handler := NewNezhaHandler()
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	if _, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}); !errors.Is(err, context.Canceled) {
+		t.Fatalf("cancel error: %v", err)
+	}
+}
+
+func TestWaitForIOStreamStateAlreadySatisfiedReturnsSnapshot(t *testing.T) {
+	handler := NewNezhaHandler()
+	state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0)})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if state != (IOStreamState{}) {
+		t.Fatalf("already satisfied state: %+v", state)
+	}
+}
+
+func TestWaitForIOStreamStateRejectsInvalidAndHonorsCancellation(t *testing.T) {
+	handler := NewNezhaHandler()
+	if _, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(-1)}); !errors.Is(err, ErrInvalidIOStreamStateExpectation) {
+		t.Fatalf("invalid count error: %v", err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	if _, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}); !errors.Is(err, context.Canceled) {
+		t.Fatalf("cancel error: %v", err)
+	}
+}

service/rpc/io_stream_state_test.go

@@ -0,0 +1,68 @@
+package rpc
+
+import (
+	"errors"
+	"testing"
+)
+
+func TestIOStreamStateSnapshotAndGeneration(t *testing.T) {
+	handler := NewNezhaHandler()
+	initial := handler.SnapshotIOStreamState()
+	if initial.Count != 0 || initial.Generation != 0 {
+		t.Fatalf("unexpected initial state: %+v", initial)
+	}
+	if err := handler.CreateStream("state-stream", 1, 1); err != nil {
+		t.Fatal(err)
+	}
+	created := handler.SnapshotIOStreamState()
+	if created.Count != 1 || created.Generation != 1 {
+		t.Fatalf("unexpected created state: %+v", created)
+	}
+	if err := handler.CreateStream("state-stream", 2, 2); !errors.Is(err, ErrStreamAlreadyExists) {
+		t.Fatalf("duplicate create error: %v", err)
+	}
+	if got := handler.SnapshotIOStreamState(); got != created {
+		t.Fatalf("duplicate create changed state: %+v", got)
+	}
+	if err := handler.CloseStream("unknown"); err != nil {
+		t.Fatal(err)
+	}
+	if got := handler.SnapshotIOStreamState(); got != created {
+		t.Fatalf("unknown close changed state: %+v", got)
+	}
+}
+
+func TestIOStreamStateRevocationPublishesOncePerBatch(t *testing.T) {
+	handler := NewNezhaHandler()
+	if err := handler.CreateStreamWithPurpose("purpose-a", 0, 1, PurposeMCPTransfer); err != nil {
+		t.Fatal(err)
+	}
+	if err := handler.CreateStreamWithPurpose("purpose-b", 0, 1, PurposeMCPTransfer); err != nil {
+		t.Fatal(err)
+	}
+	if err := handler.CreateStream("server-a", 0, 2); err != nil {
+		t.Fatal(err)
+	}
+	before := handler.SnapshotIOStreamState()
+	if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 2 {
+		t.Fatalf("revoked purpose streams: %d", revoked)
+	}
+	afterPurpose := handler.SnapshotIOStreamState()
+	if afterPurpose.Generation != before.Generation+1 || afterPurpose.Count != 1 {
+		t.Fatalf("purpose revocation state: before=%+v after=%+v", before, afterPurpose)
+	}
+	if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 0 {
+		t.Fatalf("repeat purpose revocation: %d", revoked)
+	}
+	if got := handler.SnapshotIOStreamState(); got != afterPurpose {
+		t.Fatalf("empty purpose revocation changed state: %+v", got)
+	}
+	handler.RevokeStreamsForServer(2)
+	if got := handler.SnapshotIOStreamState(); got.Generation != afterPurpose.Generation+1 || got.Count != 0 {
+		t.Fatalf("server revocation state: %+v", got)
+	}
+	handler.RevokeStreamsForServer(2)
+	if got := handler.SnapshotIOStreamState(); got.Generation != afterPurpose.Generation+1 {
+		t.Fatalf("empty server revocation changed generation: %+v", got)
+	}
+}

service/rpc/io_stream_state_wait_test.go

@@ -0,0 +1,220 @@
+package rpc
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestWaitForIOStreamStateWakesOnCloseAndAbsence(t *testing.T) {
+	handler := NewNezhaHandler()
+	if err := handler.CreateStream("wait-state", 1, 1); err != nil {
+		t.Fatal(err)
+	}
+	result := make(chan IOStreamState, 1)
+	go func() {
+		state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0), AbsentStreamID: "wait-state"})
+		if err != nil {
+			t.Errorf("wait failed: %v", err)
+			return
+		}
+		result <- state
+	}()
+	if err := handler.CloseStream("wait-state"); err != nil {
+		t.Fatal(err)
+	}
+	state := <-result
+	if state.Count != 0 || state.Generation != 2 {
+		t.Fatalf("unexpected waited state: %+v", state)
+	}
+}
+
+func TestWaitForIOStreamStateCreateWakeUsesCapturedNotification(t *testing.T) {
+	handler := NewNezhaHandler()
+	waitReady := make(chan struct{})
+	handler.ioStreamWaitLockedHook = func() {
+		select {
+		case <-waitReady:
+		default:
+			close(waitReady)
+		}
+	}
+	result := make(chan IOStreamState, 1)
+	go func() {
+		state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)})
+		if err == nil {
+			result <- state
+		}
+	}()
+	select {
+	case <-waitReady:
+	case <-time.After(time.Second):
+		t.Fatal("waiter did not capture its notification channel")
+	}
+	if err := handler.CreateStream("create-wake", 1, 1); err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case state := <-result:
+		if state.Count != 1 || state.Generation != 1 {
+			t.Fatalf("unexpected created state: %+v", state)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("create did not wake waiter")
+	}
+}
+
+func TestWaitForIOStreamStateCloseWakeUsesCapturedNotification(t *testing.T) {
+	handler := NewNezhaHandler()
+	if err := handler.CreateStream("close-wake", 1, 1); err != nil {
+		t.Fatal(err)
+	}
+	waitReady := make(chan struct{})
+	handler.ioStreamWaitLockedHook = func() {
+		select {
+		case <-waitReady:
+		default:
+			close(waitReady)
+		}
+	}
+	result := make(chan IOStreamState, 1)
+	go func() {
+		state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0), AbsentStreamID: "close-wake"})
+		if err == nil {
+			result <- state
+		}
+	}()
+	select {
+	case <-waitReady:
+	case <-time.After(time.Second):
+		t.Fatal("waiter did not capture its notification channel")
+	}
+	if err := handler.CloseStream("close-wake"); err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case state := <-result:
+		if state.Count != 0 || state.Generation != 2 {
+			t.Fatalf("unexpected closed state: %+v", state)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("close did not wake waiter")
+	}
+}
+
+func TestWaitForIOStreamStateDoesNotMissMutationBetweenSnapshotAndWait(t *testing.T) {
+	handler := NewNezhaHandler()
+	hookCalled := make(chan struct{})
+	mutationDone := make(chan error, 1)
+	var hookOnce sync.Once
+	handler.ioStreamWaitLockedHook = func() {
+		hookOnce.Do(func() {
+			close(hookCalled)
+			go func() {
+				mutationDone <- handler.CreateStream("lost-wakeup", 1, 1)
+			}()
+		})
+	}
+	result := make(chan IOStreamState, 1)
+	go func() {
+		state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)})
+		if err == nil {
+			result <- state
+		}
+	}()
+	select {
+	case <-hookCalled:
+	case <-time.After(time.Second):
+		t.Fatal("waiter did not reach deterministic mutation seam")
+	}
+	select {
+	case err := <-mutationDone:
+		if err != nil {
+			t.Fatal(err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("mutation did not complete")
+	}
+	select {
+	case state := <-result:
+		if state.Count != 1 || state.Generation != 1 {
+			t.Fatalf("unexpected mutation state: %+v", state)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("waiter missed mutation published during wait setup")
+	}
+}
+
+func TestWaitForIOStreamStateConcurrentCreateCloseWaiters(t *testing.T) {
+	handler := NewNezhaHandler()
+	created := make(chan IOStreamState, 1)
+	closed := make(chan IOStreamState, 1)
+	go func() {
+		state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)})
+		if err == nil {
+			created <- state
+		}
+	}()
+	if err := handler.CreateStream("concurrent", 1, 1); err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case state := <-created:
+		if state.Count != 1 {
+			t.Fatalf("created waiter state: %+v", state)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("created waiter did not wake")
+	}
+	go func() {
+		state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0), AbsentStreamID: "concurrent"})
+		if err == nil {
+			closed <- state
+		}
+	}()
+	if err := handler.CloseStream("concurrent"); err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case state := <-closed:
+		if state.Count != 0 {
+			t.Fatalf("closed waiter state: %+v", state)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("closed waiter did not wake")
+	}
+}
+
+func TestWaitForIOStreamStateDoesNotAcceptUnrelatedSameCountForPresentID(t *testing.T) {
+	handler := NewNezhaHandler()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	result := make(chan error, 1)
+	go func() {
+		_, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{
+			ExpectedCount:   ExpectedIOStreamCount(1),
+			PresentStreamID: "wanted",
+		})
+		result <- err
+	}()
+	if err := handler.CreateStream("unrelated", 1, 1); err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case err := <-result:
+		if err == nil {
+			t.Fatal("same-count unrelated stream satisfied identity expectation")
+		}
+	default:
+	}
+	cancel()
+	select {
+	case err := <-result:
+		if err == nil {
+			t.Fatal("identity waiter unexpectedly succeeded")
+		}
+	case <-time.After(time.Second):
+		t.Fatal("identity waiter did not observe cancellation")
+	}
+}

service/rpc/io_stream_test.go

@@ -1,8 +1,6 @@
 package rpc
 
 import (
-	"errors"
-	"fmt"
 	"io"
 	"reflect"
 	"testing"
@@ -100,130 +98,6 @@ func TestIOStream(t *testing.T) {
 	})
 }
 
-// The WebSocket stream endpoints (terminal / fm) were unbounded: an
-// authenticated member could open thousands of streams, each spawning
-// goroutines, a 1 MiB buffer, and an agent-side PTY, exhausting dashboard and
-// agent resources (GHSA-jg62-j5h6-8mpq). CreateStream now caps concurrent
-// streams per user and per server. These tests pin the caps and the
-// dashboard-internal (uid==0) exemption.
-
-// Baseline: a normal operator opening a terminal and a file-manager session
-// against one server (the everyday case) must always succeed — the cap exists
-// to stop floods, not to interfere with ordinary use.
-func TestCreateStreamNormalUserEverydayUseSucceeds(t *testing.T) {
-	h := NewNezhaHandler()
-	const uid, serverID = uint64(7), uint64(1)
-
-	if err := h.CreateStream("term", uid, serverID); err != nil {
-		t.Fatalf("opening a terminal must succeed for a normal user, got %v", err)
-	}
-	if err := h.CreateStream("fm", uid, serverID); err != nil {
-		t.Fatalf("opening a file manager alongside a terminal must succeed, got %v", err)
-	}
-}
-
-// Several normal users working at the same time must not interfere: one user's
-// streams do not consume another user's per-user budget.
-func TestCreateStreamNormalUsersAreIndependent(t *testing.T) {
-	h := NewNezhaHandler()
-
-	for u := uint64(1); u <= 5; u++ {
-		for i := 0; i < maxStreamsPerUser; i++ {
-			id := fmt.Sprintf("u%d-s%d", u, i)
-			if err := h.CreateStream(id, u, 100+u); err != nil {
-				t.Fatalf("user %d stream %d must succeed; per-user budgets must be independent, got %v", u, i, err)
-			}
-		}
-	}
-}
-
-func TestCreateStreamEnforcesPerUserCap(t *testing.T) {
-	h := NewNezhaHandler()
-	const uid = uint64(42)
-
-	for i := 0; i < maxStreamsPerUser; i++ {
-		if err := h.CreateStream(fmt.Sprintf("u-%d", i), uid, uint64(i)); err != nil {
-			t.Fatalf("stream %d within the per-user cap must succeed, got %v", i, err)
-		}
-	}
-
-	err := h.CreateStream("u-over", uid, 9999)
-	if !errors.Is(err, ErrTooManyStreamsForUser) {
-		t.Fatalf("the (maxStreamsPerUser+1)-th stream must be rejected with ErrTooManyStreamsForUser, got %v", err)
-	}
-}
-
-func TestCreateStreamEnforcesPerServerCap(t *testing.T) {
-	h := NewNezhaHandler()
-	const serverID = uint64(7)
-
-	for i := 0; i < maxStreamsPerServer; i++ {
-		if err := h.CreateStream(fmt.Sprintf("s-%d", i), uint64(i+1), serverID); err != nil {
-			t.Fatalf("stream %d within the per-server cap must succeed, got %v", i, err)
-		}
-	}
-
-	err := h.CreateStream("s-over", 99999, serverID)
-	if !errors.Is(err, ErrTooManyStreamsForServer) {
-		t.Fatalf("the (maxStreamsPerServer+1)-th stream to one server must be rejected with ErrTooManyStreamsForServer, got %v", err)
-	}
-}
-
-// Dashboard-internal streams (NAT, server transfer, MCP transfer) pass
-// creatorUserID==0. They must NOT be capped per user, or those features would
-// throttle themselves; but they must still count toward the per-server cap so
-// no single server can be flooded regardless of the originating path.
-func TestCreateStreamExemptsInternalStreamsFromPerUserCap(t *testing.T) {
-	h := NewNezhaHandler()
-
-	for i := 0; i < maxStreamsPerUser*3; i++ {
-		if err := h.CreateStream(fmt.Sprintf("internal-%d", i), 0, uint64(i)); err != nil {
-			t.Fatalf("internal stream %d (uid==0) must never hit the per-user cap, got %v", i, err)
-		}
-	}
-}
-
-func TestCreateStreamInternalStreamsStillCountTowardPerServerCap(t *testing.T) {
-	h := NewNezhaHandler()
-	const serverID = uint64(3)
-
-	for i := 0; i < maxStreamsPerServer; i++ {
-		if err := h.CreateStream(fmt.Sprintf("internal-s-%d", i), 0, serverID); err != nil {
-			t.Fatalf("internal stream %d within the per-server cap must succeed, got %v", i, err)
-		}
-	}
-
-	err := h.CreateStream("internal-s-over", 0, serverID)
-	if !errors.Is(err, ErrTooManyStreamsForServer) {
-		t.Fatalf("internal streams must still be subject to the per-server cap, got %v", err)
-	}
-}
-
-// Closing a stream must free its slot so a user who hit the cap can open new
-// streams after old ones end — otherwise normal churn would permanently lock
-// a user out.
-func TestCreateStreamFreesSlotAfterClose(t *testing.T) {
-	h := NewNezhaHandler()
-	const uid = uint64(55)
-
-	for i := 0; i < maxStreamsPerUser; i++ {
-		if err := h.CreateStream(fmt.Sprintf("c-%d", i), uid, 1); err != nil {
-			t.Fatalf("setup stream %d must succeed, got %v", i, err)
-		}
-	}
-	if err := h.CreateStream("c-over", uid, 1); !errors.Is(err, ErrTooManyStreamsForUser) {
-		t.Fatalf("expected per-user cap to be hit, got %v", err)
-	}
-
-	if err := h.CloseStream("c-0"); err != nil {
-		t.Fatalf("CloseStream failed: %v", err)
-	}
-
-	if err := h.CreateStream("c-after-close", uid, 1); err != nil {
-		t.Fatalf("after closing one stream the user must be able to open another, got %v", err)
-	}
-}
-
 func newPipeReadWriter() io.ReadWriteCloser {
 	r, w := io.Pipe()
 	return struct {

service/rpc/nezha.go

@@ -5,13 +5,10 @@ import (
 	"errors"
 	"fmt"
 	"log"
-	"net"
 	"sync"
 	"time"
 
 	"github.com/jinzhu/copier"
-	geoipx "github.com/nezhahq/nezha/pkg/geoip"
-	"github.com/nezhahq/nezha/pkg/grpcx"
 	"github.com/nezhahq/nezha/pkg/tsdb"
 
 	"github.com/nezhahq/nezha/model"
@@ -24,17 +21,36 @@ var _ pb.NezhaServiceServer = (*NezhaHandler)(nil)
 var NezhaHandlerSingleton *NezhaHandler
 
 type NezhaHandler struct {
-	Auth          *authHandler
-	ioStreams     map[string]*ioStreamContext
-	ioStreamMutex *sync.RWMutex
+	Auth                   *authHandler
+	ioStreams              map[string]*ioStreamContext
+	ioStreamMutex          *sync.RWMutex
+	ioStreamGeneration     uint64
+	ioStreamNotify         chan struct{}
+	ioStreamWaitLockedHook func()
+	// Capability authorization and exact stream deletion share ioStreamMutex to avoid TOCTOU.
+	agentCompatCapabilities agentCompatCapabilityState
+}
+
+type serverMetricsWriter func(*tsdb.ServerMetrics) error
+
+var writeServerMetrics serverMetricsWriter = writeServerMetricsToTSDB
+
+func writeServerMetricsToTSDB(metrics *tsdb.ServerMetrics) error {
+	if !singleton.TSDBEnabled() {
+		return nil
+	}
+	return singleton.TSDBShared.WriteServerMetrics(metrics)
 }
 
 func NewNezhaHandler() *NezhaHandler {
-	return &NezhaHandler{
-		Auth:          &authHandler{},
-		ioStreamMutex: new(sync.RWMutex),
-		ioStreams:     make(map[string]*ioStreamContext),
+	handler := &NezhaHandler{
+		Auth:           &authHandler{},
+		ioStreamMutex:  new(sync.RWMutex),
+		ioStreams:      make(map[string]*ioStreamContext),
+		ioStreamNotify: make(chan struct{}),
 	}
+	handler.initializeAgentCompatCapabilities()
+	return handler
 }
 
 // attachRequestTaskStream resolves the server for clientID and publishes the
@@ -161,84 +177,91 @@ func (s *NezhaHandler) ReportSystemState(stream pb.NezhaService_ReportSystemStat
 	if err != nil {
 		return err
 	}
+	server, ok := singleton.ServerShared.Get(clientID)
+	if !ok || server == nil {
+		return errors.New("server not found")
+	}
+	lease := server.AttachStateStream(stream)
+	defer lease.Clear()
 	var state *pb.State
+	var stateCount uint64
 	for {
 		state, err = stream.Recv()
 		if err != nil {
 			log.Printf("NEZHA>> ReportSystemState error: %v, clientID: %d\n", err, clientID)
 			return err
 		}
+		stateCount++
 		innerState := model.PB2State(state)
 
-		server, ok := singleton.ServerShared.Get(clientID)
-		if !ok || server == nil {
-			return errors.New("server not found")
-		}
-
-		server.LastActive = time.Now()
-		server.State = &innerState
-
-		if singleton.TSDBEnabled() {
-			maxTemp := 0.0
-			for _, t := range innerState.Temperatures {
-				if t.Temperature > maxTemp {
-					maxTemp = t.Temperature
+		lastActive := time.Now()
+		accepted := lease.UpdateStateWithSideEffect(&innerState, lastActive, func() error {
+			{
+				maxTemp := 0.0
+				for _, t := range innerState.Temperatures {
+					if t.Temperature > maxTemp {
+						maxTemp = t.Temperature
+					}
+				}
+				maxGPU := 0.0
+				for _, g := range innerState.GPU {
+					if g > maxGPU {
+						maxGPU = g
+					}
+				}
+				if err := writeServerMetrics(&tsdb.ServerMetrics{
+					ServerID:       clientID,
+					Timestamp:      lastActive,
+					CPU:            innerState.CPU,
+					MemUsed:        innerState.MemUsed,
+					SwapUsed:       innerState.SwapUsed,
+					DiskUsed:       innerState.DiskUsed,
+					NetInSpeed:     innerState.NetInSpeed,
+					NetOutSpeed:    innerState.NetOutSpeed,
+					NetInTransfer:  innerState.NetInTransfer,
+					NetOutTransfer: innerState.NetOutTransfer,
+					Load1:          innerState.Load1,
+					Load5:          innerState.Load5,
+					Load15:         innerState.Load15,
+					TCPConnCount:   innerState.TcpConnCount,
+					UDPConnCount:   innerState.UdpConnCount,
+					ProcessCount:   innerState.ProcessCount,
+					Temperature:    maxTemp,
+					Uptime:         innerState.Uptime,
+					GPU:            maxGPU,
+				}); err != nil {
+					log.Printf("NEZHA>> Failed to write server metrics to TSDB: %v", err)
 				}
 			}
-			maxGPU := 0.0
-			for _, g := range innerState.GPU {
-				if g > maxGPU {
-					maxGPU = g
-				}
-			}
-			if err := singleton.TSDBShared.WriteServerMetrics(&tsdb.ServerMetrics{
-				ServerID:       clientID,
-				Timestamp:      time.Now(),
-				CPU:            innerState.CPU,
-				MemUsed:        innerState.MemUsed,
-				SwapUsed:       innerState.SwapUsed,
-				DiskUsed:       innerState.DiskUsed,
-				NetInSpeed:     innerState.NetInSpeed,
-				NetOutSpeed:    innerState.NetOutSpeed,
-				NetInTransfer:  innerState.NetInTransfer,
-				NetOutTransfer: innerState.NetOutTransfer,
-				Load1:          innerState.Load1,
-				Load5:          innerState.Load5,
-				Load15:         innerState.Load15,
-				TCPConnCount:   innerState.TcpConnCount,
-				UDPConnCount:   innerState.UdpConnCount,
-				ProcessCount:   innerState.ProcessCount,
-				Temperature:    maxTemp,
-				Uptime:         innerState.Uptime,
-				GPU:            maxGPU,
-			}); err != nil {
-				log.Printf("NEZHA>> Failed to write server metrics to TSDB: %v", err)
-			}
+			return nil
+		})
+		if !accepted {
+			return errors.New("state stream superseded")
 		}
 
-		// 应对 dashboard / agent 重启的情况，如果从未记录过，先打点，等到小时时间点时入库
-		if server.PrevTransferInSnapshot == 0 || server.PrevTransferOutSnapshot == 0 {
-			server.PrevTransferInSnapshot = state.NetInTransfer
-			server.PrevTransferOutSnapshot = state.NetOutTransfer
+		if err := notifyStateReceived(clientID, server.UUID, lease.Generation(), stateCount); err != nil {
+			return err
+		}
+		if err := notifyReceiptAccepted(clientID, server.UUID, lease.Generation(), stateCount); err != nil {
+			return err
 		}
-
 		if err = stream.Send(&pb.Receipt{Proced: true}); err != nil {
 			return err
 		}
 	}
 }
 
-func (s *NezhaHandler) onReportSystemInfo(c context.Context, r *pb.Host) error {
+func (s *NezhaHandler) onReportSystemInfo(c context.Context, r *pb.Host) (model.HostReportResult, error) {
 	var clientID uint64
 	var err error
 	if clientID, err = s.Auth.Check(c); err != nil {
-		return err
+		return model.HostReportResult{}, err
 	}
 	host := model.PB2Host(r)
 
 	server, ok := singleton.ServerShared.Get(clientID)
 	if !ok || server == nil {
-		return errors.New("server not found")
+		return model.HostReportResult{}, errors.New("server not found")
 	}
 
 	/**
@@ -246,170 +269,23 @@ func (s *NezhaHandler) onReportSystemInfo(c context.Context, r *pb.Host) error {
 	 * 当 agent 重启时，bootTime 变大，agent 端会先上报 host 信息，然后上报 state 信息
 	 * 这时可以借助上报顺序的空档，立即记录停机前的数据并重置 Prev* 数据，并由接下来的 state 方法重新赋值
 	 */
-	if !server.LastActive.IsZero() && host.BootTime > server.Host.BootTime {
-		singleton.RecordTransferHourlyUsage(server)
-		server.PrevTransferInSnapshot = 0
-		server.PrevTransferOutSnapshot = 0
-	}
-
-	server.Host = &host
-	return nil
+	return server.RuntimeHandle().ApplyHostReport(&host, time.Now(), singleton.PersistTransfer)
 }
 
 func (s *NezhaHandler) ReportSystemInfo(c context.Context, r *pb.Host) (*pb.Receipt, error) {
-	if err := s.onReportSystemInfo(c, r); err != nil {
+	if _, err := s.onReportSystemInfo(c, r); err != nil {
 		return nil, err
 	}
 	return &pb.Receipt{Proced: true}, nil
 }
 
 func (s *NezhaHandler) ReportSystemInfo2(c context.Context, r *pb.Host) (*pb.Uint64Receipt, error) {
-	if err := s.onReportSystemInfo(c, r); err != nil {
+	result, err := s.onReportSystemInfo(c, r)
+	if err != nil {
+		return nil, err
+	}
+	if err := notifyInfo2(result.ServerID, result.UUID); err != nil {
 		return nil, err
 	}
 	return &pb.Uint64Receipt{Data: singleton.DashboardBootTime}, nil
 }
-
-func (s *NezhaHandler) IOStream(stream pb.NezhaService_IOStreamServer) error {
-	clientID, err := s.Auth.Check(stream.Context())
-	if err != nil {
-		return err
-	}
-	id, err := stream.Recv()
-	if err != nil {
-		return err
-	}
-
-	// ff05ff05 是 Nezha 的魔数，用于标识流 ID。校验由 isValidIOStreamMagic 完成，
-	// 历史 inline 检查曾因 && 短路放过几乎全部非魔数 payload (byte0==0xff 即通过)。
-	if id == nil || !isValidIOStreamMagic(id.Data) {
-		return fmt.Errorf("invalid stream id")
-	}
-
-	streamId := string(id.Data[4:])
-
-	// agent 侧归属校验：只有 createTerminal / createFM / ServeNAT 选定的目标 server
-	// 才能接管该 stream。漏掉这一步等同于把 terminal / fm / NAT 会话向所有合法 agent
-	// 开放（任何获得 streamId 的 agent 都能抢答），构成 session-hijack RCE 中介。
-	// 这是 commit 6661d6a（user 侧归属校验）的对偶补丁。先校验后启 keepalive，
-	// 避免未授权 agent 触发悬空 goroutine 持续向其发心跳。
-	if !s.IsStreamAuthorizedForAgent(streamId, clientID) {
-		return fmt.Errorf("stream not authorized for agent")
-	}
-
-	if _, err := s.GetStream(streamId); err != nil {
-		return err
-	}
-	iw := grpcx.NewIOStreamWrapper(stream)
-
-	// Keepalive MUST go through the wrapper so it shares the same sendMu as
-	// MCP fs.transfer / terminal / fm Writers. Calling stream.Send directly
-	// here used to race those Writers — grpc-go forbids concurrent SendMsg
-	// on the same stream. The wrapper's sendMu is the dashboard-side dual
-	// of agent/cmd/agent/mcp_fs_transfer.go's serialIOStreamSender.
-	keepaliveDone := make(chan struct{})
-	go func() {
-		defer close(keepaliveDone)
-		ticker := time.NewTicker(time.Second * 30)
-		defer ticker.Stop()
-		for {
-			select {
-			case <-iw.Context().Done():
-				return
-			case <-iw.Done():
-				// 业务侧（CloseStream / RevokeStreamsForPurpose）调过
-				// iw.Close()。即便底层 gRPC stream context 尚未取消，也
-				// 必须立刻收手——否则要再等一整个 30s tick，handler 在
-				// iw.Wait() 之后又得多等一拍 keepaliveDone 才能返回。
-				return
-			case <-ticker.C:
-				if err := iw.SendKeepalive(); err != nil {
-					log.Printf("NEZHA>> IOStream keepAlive error: %v\n", err)
-					return
-				}
-			}
-		}
-	}()
-
-	if err := s.AgentConnected(streamId, iw); err != nil {
-		return err
-	}
-	iw.Wait()
-	<-keepaliveDone
-	return nil
-}
-
-func (s *NezhaHandler) ReportGeoIP(c context.Context, r *pb.GeoIP) (*pb.GeoIP, error) {
-	var clientID uint64
-	var err error
-	if clientID, err = s.Auth.Check(c); err != nil {
-		return nil, err
-	}
-
-	geoip := model.PB2GeoIP(r)
-	use6 := r.GetUse6()
-
-	if geoip.IP.IPv4Addr == "" && geoip.IP.IPv6Addr == "" {
-		ip, _ := c.Value(model.CtxKeyRealIP{}).(string)
-		if ip == "" {
-			ip, _ = c.Value(model.CtxKeyConnectingIP{}).(string)
-		}
-		geoip.IP.IPv4Addr = ip
-	}
-
-	joinedIP := geoip.IP.Join()
-
-	server, ok := singleton.ServerShared.Get(clientID)
-	if !ok || server == nil {
-		return nil, fmt.Errorf("server not found")
-	}
-
-	// 检查并更新DDNS
-	if server.EnableDDNS && joinedIP != "" &&
-		(server.GeoIP == nil || server.GeoIP.IP != geoip.IP) {
-		ipv4 := geoip.IP.IPv4Addr
-		ipv6 := geoip.IP.IPv6Addr
-
-		if err := singleton.ServerShared.UpdateDDNS(server, &model.IP{IPv4Addr: ipv4, IPv6Addr: ipv6}); err != nil {
-			log.Printf("NEZHA>> Failed to update DDNS for server %d: %v", err, server.ID)
-		}
-	}
-
-	// 发送IP变动通知
-	if server.GeoIP != nil && singleton.Conf.EnableIPChangeNotification &&
-		((singleton.Conf.Cover == model.ConfigCoverAll && !singleton.Conf.IgnoredIPNotificationServerIDs[clientID]) ||
-			(singleton.Conf.Cover == model.ConfigCoverIgnoreAll && singleton.Conf.IgnoredIPNotificationServerIDs[clientID])) &&
-		server.GeoIP.IP.Join() != "" &&
-		joinedIP != "" &&
-		server.GeoIP.IP != geoip.IP {
-
-		singleton.NotificationShared.SendNotification(singleton.Conf.IPChangeNotificationGroupID,
-			fmt.Sprintf(
-				"[%s] %s, %s => %s",
-				singleton.Localizer.T("IP Changed"),
-				server.Name, singleton.IPDesensitize(server.GeoIP.IP.Join()),
-				singleton.IPDesensitize(joinedIP),
-			),
-			"")
-	}
-
-	// 根据内置数据库查询 IP 地理位置
-	var ip string
-	if geoip.IP.IPv6Addr != "" && (use6 || geoip.IP.IPv4Addr == "") {
-		ip = geoip.IP.IPv6Addr
-	} else {
-		ip = geoip.IP.IPv4Addr
-	}
-
-	netIP := net.ParseIP(ip)
-	location, err := geoipx.Lookup(netIP)
-	if err != nil {
-		log.Printf("NEZHA>> geoip.Lookup: %v", err)
-	}
-	geoip.CountryCode = location
-
-	// 将地区码写入到 Host
-	server.GeoIP = &geoip
-
-	return &pb.GeoIP{Ip: nil, CountryCode: location, DashboardBootTime: singleton.DashboardBootTime}, nil
-}

service/rpc/request_task_security_test.go

@@ -11,6 +11,7 @@ import (
 	"gorm.io/gorm"
 
 	"github.com/nezhahq/nezha/model"
+	"github.com/nezhahq/nezha/pkg/i18n"
 	pb "github.com/nezhahq/nezha/proto"
 	"github.com/nezhahq/nezha/service/singleton"
 )
@@ -257,7 +258,10 @@ func setupRequestTaskSecurityFixture(t *testing.T, servers []*model.Server, cron
 	originalDB := singleton.DB
 	originalConf := singleton.Conf
 	originalLoc := singleton.Loc
+	originalLocalizer := singleton.Localizer
+	originalNotification := singleton.NotificationShared
 	originalServerShared := singleton.ServerShared
+	originalServiceSentinel := singleton.ServiceSentinelShared
 	originalCronShared := singleton.CronShared
 	originalUserInfoMap := singleton.UserInfoMap
 	originalAgentSecretToUserID := singleton.AgentSecretToUserId
@@ -275,6 +279,8 @@ func setupRequestTaskSecurityFixture(t *testing.T, servers []*model.Server, cron
 	singleton.DB = db
 	singleton.Conf = &singleton.ConfigClass{Config: &model.Config{}}
 	singleton.Loc = time.UTC
+	singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations)
+	singleton.NotificationShared = singleton.NewEmptyNotificationClassForTest()
 	if err := singleton.DB.AutoMigrate(model.Server{}, model.Cron{}); err != nil {
 		t.Fatal(err)
 	}
@@ -297,13 +303,14 @@ func setupRequestTaskSecurityFixture(t *testing.T, servers []*model.Server, cron
 	singleton.CronShared = singleton.NewCronClass()
 
 	t.Cleanup(func() {
-		if singleton.CronShared != nil && singleton.CronShared.Cron != nil {
-			singleton.CronShared.Stop()
-		}
-		sqlDB.Close()
+		singleton.CronShared.Close()
+		_ = sqlDB.Close()
 		singleton.DB = originalDB
 		singleton.Conf = originalConf
 		singleton.Loc = originalLoc
+		singleton.Localizer = originalLocalizer
+		singleton.NotificationShared = originalNotification
+		singleton.ServiceSentinelShared = originalServiceSentinel
 		singleton.ServerShared = originalServerShared
 		singleton.CronShared = originalCronShared
 		singleton.UserLock.Lock()

service/rpc/state_metrics_test.go

@@ -0,0 +1,42 @@
+package rpc
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/nezhahq/nezha/model"
+	"github.com/nezhahq/nezha/pkg/tsdb"
+)
+
+func TestStateMetricsWriterRunsOnlyForCurrentGeneration(t *testing.T) {
+	// Given
+	server := &model.Server{}
+	model.InitServer(server)
+	oldLease := server.AttachStateStream(stateGenerationStream{})
+	newLease := server.AttachStateStream(stateGenerationStream{})
+	oldCalls := 0
+	newCalls := 0
+	oldWriter := writeServerMetrics
+	writeServerMetrics = func(*tsdb.ServerMetrics) error {
+		newCalls++
+		return nil
+	}
+	t.Cleanup(func() { writeServerMetrics = oldWriter })
+
+	// When
+	oldAccepted := server.UpdateStateIfCurrentWithSideEffect(oldLease, &model.HostState{Uptime: 11}, time.Unix(100, 0), func() error {
+		oldCalls++
+		return writeServerMetrics(&tsdb.ServerMetrics{ServerID: 7, Timestamp: time.Unix(100, 0)})
+	})
+	newAccepted := server.UpdateStateIfCurrentWithSideEffect(newLease, &model.HostState{Uptime: 22}, time.Unix(200, 0), func() error {
+		return writeServerMetrics(&tsdb.ServerMetrics{ServerID: 7, Timestamp: time.Unix(200, 0)})
+	})
+
+	// Then
+	require.False(t, oldAccepted)
+	require.True(t, newAccepted)
+	require.Zero(t, oldCalls)
+	require.Equal(t, 1, newCalls)
+}

service/rpc/state_stream_generation_test.go

@@ -0,0 +1,206 @@
+package rpc
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/metadata"
+
+	"github.com/nezhahq/nezha/model"
+	"github.com/nezhahq/nezha/pkg/tsdb"
+	pb "github.com/nezhahq/nezha/proto"
+	"github.com/nezhahq/nezha/service/singleton"
+)
+
+func TestReportSystemState_HandlerWaitsForMetricsBeforeReceipt(t *testing.T) {
+	// Given
+	reporter := requestTaskSecurityServer(9, 200, "ffffffff-ffff-ffff-ffff-ffffffffffff")
+	setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{200: {Role: model.RoleMember}}, map[string]uint64{"reporter-secret": 200})
+	stop := make(chan struct{})
+	stream := &stateGenerationHandlerStream{
+		ctx:    metadata.NewIncomingContext(context.Background(), metadata.Pairs("client_secret", "reporter-secret", "client_uuid", reporter.UUID)),
+		states: make(chan *pb.State, 1), receipts: make(chan *pb.Receipt, 1), stop: stop,
+	}
+	stream.states <- &pb.State{Uptime: 44}
+	metricsStarted := make(chan *tsdb.ServerMetrics, 1)
+	metricsRelease := make(chan struct{})
+	oldWriter := writeServerMetrics
+	writeServerMetrics = func(metrics *tsdb.ServerMetrics) error {
+		metricsStarted <- metrics
+		<-metricsRelease
+		return nil
+	}
+	t.Cleanup(func() { writeServerMetrics = oldWriter })
+
+	// When
+	done := make(chan error, 1)
+	go func() { done <- NewNezhaHandler().ReportSystemState(stream) }()
+	metrics := <-metricsStarted
+	select {
+	case <-stream.receipts:
+		t.Fatal("receipt sent before metrics writer completed")
+	default:
+	}
+	close(metricsRelease)
+
+	// Then
+	require.Equal(t, reporter.ID, metrics.ServerID)
+	require.Equal(t, uint64(44), metrics.Uptime)
+	require.NotNil(t, <-stream.receipts)
+	current, ok := singleton.ServerShared.Get(reporter.ID)
+	require.True(t, ok)
+	require.Equal(t, current.RuntimeSnapshot().LastActive, metrics.Timestamp)
+	close(stop)
+	require.ErrorIs(t, <-done, context.Canceled)
+}
+
+type stateGenerationStream struct{}
+
+func (stateGenerationStream) Send(*pb.Receipt) error       { return nil }
+func (stateGenerationStream) Recv() (*pb.State, error)     { return nil, nil }
+func (stateGenerationStream) SetHeader(metadata.MD) error  { return nil }
+func (stateGenerationStream) SendHeader(metadata.MD) error { return nil }
+func (stateGenerationStream) SetTrailer(metadata.MD)       {}
+func (stateGenerationStream) Context() context.Context     { return context.Background() }
+func (stateGenerationStream) SendMsg(any) error            { return nil }
+func (stateGenerationStream) RecvMsg(any) error            { return nil }
+
+type stateGenerationHandlerStream struct {
+	ctx      context.Context
+	states   chan *pb.State
+	receipts chan *pb.Receipt
+	stop     <-chan struct{}
+}
+
+func (s *stateGenerationHandlerStream) Send(receipt *pb.Receipt) error {
+	s.receipts <- receipt
+	return nil
+}
+
+func (s *stateGenerationHandlerStream) Recv() (*pb.State, error) {
+	select {
+	case state := <-s.states:
+		return state, nil
+	case <-s.stop:
+		return nil, context.Canceled
+	}
+}
+
+func (s *stateGenerationHandlerStream) SetHeader(metadata.MD) error  { return nil }
+func (s *stateGenerationHandlerStream) SendHeader(metadata.MD) error { return nil }
+func (s *stateGenerationHandlerStream) SetTrailer(metadata.MD)       {}
+func (s *stateGenerationHandlerStream) Context() context.Context     { return s.ctx }
+func (s *stateGenerationHandlerStream) SendMsg(any) error            { return nil }
+func (s *stateGenerationHandlerStream) RecvMsg(any) error            { return nil }
+
+func TestReportSystemState_HandlerOldStreamCannotClearNewerState(t *testing.T) {
+	// Given
+	reporter := requestTaskSecurityServer(7, 200, "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee")
+	setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{
+		200: {Role: model.RoleMember},
+	}, map[string]uint64{"reporter-secret": 200})
+	oldStop := make(chan struct{})
+	newStop := make(chan struct{})
+	oldStream := &stateGenerationHandlerStream{
+		ctx:    metadata.NewIncomingContext(context.Background(), metadata.Pairs("client_secret", "reporter-secret", "client_uuid", reporter.UUID)),
+		states: make(chan *pb.State, 1), receipts: make(chan *pb.Receipt, 1), stop: oldStop,
+	}
+	newStream := &stateGenerationHandlerStream{
+		ctx:    metadata.NewIncomingContext(context.Background(), metadata.Pairs("client_secret", "reporter-secret", "client_uuid", reporter.UUID)),
+		states: make(chan *pb.State, 1), receipts: make(chan *pb.Receipt, 1), stop: newStop,
+	}
+	oldStream.states <- &pb.State{Uptime: 11}
+	newStream.states <- &pb.State{Uptime: 22}
+	handler := NewNezhaHandler()
+	oldDone := make(chan error, 1)
+	newDone := make(chan error, 1)
+	go func() { oldDone <- handler.ReportSystemState(oldStream) }()
+	<-oldStream.receipts
+
+	// When
+	go func() { newDone <- handler.ReportSystemState(newStream) }()
+	<-newStream.receipts
+	close(newStop)
+	require.ErrorIs(t, <-newDone, context.Canceled)
+	close(oldStop)
+	require.ErrorIs(t, <-oldDone, context.Canceled)
+
+	// Then
+	server, ok := singleton.ServerShared.Get(reporter.ID)
+	require.True(t, ok)
+	require.Equal(t, uint64(22), server.State.Uptime)
+	require.True(t, server.LastActive.IsZero())
+}
+
+func TestReportSystemState_OldStreamCannotUpdateNewerGeneration(t *testing.T) {
+	// Given
+	server := &model.Server{}
+	model.InitServer(server)
+	oldStream := stateGenerationStream{}
+	newStream := stateGenerationStream{}
+	oldLease := server.AttachStateStream(oldStream)
+	updateGate := make(chan struct{})
+	updateDone := make(chan bool, 1)
+	oldState := &model.HostState{Uptime: 11}
+	newState := &model.HostState{Uptime: 22}
+	oldTime := time.Unix(100, 0)
+	newTime := time.Unix(200, 0)
+	var waitGroup sync.WaitGroup
+	waitGroup.Add(1)
+	go func() {
+		defer waitGroup.Done()
+		<-updateGate
+		updateDone <- server.UpdateStateIfCurrent(oldLease, oldState, oldTime)
+	}()
+
+	// When
+	newLease := server.AttachStateStream(newStream)
+	close(updateGate)
+	oldUpdateAccepted := <-updateDone
+	newUpdateAccepted := server.UpdateStateIfCurrent(newLease, newState, newTime)
+	waitGroup.Wait()
+
+	// Then
+	require.False(t, oldUpdateAccepted)
+	require.True(t, newUpdateAccepted)
+	require.Equal(t, newState, server.State)
+	require.Equal(t, newTime, server.LastActive)
+}
+
+func TestReportSystemState_OldCleanupCannotClearNewerGeneration(t *testing.T) {
+	// Given
+	server := &model.Server{}
+	model.InitServer(server)
+	oldLease := server.AttachStateStream(stateGenerationStream{})
+	newLease := server.AttachStateStream(stateGenerationStream{})
+	state := &model.HostState{Uptime: 22}
+	activeAt := time.Unix(200, 0)
+	require.True(t, server.UpdateStateIfCurrent(newLease, state, activeAt))
+
+	// When
+	oldCleanup := server.ClearStateStreamIfCurrent(oldLease)
+
+	// Then
+	require.False(t, oldCleanup)
+	require.Equal(t, state, server.State)
+	require.Equal(t, activeAt, server.LastActive)
+}
+
+func TestReportSystemState_CurrentCleanupClearsOnlineVisibility(t *testing.T) {
+	// Given
+	server := &model.Server{}
+	model.InitServer(server)
+	lease := server.AttachStateStream(stateGenerationStream{})
+	activeAt := time.Unix(300, 0)
+	require.True(t, server.UpdateStateIfCurrent(lease, &model.HostState{Uptime: 33}, activeAt))
+
+	// When
+	cleared := server.ClearStateStreamIfCurrent(lease)
+
+	// Then
+	require.True(t, cleared)
+	require.True(t, server.LastActive.IsZero())
+}

service/rpc/wait_for_agent_revoke_test.go

@@ -16,7 +16,9 @@ import (
 func TestWaitForAgent_RevokeWakesUpWaiter(t *testing.T) {
 	h := NewNezhaHandler()
 	const streamID = "kill-switch-wait"
-	h.CreateStreamWithPurpose(streamID, 0, 7, PurposeMCPTransfer)
+	if err := h.CreateStreamWithPurpose(streamID, 0, 7, PurposeMCPTransfer); err != nil {
+		t.Fatalf("create waiter stream: %v", err)
+	}
 
 	done := make(chan struct {
 		io  any
@@ -35,9 +37,16 @@ func TestWaitForAgent_RevokeWakesUpWaiter(t *testing.T) {
 			dur time.Duration
 		}{stream, ok, time.Since(start)}
 	}()
+	waiter, err := h.GetStream(streamID)
+	if err != nil {
+		t.Fatalf("get waiter context: %v", err)
+	}
+	select {
+	case <-waiter.waitStartedCh:
+	case <-time.After(time.Second):
+		t.Fatal("WaitForAgent did not enter its blocking select")
+	}
 
-	// 让 WaitForAgent 真的进入 select 等待，再触发 kill switch。
-	time.Sleep(50 * time.Millisecond)
 	if revoked := h.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 1 {
 		t.Fatalf("expected to revoke exactly 1 MCP stream, got %d", revoked)
 	}
@@ -54,3 +63,46 @@ func TestWaitForAgent_RevokeWakesUpWaiter(t *testing.T) {
 		t.Fatalf("WaitForAgent never returned after revoke; kill switch did not wake the waiter")
 	}
 }
+
+func TestRevokeStreamsForServerWakesWaitForAgentAndPreservesNewGeneration(t *testing.T) {
+	h := NewNezhaHandler()
+	const streamID = "server-revoke-generation"
+	if err := h.CreateStream(streamID, 0, 7); err != nil {
+		t.Fatalf("create waiter stream: %v", err)
+	}
+	done := make(chan bool, 1)
+	go func() {
+		_, ok := h.WaitForAgent(context.Background(), streamID, time.Minute)
+		done <- ok
+	}()
+	waiter, err := h.GetStream(streamID)
+	if err != nil {
+		t.Fatalf("get waiter stream: %v", err)
+	}
+	select {
+	case <-waiter.waitStartedCh:
+	case <-time.After(time.Second):
+		t.Fatal("WaitForAgent did not reach its blocking select")
+	}
+
+	h.RevokeStreamsForServer(7)
+	select {
+	case ok := <-done:
+		if ok {
+			t.Fatal("WaitForAgent must return false after server revocation")
+		}
+	case <-time.After(time.Second):
+		t.Fatal("server revocation did not wake WaitForAgent")
+	}
+	h.RevokeStreamsForServer(7)
+	if err := h.CreateStream(streamID, 0, 8); err != nil {
+		t.Fatalf("new generation must reuse released ID: %v", err)
+	}
+	h.RevokeStreamsForServer(7)
+	if h.StreamCount() != 1 {
+		t.Fatalf("new generation must remain tracked, got %d streams", h.StreamCount())
+	}
+	if err := h.CloseStream(streamID); err != nil {
+		t.Fatalf("cleanup new generation: %v", err)
+	}
+}

service/singleton/crontask.go

@@ -24,6 +24,18 @@ type CronClass struct {
 	*cron.Cron
 	pendingAlertTriggerTasksMu sync.Mutex
 	pendingAlertTriggerTasks   map[uint64]map[uint64][]time.Time
+	closeOnce                  sync.Once
+}
+
+// Close stops the scheduler and joins every job before a test restores globals.
+// The embedded cron.Stop only exposes the completion context; callers must await it.
+func (c *CronClass) Close() {
+	if c == nil || c.Cron == nil {
+		return
+	}
+	c.closeOnce.Do(func() {
+		<-c.Cron.Stop().Done()
+	})
 }
 
 func NewCronClass() *CronClass {

service/singleton/crontask_lifecycle_test.go

@@ -0,0 +1,97 @@
+package singleton
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/robfig/cron/v3"
+	"github.com/stretchr/testify/require"
+)
+
+const lifecycleTestTimeout = time.Second
+
+func TestCronClassClose_waitsForRunningJobs(t *testing.T) {
+	// Given
+	started := make(chan struct{})
+	release := make(chan struct{})
+	events := make(chan string, 9)
+	cronClass := &CronClass{Cron: cron.New(cron.WithSeconds())}
+	_, err := cronClass.AddFunc("@every 1ns", func() {
+		defer func() { events <- "job" }()
+		close(started)
+		<-release
+	})
+	require.NoError(t, err)
+	cronClass.Start()
+	<-started
+
+	// When
+	closed := make(chan struct{})
+	for range 8 {
+		go func() {
+			cronClass.Close()
+			events <- "close"
+			closed <- struct{}{}
+		}()
+	}
+
+	// Then
+	close(release)
+	firstEvent := awaitCronLifecycleEvent(t, events, "cron lifecycle did not complete")
+	if firstEvent != "job" {
+		t.Fatalf("Close returned before the running cron job returned: first event=%q", firstEvent)
+	}
+	for range 8 {
+		awaitCronLifecycleSignal(t, closed, "concurrent Close call did not return")
+	}
+	for range 8 {
+		if event := awaitCronLifecycleEvent(t, events, "concurrent Close call did not complete"); event != "close" {
+			t.Fatalf("unexpected cron lifecycle event: %q", event)
+		}
+	}
+}
+
+func TestCronClassClose_isIdempotentAndNilSafe(t *testing.T) {
+	cronClass := &CronClass{Cron: cron.New(cron.WithSeconds())}
+	cronClass.Start()
+
+	closed := make(chan struct{})
+	for range 8 {
+		go func() {
+			cronClass.Close()
+			closed <- struct{}{}
+		}()
+	}
+	for range 8 {
+		awaitCronLifecycleSignal(t, closed, "concurrent Close call did not return")
+	}
+	cronClass.Close()
+	var nilCronClass *CronClass
+	nilCronClass.Close()
+	(&CronClass{}).Close()
+}
+
+func awaitCronLifecycleSignal(t *testing.T, signal <-chan struct{}, message string) {
+	t.Helper()
+	ctx, cancel := context.WithTimeout(context.Background(), lifecycleTestTimeout)
+	defer cancel()
+	select {
+	case <-signal:
+	case <-ctx.Done():
+		t.Fatal(message)
+	}
+}
+
+func awaitCronLifecycleEvent(t *testing.T, events <-chan string, message string) string {
+	t.Helper()
+	ctx, cancel := context.WithTimeout(context.Background(), lifecycleTestTimeout)
+	defer cancel()
+	select {
+	case event := <-events:
+		return event
+	case <-ctx.Done():
+		t.Fatal(message)
+		return ""
+	}
+}