From 6b4f5f4677532e20bdc599e48d74e12619300601 Mon Sep 17 00:00:00 2001
From: KingKDot <why_are_you_looking_at_this@stoppeeping.com>
Date: Fri, 3 Apr 2026 20:01:34 -0500
Subject: [PATCH] fix issue with writing getting caught up under very rare
 conditions

---
 Overlord-Client/cmd/agent/handlers/ping.go    |  1 -
 Overlord-Client/cmd/agent/session.go          |  1 +
 Overlord-Client/cmd/agent/wire/safe_writer.go | 25 +++++++++++++------
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/Overlord-Client/cmd/agent/handlers/ping.go b/Overlord-Client/cmd/agent/handlers/ping.go
index 721b02a..e7b2cb3 100644
--- a/Overlord-Client/cmd/agent/handlers/ping.go
+++ b/Overlord-Client/cmd/agent/handlers/ping.go
@@ -9,7 +9,6 @@ import (
 )
 
 func HandlePing(ctx context.Context, env *runtime.Env, envelope map[string]interface{}) error {
-	env.SetLastPong(time.Now().UnixMilli())
 
 	ts := extractTimestamp(envelope["ts"])
 	if ts == 0 {
diff --git a/Overlord-Client/cmd/agent/session.go b/Overlord-Client/cmd/agent/session.go
index 77e17ee..52d9fd5 100644
--- a/Overlord-Client/cmd/agent/session.go
+++ b/Overlord-Client/cmd/agent/session.go
@@ -589,6 +589,7 @@ func runSession(ctx context.Context, cancel context.CancelFunc, conn *websocket.
 					last := env.LastPong()
 					if !last.IsZero() && time.Since(last) > grace {
 						log.Printf("ping: no pong for %s, forcing reconnect", time.Since(last))
+						conn.Close(websocket.StatusGoingAway, "pong timeout")
 						cancel()
 						return
 					}
diff --git a/Overlord-Client/cmd/agent/wire/safe_writer.go b/Overlord-Client/cmd/agent/wire/safe_writer.go
index 146401b..b68d6fa 100644
--- a/Overlord-Client/cmd/agent/wire/safe_writer.go
+++ b/Overlord-Client/cmd/agent/wire/safe_writer.go
@@ -2,22 +2,33 @@ package wire
 
 import (
 	"context"
-	"sync"
+	"time"
 
 	"nhooyr.io/websocket"
 )
 
+const writeTimeout = 30 * time.Second
+
 type SafeWriter struct {
-	mu sync.Mutex
-	w  Writer
+	sem chan struct{}
+	w   Writer
 }
 
 func NewSafeWriter(w Writer) *SafeWriter {
-	return &SafeWriter{w: w}
+	sem := make(chan struct{}, 1)
+	sem <- struct{}{}
+	return &SafeWriter{sem: sem, w: w}
 }
 
 func (s *SafeWriter) Write(ctx context.Context, messageType websocket.MessageType, p []byte) error {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	return s.w.Write(ctx, messageType, p)
+	select {
+	case <-s.sem:
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+	defer func() { s.sem <- struct{}{} }()
+
+	writeCtx, cancel := context.WithTimeout(ctx, writeTimeout)
+	defer cancel()
+	return s.w.Write(writeCtx, messageType, p)
 }