fix(daemon-ws): WS RPC disconnect-race panic + batch stale-comment-plan repair (MUL-4257)

Two PR #5193 review blockers:

1) WS RPC send-on-closed-channel race, both ends:
   - server: give each connection a cancelable ctx (cancelled on readPump
     teardown) and run the RPC handler under it, so a slow claim stops on
     disconnect; guard c.send with sendMu/sendClosed (trySend) so a late RPC
     response goroutine never writes to the closed channel. Heartbeat ack routed
     through the same guard.
   - daemon: wsRPCClient.deliver now sends under the mutex, serialized with
     attach(nil)'s close+delete, so a delivered response can't hit a channel
     the detach path just closed.
   - regressions (-race): daemon deliver-vs-detach; server
     disconnect-during-handler-response.

2) batch claim now runs the stale-comment-plan repair: extracted the
   per-runtime handler's repair (trigger deleted, only coalesced survive ->
   cancel + replay survivors) into shared repairStaleCommentPlanIfNeeded, called
   by both claim paths. Prevents the batch path (now the default poller) from
   finalizing+dispatching a task with no comment input and silently dropping the
   surviving user comment. Regression: batch omits the stale task, cancels it,
   and rebuilds the survivor into a new trigger plan.

Co-authored-by: multica-agent <github@multica.ai>

server/internal/daemon/wsrpc.go

@@ -151,16 +151,19 @@ func (c *wsRPCClient) Call(ctx context.Context, method string, reqBody, respBody
 	}
 }
 
