From 0c36588aa1e071121dc2fd6e41dd52915def1a06 Mon Sep 17 00:00:00 2001
From: Nick Chisiu <8492343+nickchisiu@users.noreply.github.com>
Date: Mon, 31 Aug 2026 18:16:01 +0300
Subject: [PATCH] fix(codegen): reconnect Go GraphQL subscriptions

---
 .../codegen/__tests__/gql-codegen.test.ts     | 12 +++++
 .../languages/go/templates/graphql/client.ejs | 54 ++++++++++---------
 2 files changed, 41 insertions(+), 25 deletions(-)

diff --git a/packages/codegen/__tests__/gql-codegen.test.ts b/packages/codegen/__tests__/gql-codegen.test.ts
index 6ed92ae..e8978dd 100644
--- a/packages/codegen/__tests__/gql-codegen.test.ts
+++ b/packages/codegen/__tests__/gql-codegen.test.ts
@@ -234,6 +234,18 @@ describe('GraphQL Codegen — All Languages', () => {
           expect(client.content).not.toContain('func (c *Gql) Pets(');
         });
 
+        it('retries subscriptions when the shared WebSocket closes during a write', async () => {
+          const files = await generateForLanguage(language);
+          const client = getFile(files, 'gql-client')!;
+          expect(client.content).toContain('wasCurrent := c.wsConn == conn');
+          expect(client.content).toContain(
+            'shouldReconnect := wasCurrent && c.reconnect && !c.disposed',
+          );
+          expect(client.content).toContain('subscribeErr = conn.WriteJSON(');
+          expect(client.content).toContain('if c.wsConn == conn {');
+          expect(client.content).toContain('c.wsConn = nil');
+        });
+
         it('generates unified result types instead of per-operation types', async () => {
           const files = await generateForLanguage(language);
           const types = getFile(files, 'gql-types')!;
diff --git a/packages/codegen/src/languages/go/templates/graphql/client.ejs b/packages/codegen/src/languages/go/templates/graphql/client.ejs
index 4234410..b78bf53 100644
--- a/packages/codegen/src/languages/go/templates/graphql/client.ejs
+++ b/packages/codegen/src/languages/go/templates/graphql/client.ejs
@@ -290,8 +290,9 @@ func (c *<%= it.clientClass %>) readWsMessages(conn *websocket.Conn) {
 		_, message, err := conn.ReadMessage()
 		if err != nil {
 			c.mu.Lock()
-			if c.wsConn == conn { c.wsConn = nil }
-			shouldReconnect := c.reconnect && !c.disposed && len(c.subs) > 0
+			wasCurrent := c.wsConn == conn
+			if wasCurrent { c.wsConn = nil }
+			shouldReconnect := wasCurrent && c.reconnect && !c.disposed && len(c.subs) > 0
 			c.mu.Unlock()
 			if shouldReconnect { go c.reconnectSubscriptions() }
 			return
@@ -372,16 +373,6 @@ func (c *<%= it.clientClass %>) Subscribe(fn func(*SubscriptionBuilder) *Subscri
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	var connectErr error
-	for attempt := 0; ; attempt++ {
-		connectErr = c.ensureWsConn()
-		if connectErr == nil { break }
-		if !c.reconnect || attempt >= c.maxReconnectAttempts {
-			return nil, connectErr
-		}
-		time.Sleep(c.reconnectInterval)
-	}
-
 	id := fmt.Sprintf("%d", c.nextSubID)
 	c.nextSubID++
 
@@ -395,19 +386,32 @@ func (c *<%= it.clientClass %>) Subscribe(fn func(*SubscriptionBuilder) *Subscri
 		handler(&envelope.Data)
 	}}
 
-	c.writeMu.Lock()
-	err := c.wsConn.WriteJSON(map[string]interface{}{
-		"type": "subscribe",
-		"id":   id,
-		"payload": map[string]interface{}{
-			"query":     query,
-			"variables": variables,
-		},
-	})
-	c.writeMu.Unlock()
-	if err != nil {
-		delete(c.subs, id)
-		return nil, fmt.Errorf("ws subscribe: %w", err)
+	var subscribeErr error
+	for attempt := 0; ; attempt++ {
+		subscribeErr = c.ensureWsConn()
+		if subscribeErr == nil {
+			conn := c.wsConn
+			c.writeMu.Lock()
+			subscribeErr = conn.WriteJSON(map[string]interface{}{
+				"type": "subscribe",
+				"id":   id,
+				"payload": map[string]interface{}{
+					"query":     query,
+					"variables": variables,
+				},
+			})
+			c.writeMu.Unlock()
+			if subscribeErr == nil { break }
+			if c.wsConn == conn {
+				_ = conn.Close()
+				c.wsConn = nil
+			}
+		}
+		if !c.reconnect || attempt >= c.maxReconnectAttempts {
+			delete(c.subs, id)
+			return nil, fmt.Errorf("ws subscribe: %w", subscribeErr)
+		}
+		time.Sleep(c.reconnectInterval)
 	}
 
 	return func() {