From 61604e18179db4d1f987b3da97e8c8e24dfb6954 Mon Sep 17 00:00:00 2001
From: Marvin Zhang <tikazyq>
Date: Mon, 20 Oct 2025 12:22:34 +0800
Subject: [PATCH] fix(task/handler): ensure latest gRPC client is used for task
 fetch/subscribe

Add svc.getGrpcClient() helper and use it when obtaining TaskClient so task fetch and
subscribe operations don't hold a stale client instance after ResetGrpcClient().
---
 core/task/handler/service.go            | 8 +++++++-
 core/task/handler/service_operations.go | 2 +-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/core/task/handler/service.go b/core/task/handler/service.go
index 9829064d..03f6c9d1 100644
--- a/core/task/handler/service.go
+++ b/core/task/handler/service.go
@@ -288,6 +288,12 @@ func (svc *Service) GetNodeConfigService() (cfgSvc interfaces.NodeConfigService)
 	return svc.cfgSvc
 }
 
+// getGrpcClient returns the current gRPC client, ensuring we always have the latest instance
+// This is important after ResetGrpcClient() calls which create new client instances
+func (svc *Service) getGrpcClient() *grpcclient.GrpcClient {
+	return grpcclient.GetGrpcClient()
+}
+
 func (svc *Service) getCurrentNodeMaxRunners() int {
 	n, err := svc.GetCurrentNode()
 	if err != nil {
@@ -450,7 +456,7 @@ func (svc *Service) fetchTask() (tid primitive.ObjectID, err error) {
 	// Use service context with timeout for fetch operation
 	ctx, cancel := context.WithTimeout(svc.ctx, svc.fetchTimeout)
 	defer cancel()
-	taskClient, err := svc.c.GetTaskClient()
+	taskClient, err := svc.getGrpcClient().GetTaskClient()
 	if err != nil {
 		return primitive.NilObjectID, fmt.Errorf("failed to get task client: %v", err)
 	}
diff --git a/core/task/handler/service_operations.go b/core/task/handler/service_operations.go
index 1b6c1b2a..ad00df18 100644
--- a/core/task/handler/service_operations.go
+++ b/core/task/handler/service_operations.go
@@ -98,7 +98,7 @@ func (svc *Service) subscribeTaskWithContext(ctx context.Context, taskId primiti
 	req := &grpc.TaskServiceSubscribeRequest{
 		TaskId: taskId.Hex(),
 	}
-	taskClient, err := svc.c.GetTaskClient()
+	taskClient, err := svc.getGrpcClient().GetTaskClient()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get task client: %v", err)
 	}