From 79a4b7a6a4fc9df513f5ca25ddc0cfc503ccc5dc Mon Sep 17 00:00:00 2001 From: Eve Date: Mon, 13 Jul 2026 15:59:20 +0800 Subject: [PATCH] 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 --- server/internal/daemon/wsrpc.go | 13 +-- server/internal/daemon/wsrpc_test.go | 29 ++++++ server/internal/daemonws/hub.go | 62 +++++++++---- server/internal/daemonws/rpc_dispatch_test.go | 39 ++++++++ server/internal/handler/daemon.go | 88 ++++++++++++------- .../daemon_batch_claim_finalize_test.go | 76 ++++++++++++++++ 6 files changed, 254 insertions(+), 53 deletions(-) diff --git a/server/internal/daemon/wsrpc.go b/server/internal/daemon/wsrpc.go index 0c81072037..2d0ee2ccb2 100644 --- a/server/internal/daemon/wsrpc.go +++ b/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 { diff --git a/server/internal/daemon/wsrpc_test.go b/server/internal/daemon/wsrpc_test.go index 9de7cab257..1e69277ae4 100644 --- a/server/internal/daemon/wsrpc_test.go +++ b/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() + } +} diff --git a/server/internal/daemonws/hub.go b/server/internal/daemonws/hub.go index 95581573d9..0dec09cdfd 100644 --- a/server/internal/daemonws/hub.go +++ b/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) } diff --git a/server/internal/daemonws/rpc_dispatch_test.go b/server/internal/daemonws/rpc_dispatch_test.go index 2e50a2780c..8844bdcb7f 100644 --- a/server/internal/daemonws/rpc_dispatch_test.go +++ b/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. +} diff --git a/server/internal/handler/daemon.go b/server/internal/handler/daemon.go index 81a5e7a9d5..dc675fc159 100644 --- a/server/internal/handler/daemon.go +++ b/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" diff --git a/server/internal/handler/daemon_batch_claim_finalize_test.go b/server/internal/handler/daemon_batch_claim_finalize_test.go index 9b4c018fa0..7068e3df13 100644 --- a/server/internal/handler/daemon_batch_claim_finalize_test.go +++ b/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) + } +}