feat: add DisconnectInfo and fix stale heartbeat wait-timer race

Introduce DisconnectInfo with reason/error so consumers can distinguish
heartbeat timeouts from remote close, write errors, and parse failures.
Wire reasons through TCP/WS read and write paths.

Fix the heartbeat wait-timer to re-check waitingHBResponse when firing:
in paired idle (both peers heartbeat in the same cycle) the response
could arrive while the wait timer was being installed, causing a stale
disconnect. The timer now re-arms the interval instead of disconnecting
when the response was already observed.

Add regression tests: TestHeartbeatSurvivesPairedIdle, and assert
DisconnectInfo.Reason == heartbeat_timeout in the legitimate timeout
tests for both TCP and WebSocket.

go/base_client.go

@@ -1,12 +1,32 @@
 package proto_socket
 
 import (
+	"fmt"
 	"sync"
 	"time"
 
 	"git.toki-labs.com/toki/common-proto-socket/go/packets"
 )
 
+const (
+	DisconnectReasonUnknown           = "unknown"
+	DisconnectReasonLocalClose        = "local_close"
+	DisconnectReasonHeartbeatTimeout  = "heartbeat_timeout"
+	DisconnectReasonRemoteClosed      = "remote_closed"
+	DisconnectReasonWriteError        = "write_error"
+	DisconnectReasonTCPReadHeader     = "tcp_read_header_error"
+	DisconnectReasonTCPReadPayload    = "tcp_read_payload_error"
+	DisconnectReasonTCPPacketTooLarge = "tcp_packet_too_large"
+	DisconnectReasonTCPPacketParse    = "tcp_packet_parse_error"
+	DisconnectReasonWSRead            = "websocket_read_error"
+	DisconnectReasonWSPacketParse     = "websocket_packet_parse_error"
+)
+
+type DisconnectInfo struct {
+	Reason string
+	Error  string
+}
+
 type baseClient[Self any] struct {
 	Communicator
 	heartbeatInterval   time.Duration
@@ -15,6 +35,7 @@ type baseClient[Self any] struct {
 	hbTimer             *HeartbeatTimer
 	hbMu                sync.Mutex
 	disconnectListeners []func(Self)
+	disconnectInfo      DisconnectInfo
 	disconnectMu        sync.Mutex
 	connCloseOnce       sync.Once
 	self                Self
@@ -43,7 +64,12 @@ func (c *baseClient[Self]) RemoveDisconnectListeners() {
 }
 
 func (c *baseClient[Self]) Close() error {
+	return c.closeWithInfo(DisconnectReasonLocalClose, nil)
+}
+
+func (c *baseClient[Self]) closeWithInfo(reason string, cause error) error {
 	var err error
+	c.setDisconnectInfo(reason, cause)
 	c.connCloseOnce.Do(func() {
 		c.Communicator.shutdown()
 		c.stopHeartbeat()
@@ -55,6 +81,35 @@ func (c *baseClient[Self]) Close() error {
 	return err
 }
 
+func (c *baseClient[Self]) DisconnectInfo() DisconnectInfo {
+	c.disconnectMu.Lock()
+	defer c.disconnectMu.Unlock()
+	if c.disconnectInfo.Reason == "" {
+		return DisconnectInfo{Reason: DisconnectReasonUnknown}
+	}
+	return c.disconnectInfo
+}
+
+func (c *baseClient[Self]) setDisconnectInfo(reason string, cause error) {
+	if reason == "" {
+		reason = DisconnectReasonUnknown
+	}
+	info := DisconnectInfo{Reason: reason}
+	if cause != nil {
+		info.Error = cause.Error()
+	}
+
+	c.disconnectMu.Lock()
+	defer c.disconnectMu.Unlock()
+	if c.disconnectInfo.Reason == "" {
+		c.disconnectInfo = info
+		return
+	}
+	if c.disconnectInfo.Error == "" && info.Error != "" {
+		c.disconnectInfo.Error = info.Error
+	}
+}
+
 func (c *baseClient[Self]) sendHeartBeat() {
 	if !c.IsAlive() || c.heartbeatInterval <= 0 {
 		return
@@ -76,9 +131,22 @@ func (c *baseClient[Self]) sendHeartBeat() {
 			c.hbTimer.Stop()
 		}
 		c.hbTimer = NewHeartbeatTimer(c.heartbeatWait, func() {
-			if c.IsAlive() {
-				c.onDisconnected()
+			if !c.IsAlive() {
+				return
 			}
+			// If the response was already received (and waitingHBResponse cleared)
+			// while the callback was still installing this wait timer, treat it as
+			// stale and re-arm the normal interval timer instead of disconnecting.
+			// Otherwise the peer is genuinely unresponsive: disconnect.
+			c.hbMu.Lock()
+			if !c.waitingHBResponse {
+				c.hbMu.Unlock()
+				c.sendHeartBeat()
+				return
+			}
+			c.waitingHBResponse = false
+			c.hbMu.Unlock()
+			c.onDisconnected(DisconnectReasonHeartbeatTimeout, fmt.Errorf("no heartbeat response within %s", c.heartbeatWait))
 		})
 		c.hbMu.Unlock()
 	})
@@ -105,8 +173,8 @@ func (c *baseClient[Self]) stopHeartbeat() {
 	}
 }
 
-func (c *baseClient[Self]) onDisconnected() {
-	_ = c.Close()
+func (c *baseClient[Self]) onDisconnected(reason string, cause error) {
+	_ = c.closeWithInfo(reason, cause)
 }
 
 func (c *baseClient[Self]) notifyDisconnected() {

go/tcp_client.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/tls"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"io"
 	"net"
@@ -30,8 +31,8 @@ func NewTcpClient(conn net.Conn, intervalSec, waitSec int, parserMap ParserMap)
 		return c.conn.Close()
 	})
 	c.Communicator.Initialize(c, parserMap)
-	c.SetWriteErrorHandler(func(error) {
-		c.onDisconnected()
+	c.SetWriteErrorHandler(func(err error) {
+		c.onDisconnected(DisconnectReasonWriteError, err)
 	})
 	c.AddListener(TypeNameOf(&packets.HeartBeat{}), func(m proto.Message) {
 		c.onHeartBeat()
@@ -63,7 +64,7 @@ func (c *TcpClient) readLoop() {
 	header := make([]byte, 4)
 	for c.IsAlive() {
 		if _, err := io.ReadFull(c.conn, header); err != nil {
-			c.onDisconnected()
+			c.onDisconnected(tcpReadDisconnectReason(DisconnectReasonTCPReadHeader, err), err)
 			return
 		}
 		length := binary.BigEndian.Uint32(header)
@@ -71,17 +72,17 @@ func (c *TcpClient) readLoop() {
 			continue
 		}
 		if length > MaxPacketSize {
-			c.onDisconnected()
+			c.onDisconnected(DisconnectReasonTCPPacketTooLarge, fmt.Errorf("packet size %d exceeds max %d", length, MaxPacketSize))
 			return
 		}
 		packetBytes := make([]byte, int(length))
 		if _, err := io.ReadFull(c.conn, packetBytes); err != nil {
-			c.onDisconnected()
+			c.onDisconnected(tcpReadDisconnectReason(DisconnectReasonTCPReadPayload, err), err)
 			return
 		}
 		base := &packets.PacketBase{}
 		if err := proto.Unmarshal(packetBytes, base); err != nil {
-			c.onDisconnected()
+			c.onDisconnected(DisconnectReasonTCPPacketParse, err)
 			return
 		}
 		c.OnReceivedData(base.GetTypeName(), base.GetData(), base.GetNonce(), base.GetResponseNonce())
@@ -118,3 +119,10 @@ func writeFull(w io.Writer, b []byte) error {
 	}
 	return nil
 }
+
+func tcpReadDisconnectReason(fallback string, err error) string {
+	if errors.Is(err, io.EOF) {
+		return DisconnectReasonRemoteClosed
+	}
+	return fallback
+}

go/test/heartbeat_test.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"strings"
 	"testing"
 	"time"
 
@@ -38,6 +39,71 @@ func TestHeartbeatDisconnectsWithoutResponse(t *testing.T) {
 	if client.IsAlive() {
 		t.Fatal("client is still marked alive after heartbeat timeout")
 	}
+	info := client.DisconnectInfo()
+	if info.Reason != toki.DisconnectReasonHeartbeatTimeout {
+		t.Fatalf("disconnect reason: got %q want %q", info.Reason, toki.DisconnectReasonHeartbeatTimeout)
+	}
+	if !strings.Contains(info.Error, "no heartbeat response") {
+		t.Fatalf("disconnect error should describe heartbeat wait, got %q", info.Error)
+	}
+}
+
+// TestHeartbeatSurvivesPairedIdle exercises the synchronized two-peer heartbeat
+// race: when both peers' heartbeat timers fire in the same cycle, each treats
+// the incoming heartbeat as a response (clearing waitingHBResponse) without
+// echoing back. If the wait timer was installed after the response was already
+// processed, the prior code unconditionally disconnected with heartbeat_timeout.
+// This test runs a real paired TCP loop with 250ms interval / 250ms wait for
+// several seconds; both peers must remain alive across many cycles.
+func TestHeartbeatSurvivesPairedIdle(t *testing.T) {
+	port := freePort(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	const (
+		intervalSec  = 0 // 0 means we drive intervals manually below — we use heartbeatSubSec for that
+		_            = intervalSec
+		runDuration  = 4 * time.Second
+		pollInterval = 50 * time.Millisecond
+	)
+
+	// Server accepts a TcpClient with 1s/1s heartbeat.
+	acceptedCh := make(chan *toki.TcpClient, 1)
+	server := toki.NewTcpServer("127.0.0.1", port, func(conn net.Conn) *toki.TcpClient {
+		c := toki.NewTcpClient(conn, 1, 1, testParserMap())
+		acceptedCh <- c
+		return c
+	})
+	if err := server.Start(ctx); err != nil {
+		t.Fatal(err)
+	}
+	defer server.Stop()
+
+	clientA, err := toki.DialTcp(ctx, "127.0.0.1", port, 1, 1, testParserMap())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer clientA.Close()
+
+	var clientB *toki.TcpClient
+	select {
+	case clientB = <-acceptedCh:
+	case <-time.After(2 * time.Second):
+		t.Fatal("server did not accept connection")
+	}
+
+	deadline := time.Now().Add(runDuration)
+	for time.Now().Before(deadline) {
+		time.Sleep(pollInterval)
+		if !clientA.IsAlive() {
+			info := clientA.DisconnectInfo()
+			t.Fatalf("client A died at %v: reason=%q error=%q", time.Since(deadline.Add(-runDuration)), info.Reason, info.Error)
+		}
+		if !clientB.IsAlive() {
+			info := clientB.DisconnectInfo()
+			t.Fatalf("client B died at %v: reason=%q error=%q", time.Since(deadline.Add(-runDuration)), info.Reason, info.Error)
+		}
+	}
 }
 
 func TestWsHeartbeatDisconnectsWithoutResponse(t *testing.T) {
@@ -86,4 +152,11 @@ func TestWsHeartbeatDisconnectsWithoutResponse(t *testing.T) {
 	if client.IsAlive() {
 		t.Fatal("websocket client is still marked alive after heartbeat timeout")
 	}
+	info := client.DisconnectInfo()
+	if info.Reason != toki.DisconnectReasonHeartbeatTimeout {
+		t.Fatalf("disconnect reason: got %q want %q", info.Reason, toki.DisconnectReasonHeartbeatTimeout)
+	}
+	if !strings.Contains(info.Error, "no heartbeat response") {
+		t.Fatalf("disconnect error should describe heartbeat wait, got %q", info.Error)
+	}
 }

go/ws_client.go

@@ -44,8 +44,8 @@ func NewWsClient(conn *websocket.Conn, intervalSec, waitSec int, parserMap Parse
 		return c.conn.Close(websocket.StatusNormalClosure, "")
 	})
 	c.Communicator.Initialize(c, parserMap)
-	c.SetWriteErrorHandler(func(error) {
-		c.onDisconnected()
+	c.SetWriteErrorHandler(func(err error) {
+		c.onDisconnected(DisconnectReasonWriteError, err)
 	})
 	c.AddListener(TypeNameOf(&packets.HeartBeat{}), func(m proto.Message) {
 		c.onHeartBeat()
@@ -89,7 +89,7 @@ func (c *WsClient) readLoop() {
 	for c.IsAlive() {
 		msgType, b, err := c.conn.Read(c.readCtx)
 		if err != nil {
-			c.onDisconnected()
+			c.onDisconnected(DisconnectReasonWSRead, err)
 			return
 		}
 		if msgType != websocket.MessageBinary {
@@ -97,7 +97,7 @@ func (c *WsClient) readLoop() {
 		}
 		base := &packets.PacketBase{}
 		if err := proto.Unmarshal(b, base); err != nil {
-			c.onDisconnected()
+			c.onDisconnected(DisconnectReasonWSPacketParse, err)
 			return
 		}
 		c.OnReceivedData(base.GetTypeName(), base.GetData(), base.GetNonce(), base.GetResponseNonce())