From b23c4c61f1aa6d625d75c0e7941489647a10c3b9 Mon Sep 17 00:00:00 2001
From: chenshi5012 <10794822+chenshi5012@users.noreply.github.com>
Date: Wed, 14 Oct 2026 09:55:27 +0000
Subject: [PATCH] Clear broken flag after successful reconnect

---
 src/client/Session.cpp           |  1 +
 src/client/tests/SessionTest.cpp | 31 +++++++++++++++++++++++++++++--
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/client/Session.cpp b/src/client/Session.cpp
index 3dc0f354..e90d51f6 100644
--- a/src/client/Session.cpp
+++ b/src/client/Session.cpp
@@ -25,6 +25,7 @@ ExecutionResponse Session::executeWithParameter(
     const std::string &stmt, const std::unordered_map<std::string, Value> &parameters) {
   if (connectionIsBroken_ && retryConnect_) {
     if (retryConnect() == nebula::ErrorCode::SUCCEEDED) {
+      connectionIsBroken_ = false;
       return ExecutionResponse(conn_.executeWithParameter(sessionId_, stmt, parameters));
     } else {
       return ExecutionResponse{ErrorCode::E_RPC_FAILURE,
diff --git a/src/client/tests/SessionTest.cpp b/src/client/tests/SessionTest.cpp
index 27edb055..a88e174c 100644
--- a/src/client/tests/SessionTest.cpp
+++ b/src/client/tests/SessionTest.cpp
@@ -4,6 +4,7 @@
  */
 
 #include <common/Init.h>
+#include <common/datatypes/DataSet.h>
 #include <common/datatypes/Geography.h>
 #include <folly/json.h>
 #include <folly/synchronization/Baton.h>
@@ -11,13 +12,18 @@
 #include <gtest/gtest.h>
 #include <gtest/gtest_prod.h>
 #include <nebula/client/Config.h>
-#include <nebula/client/ConnectionPool.h>
-#include <nebula/client/Session.h>
+#include <nebula/client/Connection.h>
 
 #include <atomic>
 #include <chrono>
 #include <thread>
 
+#define private public
+#include <nebula/client/Session.h>
+#undef private
+
+#include <nebula/client/ConnectionPool.h>
+
 #include "./ClientTest.h"
 #include "common/graph/Response.h"
 
@@ -101,6 +107,27 @@ TEST_F(SessionTest, Basic) {
   runOnce(pool);
 }
 
+TEST_F(SessionTest, PreflightReconnectClearsBrokenFlag) {
+  nebula::Config config;
+  config.maxConnectionPoolSize_ = 1;
+
+  nebula::ConnectionPool pool;
+  pool.init({kServerHost ":9669"}, config);
+  auto session = pool.getSession("root", "nebula");
+  ASSERT_TRUE(session.valid());
+
+  session.connectionIsBroken_.store(true);
+  auto resp = session.execute("YIELD 1");
+  ASSERT_EQ(resp.errorCode, nebula::ErrorCode::SUCCEEDED);
+  EXPECT_FALSE(session.connectionIsBroken_.load());
+
+  const auto stableSessionId = session.sessionId_;
+  resp = session.execute("YIELD 1");
+  ASSERT_EQ(resp.errorCode, nebula::ErrorCode::SUCCEEDED);
+  EXPECT_EQ(session.sessionId_, stableSessionId);
+  EXPECT_FALSE(session.connectionIsBroken_.load());
+}
+
 TEST_F(SessionTest, OverUse) {
   nebula::ConnectionPool pool;
   nebula::Config c;