From 9c9c085690d9c70364a8657b1b6205cc783ee9a1 Mon Sep 17 00:00:00 2001
From: Siyuan <siyuan@buildjet.com>
Date: Tue, 18 Nov 2025 09:39:02 +0000
Subject: [PATCH] chore: move setStream to handleEventStream

---
 internal/native/grpc_client.go | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/internal/native/grpc_client.go b/internal/native/grpc_client.go
index 1029e06f..c28ceb6c 100644
--- a/internal/native/grpc_client.go
+++ b/internal/native/grpc_client.go
@@ -75,13 +75,14 @@ func NewGRPCClient(opts grpcClientOptions) (*GRPCClient, error) {
 	return grpcClient, nil
 }
 
-func (c *GRPCClient) setStream(stream pb.NativeService_StreamEventsClient) {
-	c.eventM.Lock()
-	defer c.eventM.Unlock()
-	c.eventStream = stream
-}
-
 func (c *GRPCClient) handleEventStream(stream pb.NativeService_StreamEventsClient) {
+	c.eventM.Lock()
+	c.eventStream = stream
+	defer func() {
+		c.eventStream = nil
+		c.eventM.Unlock()
+	}()
+
 	logger := *c.logger
 	for {
 		if stream == nil {
@@ -135,9 +136,7 @@ func (c *GRPCClient) startEventStream() {
 			continue
 		}
 
-		c.setStream(stream)
 		c.handleEventStream(stream)
-		c.setStream(nil)
 
 		// Wait before retrying
 		time.Sleep(1 * time.Second)