-// deliver routes an inbound rpc_response frame to the waiting Call. Unknown
-// request ids (already timed out / detached) are dropped.
+// deliver routes an inbound rpc_response frame to the waiting Call. The send
+// happens under the mutex so it is serialized with attach(nil)'s close+delete:
+// a channel present in pending is guaranteed not yet closed, so this never
+// sends on a closed channel. Unknown request ids (already timed out / detached)
+// are dropped.
 func (c *wsRPCClient) deliver(resp protocol.RPCResponsePayload) {
 	if c == nil {
 		return
 	}
 	c.mu.Lock()
-	ch := c.pending[resp.RequestID]
-	c.mu.Unlock()
-	if ch == nil {
+	defer c.mu.Unlock()
+	ch, ok := c.pending[resp.RequestID]
+	if !ok {
 		return
 	}
 	select {

server/internal/daemon/wsrpc_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"sync"
 	"testing"
 	"time"
 
@@ -108,3 +109,31 @@ func TestWSRPCClient_DetachFailsPending(t *testing.T) {
 		t.Fatal("Call did not return after detach")
 	}
 }
+
+// TestWSRPCClient_DeliverDetachRaceNoPanic hammers deliver racing with
+// attach(nil) (disconnect). Before the fix, deliver could send on a channel
+// attach(nil) had just closed → "send on closed channel" panic. Run under
+// -race; passing means the two are serialized under the mutex.
+func TestWSRPCClient_DeliverDetachRaceNoPanic(t *testing.T) {
+	for iter := 0; iter < 300; iter++ {
+		c := newWSRPCClient(time.Second)
+		c.attach(func([]byte) error { return nil })
+		id := "req"
+		ch := make(chan protocol.RPCResponsePayload, 1)
+		c.mu.Lock()
+		c.pending[id] = ch
+		c.mu.Unlock()
+
+		var wg sync.WaitGroup
+		wg.Add(2)
+		go func() {
+			defer wg.Done()
+			c.deliver(protocol.RPCResponsePayload{RequestID: id, Status: 200})
+		}()
+		go func() {
+			defer wg.Done()
+			c.attach(nil) // closes + deletes pending under the same mutex
+		}()
+		wg.Wait()
+	}
+}

server/internal/daemonws/hub.go

@@ -92,6 +92,18 @@ type client struct {
 	identity ClientIdentity
 	runtimes map[string]struct{}
 
+	// ctx is cancelled when the connection tears down, so async RPC handlers
+	// stop instead of running against a dead socket. cancel is invoked from
+	// readPump's defer.
+	ctx    context.Context
+	cancel context.CancelFunc
+
+	// sendMu guards sendClosed so a late async send (e.g. an RPC response
+	// goroutine) can never write to the closed send channel. teardown flips
+	// sendClosed under sendMu before closing send.
+	sendMu     sync.Mutex
+	sendClosed bool
+
 	dedupMu  sync.Mutex
 	seenIDs  map[string]struct{}
 	seenList []string
@@ -100,6 +112,23 @@ type client struct {
 	rpcSem chan struct{}
 }
 
+// trySend delivers frame to the write pump without blocking and without ever
+// writing to a closed channel (safe against concurrent teardown). Returns false
+// when the buffer is full or the connection is closing.
+func (c *client) trySend(frame []byte) bool {
+	c.sendMu.Lock()
+	defer c.sendMu.Unlock()
+	if c.sendClosed {
+		return false
+	}
+	select {
+	case c.send <- frame:
+		return true
+	default:
+		return false
+	}
+}
+
 const eventDedupCapacity = 128
 
 // markSeen records eventID as already delivered to this client. Empty event IDs
@@ -279,6 +308,7 @@ func (h *Hub) HandleWebSocket(w http.ResponseWriter, r *http.Request, identity C
 		runtimes: runtimes,
 		rpcSem:   make(chan struct{}, maxInFlightRPCPerClient),
 	}
+	c.ctx, c.cancel = context.WithCancel(context.Background())
 	h.register(c)
 
 	go c.writePump()
@@ -523,7 +553,10 @@ func (h *Hub) unregister(c *client) {
 			}
 		}
 	}
+	c.sendMu.Lock()
+	c.sendClosed = true
 	close(c.send)
+	c.sendMu.Unlock()
 	total := len(h.clients)
 	h.mu.Unlock()
 
@@ -541,6 +574,9 @@ func (h *Hub) unregister(c *client) {
 
 func (c *client) readPump() {
 	defer func() {
+		// Cancel first so async RPC handlers stop before we close the send
+		// channel, then unregister (which marks send closed).
+		c.cancel()
 		c.hub.unregister(c)
 		c.conn.Close()
 	}()
@@ -623,9 +659,10 @@ func (c *client) handleRPCFrame(raw json.RawMessage) {
 	}
 	go func() {
 		defer func() { <-c.rpcSem }()
-		// Bound by the connection lifetime; the conn closing unblocks any
-		// wedged handler via the daemon's own request timeout + fallback.
-		status, body, err := handler(context.Background(), c.identity, req.Method, req.Body)
+		// Use the connection ctx so a slow handler (e.g. a DB-bound claim) is
+		// cancelled when the socket tears down, rather than running on against
+		// a dead connection with context.Background().
+		status, body, err := handler(c.ctx, c.identity, req.Method, req.Body)
 		if err != nil {
 			if status < 400 {
 				status = http.StatusInternalServerError
@@ -651,12 +688,10 @@ func (c *client) sendRPCResponse(requestID string, status int, body json.RawMess
 		slog.Debug("daemon websocket rpc response marshal failed", "error", err)
 		return
 	}
-	select {
-	case c.send <- frame:
-	default:
-		// Send buffer full — drop the response; the daemon's per-request
-		// timeout fires and it falls back to HTTP.
-		slog.Debug("daemon websocket rpc response dropped: send buffer full",
+	if !c.trySend(frame) {
+		// Send buffer full or connection closing — drop the response; the
+		// daemon's per-request timeout fires and it falls back to HTTP.
+		slog.Debug("daemon websocket rpc response dropped",
 			"daemon_id", c.identity.DaemonID, "request_id", requestID)
 	}
 }
@@ -715,12 +750,9 @@ func (c *client) handleHeartbeatFrame(raw json.RawMessage) {
 		slog.Debug("daemon websocket heartbeat ack marshal failed", "error", err)
 		return
 	}
-	select {
-	case c.send <- frame:
-	default:
-		// Send buffer is full — slow client. Don't block the read pump; the
-		// next writePump tick or notifyFrame eviction will clean up.
-		slog.Debug("daemon websocket heartbeat ack dropped: send buffer full",
+	if !c.trySend(frame) {
+		// Send buffer full or connection closing — drop; HTTP heartbeat resumes.
+		slog.Debug("daemon websocket heartbeat ack dropped",
 			"daemon_id", c.identity.DaemonID,
 			"runtime_id", payload.RuntimeID)
 	}

server/internal/daemonws/rpc_dispatch_test.go

@@ -114,3 +114,42 @@ func TestRPCDispatch_NoHandler(t *testing.T) {
 		t.Fatalf("resp = %+v, want req-3 with 503", resp)
 	}
 }
+
+// TestRPCDispatch_DisconnectDuringHandlerNoPanic pins the server-side fix: an
+// RPC handler that finishes AFTER the connection tears down must not send on
+// the closed send channel. The handler blocks until the client has
+// disconnected, then returns and attempts to write its response. Run under
+// -race; passing means the guarded send + connection-ctx teardown are safe.
+func TestRPCDispatch_DisconnectDuringHandlerNoPanic(t *testing.T) {
+	hub := NewHub()
+	release := make(chan struct{})
+	entered := make(chan struct{}, 1)
+	hub.SetRPCHandler(func(ctx context.Context, identity ClientIdentity, method string, body json.RawMessage) (int, json.RawMessage, error) {
+		select {
+		case entered <- struct{}{}:
+		default:
+		}
+		<-release // return only after the client disconnects
+		return http.StatusOK, json.RawMessage(`{"ok":true}`), nil
+	})
+	conn := dialRPCTestConn(t, hub, ClientIdentity{DaemonID: "daemon-1", RuntimeIDs: []string{"rt-1"}})
+
+	frame, _ := json.Marshal(protocol.Message{
+		Type:    protocol.EventDaemonRPCRequest,
+		Payload: mustMarshalRaw(protocol.RPCRequestPayload{RequestID: "req-1", Method: "tasks.claim"}),
+	})
+	if err := conn.WriteMessage(websocket.TextMessage, frame); err != nil {
+		t.Fatalf("write: %v", err)
+	}
+	select {
+	case <-entered:
+	case <-time.After(2 * time.Second):
+		t.Fatal("handler never started")
+	}
+
+	conn.Close() // client disconnect → server readPump exits → cancel + close send
+	time.Sleep(50 * time.Millisecond)
+	close(release) // handler returns and tries to send its response
+	time.Sleep(100 * time.Millisecond)
+	// No panic == pass.
+}

server/internal/handler/daemon.go

@@ -1321,6 +1321,45 @@ func parseRuntimeConnectedAppsForClaim(raw []byte, taskID pgtype.UUID) []runtime
 	return apps
 }
 
+// repairStaleCommentPlanIfNeeded handles the edit/delete race where a claimed
+// task's trigger_comment_id was cleared but coalesced_comment_ids survive: such
+// a task must never be dispatched as a generic assignment — its user-scoped MCP
+// overlay still belongs to the deleted author, and the prompt would read issue
+// history exposing that stale user's capabilities. When it applies, the task is
+// cancelled and its surviving comments are replayed through normal routing
+// (which recomputes originator + connected-app context).
+//
+// Returns handled=true when the task must NOT be dispatched — either a clean
+// repair (failure==nil) or a hard failure (failure!=nil, carrying the
+// status/message/outcome the per-runtime endpoint renders). handled=false means
+// proceed with a normal claim. Shared by the per-runtime and batch claim
+// handlers so the batch path can't silently drop surviving comments (MUL-4257).
+func (h *Handler) repairStaleCommentPlanIfNeeded(ctx context.Context, task *db.AgentTaskQueue, runtimeWorkspaceID string) (handled bool, failure *claimBuildFailure) {
+	if task.TriggerCommentID.Valid || len(task.CoalescedCommentIds) == 0 {
+		return false, nil
+	}
+	if !task.IssueID.Valid {
+		return true, &claimBuildFailure{outcome: "error_stale_comment_plan", status: http.StatusInternalServerError, message: "comment task has no issue"}
+	}
+	issue, loadErr := h.Queries.GetIssue(ctx, task.IssueID)
+	if loadErr != nil {
+		return true, &claimBuildFailure{outcome: "error_stale_comment_plan", status: http.StatusInternalServerError, message: "failed to repair stale comment task"}
+	}
+	if uuidToString(issue.WorkspaceID) != runtimeWorkspaceID {
+		if _, cancelErr := h.TaskService.CancelTask(ctx, task.ID); cancelErr != nil {
+			slog.Error("task claim: cancel stale cross-workspace task failed",
+				"task_id", uuidToString(task.ID), "error", cancelErr)
+		}
+		return true, &claimBuildFailure{outcome: "error_workspace", status: http.StatusInternalServerError, message: "task workspace isolation check failed"}
+	}
+	cancelled, cancelErr := h.TaskService.CancelTask(ctx, task.ID)
+	if cancelErr != nil {
+		return true, &claimBuildFailure{outcome: "error_stale_comment_plan", status: http.StatusInternalServerError, message: "failed to repair stale comment task"}
+	}
+	h.retriggerCancelledTaskSurvivors(ctx, issue, []db.AgentTaskQueue{*cancelled}, pgtype.UUID{})
+	return true, nil
+}
+
 // claimBatchMaxTasksCap bounds how many tasks a single machine-level batch
 // claim may return, so one request can neither build an unbounded payload nor
 // hold the DB for an unbounded number of per-agent claim transactions. The
@@ -1455,6 +1494,14 @@ func (h *Handler) ClaimTasksByRuntime(w http.ResponseWriter, r *http.Request) {
 			continue
 		}
 		rtWorkspaceID := uuidToString(rt.WorkspaceID)
+		// Stale comment-plan repair must run for the batch path too: otherwise a
+		// task whose trigger was deleted (only coalesced survive) would be
+		// finalized+dispatched with no comment input, silently dropping the
+		// surviving user comment. On repair (or hard failure) the task is
+		// cancelled / left for reclaim and omitted from the batch.
+		if handled, _ := h.repairStaleCommentPlanIfNeeded(r.Context(), &task, rtWorkspaceID); handled {
+			continue
+		}
 		resp, deliveredCommentIDs, _, _, failure := h.buildClaimedTaskResponse(r, &task, rt, uuidToString(task.RuntimeID), rtWorkspaceID)
 		if failure != nil {
 			// Builder rejected this task (workspace isolation / chat-input);
@@ -2365,42 +2412,17 @@ func (h *Handler) ClaimTaskByRuntime(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	if !task.TriggerCommentID.Valid && len(task.CoalescedCommentIds) > 0 {
-		// A legacy edit/delete race can leave a plan whose trigger FK was
-		// cleared while its user-scoped MCP overlay still belongs to the deleted
-		// author. Never dispatch it as a generic assignment: that prompt reads
-		// issue history and would still expose the stale user's capabilities.
-		// Cancel the stale row and replay every survivor through normal routing,
-		// which recomputes originator and connected-app context.
-		if !task.IssueID.Valid {
-			outcome = "error_stale_comment_plan"
-			writeError(w, http.StatusInternalServerError, "comment task has no issue")
-			return
-		}
-		issue, loadErr := h.Queries.GetIssue(r.Context(), task.IssueID)
-		if loadErr != nil {
-			outcome = "error_stale_comment_plan"
-			writeError(w, http.StatusInternalServerError, "failed to repair stale comment task")
-			return
-		}
-		if uuidToString(issue.WorkspaceID) != runtimeWorkspaceID {
-			outcome = "error_workspace"
-			if _, cancelErr := h.TaskService.CancelTask(r.Context(), task.ID); cancelErr != nil {
-				slog.Error("task claim: cancel stale cross-workspace task failed",
-					"task_id", uuidToString(task.ID), "error", cancelErr)
+		handled, failure := h.repairStaleCommentPlanIfNeeded(r.Context(), task, runtimeWorkspaceID)
+		if handled {
+			if failure != nil {
+				outcome = failure.outcome
+				writeError(w, failure.status, failure.message)
+				return
 			}
-			writeError(w, http.StatusInternalServerError, "task workspace isolation check failed")
+			outcome = "repaired_stale_comment_plan"
+			payloadBytes, _ = writeMeasuredJSON(w, http.StatusOK, map[string]any{"task": nil})
 			return
 		}
-		cancelled, cancelErr := h.TaskService.CancelTask(r.Context(), task.ID)
-		if cancelErr != nil {
-			outcome = "error_stale_comment_plan"
-			writeError(w, http.StatusInternalServerError, "failed to repair stale comment task")
-			return
-		}
-		h.retriggerCancelledTaskSurvivors(r.Context(), issue, []db.AgentTaskQueue{*cancelled}, pgtype.UUID{})
-		outcome = "repaired_stale_comment_plan"
-		payloadBytes, _ = writeMeasuredJSON(w, http.StatusOK, map[string]any{"task": nil})
-		return
 	}
 
 	outcome = "claimed"

server/internal/handler/daemon_batch_claim_finalize_test.go

@@ -271,3 +271,79 @@ func TestClaimTasksByRuntime_RequiresDaemonID(t *testing.T) {
 		t.Fatalf("expected 400 when daemon_id is missing, got %d: %s", w.Code, w.Body.String())
 	}
 }
+
+// TestClaimTasksByRuntime_RepairsStaleCommentPlan pins the MUL-4257 review
+// must-fix: when a claimed task's trigger comment was deleted (only coalesced
+// survivors remain), the batch path must NOT finalize+dispatch it (which would
+// silently drop the surviving comment). Instead it cancels the stale task,
+// omits it from the batch, and replays the surviving comment as a fresh plan.
+func TestClaimTasksByRuntime_RepairsStaleCommentPlan(t *testing.T) {
+	if testHandler == nil || testPool == nil {
+		t.Skip("database not available")
+	}
+	ctx := context.Background()
+	rt := createClaimReclaimRuntime(t, ctx, "Stale plan rt")
+	agentID, issueID := createClaimReclaimAgentAndIssue(t, ctx, rt, "Stale plan agent")
+	// Assign the issue to the agent so the surviving comment re-routes to it.
+	if _, err := testPool.Exec(ctx, `UPDATE issue SET assignee_type='agent', assignee_id=$1 WHERE id=$2`, agentID, issueID); err != nil {
+		t.Fatalf("assign issue: %v", err)
+	}
+
+	// A surviving member comment on the issue.
+	var survivorID string
+	if err := testPool.QueryRow(ctx, `
+		INSERT INTO comment (workspace_id, issue_id, author_type, author_id, content)
+		VALUES ($1, $2, 'member', $3, 'please still handle this')
+		RETURNING id
+	`, testWorkspaceID, issueID, testUserID).Scan(&survivorID); err != nil {
+		t.Fatalf("seed survivor comment: %v", err)
+	}
+	t.Cleanup(func() { testPool.Exec(ctx, `DELETE FROM comment WHERE id = $1`, survivorID) })
+
+	// Stale plan: trigger_comment_id NULL, only coalesced survivor remains.
+	var staleID string
+	if err := testPool.QueryRow(ctx, `
+		INSERT INTO agent_task_queue (agent_id, runtime_id, issue_id, status, priority, coalesced_comment_ids)
+		VALUES ($1, $2, $3, 'queued', 0, ARRAY[$4]::uuid[])
+		RETURNING id
+	`, agentID, rt, issueID, survivorID).Scan(&staleID); err != nil {
+		t.Fatalf("seed stale task: %v", err)
+	}
+	t.Cleanup(func() { testPool.Exec(ctx, `DELETE FROM agent_task_queue WHERE id = $1`, staleID) })
+	t.Cleanup(func() { testPool.Exec(ctx, `DELETE FROM agent_task_queue WHERE issue_id = $1`, issueID) })
+
+	w := postBatchClaim(t, testWorkspaceID, []string{rt}, 5)
+	if w.Code != http.StatusOK {
+		t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
+	}
+	var resp batchClaimReceiptResponse
+	if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+		t.Fatalf("decode: %v", err)
+	}
+	// (1) The stale task must not be returned.
+	for _, task := range resp.Tasks {
+		if task.ID == staleID {
+			t.Fatalf("stale-plan task %s was dispatched by the batch path; want it repaired/omitted", staleID)
+		}
+	}
+	// (2) The stale task must be cancelled.
+	var status string
+	if err := testPool.QueryRow(ctx, `SELECT status FROM agent_task_queue WHERE id = $1`, staleID).Scan(&status); err != nil {
+		t.Fatalf("read stale status: %v", err)
+	}
+	if status != "cancelled" {
+		t.Fatalf("stale task status = %s, want cancelled", status)
+	}
+	// (3) The surviving comment was rebuilt into a fresh plan (a new task with
+	// the survivor as its trigger).
+	var rebuilt int
+	if err := testPool.QueryRow(ctx, `
+		SELECT count(*) FROM agent_task_queue
+		WHERE issue_id = $1 AND trigger_comment_id = $2 AND id <> $3
+	`, issueID, survivorID, staleID).Scan(&rebuilt); err != nil {
+		t.Fatalf("count rebuilt: %v", err)
+	}
+	if rebuilt < 1 {
+		t.Fatalf("expected the surviving comment rebuilt into a new trigger task, found %d", rebuilt)
+	}
+}