From aa2682860e2c23aa8f37a05b5c0d344e7fc33269 Mon Sep 17 00:00:00 2001 From: Eve Date: Mon, 13 Jul 2026 15:24:03 +0800 Subject: [PATCH] feat(daemon-ws): generic WS request/response transport for daemon RPC (MUL-4257) Add a generic daemon->server request/response layer over the existing WS control connection, the transport for WS-first claim (HTTP fallback): - protocol: daemon:rpc_request / daemon:rpc_response envelopes with a correlation request_id + method + body, and an rpc-v1 capability gate. - daemonws.Hub: SetRPCHandler + goroutine-dispatched handleRPCFrame (bounded by a per-connection in-flight cap) that echoes the request_id; missing handler / saturation return non-2xx so the daemon falls back to HTTP. Read limit raised to 64KB for rpc requests carrying a runtime set. - hub tests: round-trip, handler-error->non-2xx, no-handler->503. Co-authored-by: multica-agent --- server/internal/daemonws/hub.go | 111 ++++++++++++++++- server/internal/daemonws/rpc_dispatch_test.go | 116 ++++++++++++++++++ server/pkg/protocol/events.go | 7 ++ server/pkg/protocol/messages.go | 27 ++++ 4 files changed, 260 insertions(+), 1 deletion(-) create mode 100644 server/internal/daemonws/rpc_dispatch_test.go diff --git a/server/internal/daemonws/hub.go b/server/internal/daemonws/hub.go index 6e240bb84e..64703e2256 100644 --- a/server/internal/daemonws/hub.go +++ b/server/internal/daemonws/hub.go @@ -92,6 +92,9 @@ type client struct { dedupMu sync.Mutex seenIDs map[string]struct{} seenList []string + + // rpcSem bounds concurrent RPC handlers for this connection. + rpcSem chan struct{} } const eventDedupCapacity = 128 @@ -126,6 +129,18 @@ func (c *client) markSeen(eventID string) bool { // the ack and is logged at debug level. type HeartbeatHandler func(ctx context.Context, identity ClientIdentity, runtimeID string, supportsBatchImport bool) (*protocol.DaemonHeartbeatAckPayload, error) +// RPCHandler processes a generic daemon:rpc_request (MUL-4257). It dispatches +// on method (e.g. "tasks.claim"), scoping work to identity (DaemonID + +// authenticated RuntimeIDs), and returns an HTTP-style status plus a response +// body OR an error. A returned error is surfaced to the daemon as a non-2xx +// RPC response so it can fall back to HTTP. The handler runs in its own +// goroutine, so it must not assume it owns the read pump. +type RPCHandler func(ctx context.Context, identity ClientIdentity, method string, body json.RawMessage) (status int, respBody json.RawMessage, err error) + +// maxInFlightRPCPerClient bounds concurrent RPC handlers per connection so a +// single daemon cannot fan out unbounded goroutines / DB work over one socket. +const maxInFlightRPCPerClient = 8 + // MessageKindRecorder is the optional metric hook called once per inbound // daemon WebSocket frame. kind is the protocol message type with the // "daemon:" prefix stripped (e.g. "heartbeat") or the literal "unknown" for @@ -147,6 +162,9 @@ type Hub struct { hbMu sync.RWMutex onHeartbeat HeartbeatHandler + rpcMu sync.RWMutex + onRPC RPCHandler + kindMu sync.RWMutex kindRecorder MessageKindRecorder } @@ -187,6 +205,24 @@ func (h *Hub) heartbeatHandler() HeartbeatHandler { return h.onHeartbeat } +// SetRPCHandler installs the callback used for daemon:rpc_request frames +// (MUL-4257). Like SetHeartbeatHandler it is wired after handler construction. +// A nil handler disables WS RPC — daemons fall back to the HTTP claim endpoint. +func (h *Hub) SetRPCHandler(fn RPCHandler) { + if h == nil { + return + } + h.rpcMu.Lock() + h.onRPC = fn + h.rpcMu.Unlock() +} + +func (h *Hub) rpcHandler() RPCHandler { + h.rpcMu.RLock() + defer h.rpcMu.RUnlock() + return h.onRPC +} + // SetMessageKindRecorder installs an optional callback fired exactly once per // inbound daemon WebSocket frame. Used by the metrics layer to count traffic // by handler kind without hard-coupling the hub to any specific collector. @@ -238,6 +274,7 @@ func (h *Hub) HandleWebSocket(w http.ResponseWriter, r *http.Request, identity C send: make(chan []byte, 16), identity: identity, runtimes: runtimes, + rpcSem: make(chan struct{}, maxInFlightRPCPerClient), } h.register(c) @@ -505,7 +542,9 @@ func (c *client) readPump() { c.conn.Close() }() - c.conn.SetReadLimit(4096) + // Read limit sized for daemon:rpc_request frames carrying a machine's full + // runtime_id set (MUL-4257), well above the tiny heartbeat/wakeup frames. + c.conn.SetReadLimit(64 * 1024) c.conn.SetReadDeadline(time.Now().Add(pongWait)) c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(pongWait)) @@ -543,12 +582,82 @@ func (c *client) handleFrame(raw []byte) { switch msg.Type { case protocol.EventDaemonHeartbeat: c.handleHeartbeatFrame(msg.Payload) + case protocol.EventDaemonRPCRequest: + c.handleRPCFrame(msg.Payload) default: // Unknown app messages are intentionally ignored for forward // compatibility with future daemon → server message types. } } +// handleRPCFrame processes a generic daemon:rpc_request (MUL-4257): it runs the +// registered RPC handler in its own goroutine (so a DB-bound claim does not +// stall the read pump or the next heartbeat) and writes back a +// daemon:rpc_response echoing the request id. A missing handler or a full +// in-flight slot yields a non-2xx response so the daemon falls back to HTTP. +func (c *client) handleRPCFrame(raw json.RawMessage) { + var req protocol.RPCRequestPayload + if err := json.Unmarshal(raw, &req); err != nil { + slog.Debug("daemon websocket rpc invalid payload", "error", err, "daemon_id", c.identity.DaemonID) + return + } + if req.RequestID == "" { + slog.Debug("daemon websocket rpc missing request_id", "daemon_id", c.identity.DaemonID) + return + } + handler := c.hub.rpcHandler() + if handler == nil { + c.sendRPCResponse(req.RequestID, http.StatusServiceUnavailable, nil, "rpc handler unavailable") + return + } + // Bound concurrent handlers; if saturated, tell the daemon to fall back + // rather than queueing unbounded work on one socket. + select { + case c.rpcSem <- struct{}{}: + default: + c.sendRPCResponse(req.RequestID, http.StatusTooManyRequests, nil, "too many in-flight rpc requests") + return + } + 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) + if err != nil { + if status < 400 { + status = http.StatusInternalServerError + } + c.sendRPCResponse(req.RequestID, status, nil, err.Error()) + return + } + c.sendRPCResponse(req.RequestID, status, body, "") + }() +} + +func (c *client) sendRPCResponse(requestID string, status int, body json.RawMessage, errMsg string) { + frame, err := json.Marshal(protocol.Message{ + Type: protocol.EventDaemonRPCResponse, + Payload: mustMarshalRaw(protocol.RPCResponsePayload{ + RequestID: requestID, + Status: status, + Body: body, + Error: errMsg, + }), + }) + if err != nil { + 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", + "daemon_id", c.identity.DaemonID, "request_id", requestID) + } +} + // handleHeartbeatFrame processes an inbound daemon:heartbeat from the daemon, // invokes the hub's handler, and writes back a daemon:heartbeat_ack. func (c *client) handleHeartbeatFrame(raw json.RawMessage) { diff --git a/server/internal/daemonws/rpc_dispatch_test.go b/server/internal/daemonws/rpc_dispatch_test.go new file mode 100644 index 0000000000..2e50a2780c --- /dev/null +++ b/server/internal/daemonws/rpc_dispatch_test.go @@ -0,0 +1,116 @@ +package daemonws + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/multica-ai/multica/server/pkg/protocol" +) + +// dialRPCTestConn spins up a hub-backed WS server and returns a connected +// client conn plus the hub. +func dialRPCTestConn(t *testing.T, hub *Hub, identity ClientIdentity) *websocket.Conn { + t.Helper() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hub.HandleWebSocket(w, r, identity) + })) + t.Cleanup(server.Close) + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("Dial: %v", err) + } + t.Cleanup(func() { conn.Close() }) + return conn +} + +func sendRPCRequest(t *testing.T, conn *websocket.Conn, req protocol.RPCRequestPayload) protocol.RPCResponsePayload { + t.Helper() + frame, err := json.Marshal(protocol.Message{ + Type: protocol.EventDaemonRPCRequest, + Payload: mustMarshalRaw(req), + }) + if err != nil { + t.Fatalf("marshal request: %v", err) + } + if err := conn.WriteMessage(websocket.TextMessage, frame); err != nil { + t.Fatalf("write: %v", err) + } + if err := conn.SetReadDeadline(time.Now().Add(2 * time.Second)); err != nil { + t.Fatalf("SetReadDeadline: %v", err) + } + _, raw, err := conn.ReadMessage() + if err != nil { + t.Fatalf("read: %v", err) + } + var msg protocol.Message + if err := json.Unmarshal(raw, &msg); err != nil { + t.Fatalf("unmarshal msg: %v", err) + } + if msg.Type != protocol.EventDaemonRPCResponse { + t.Fatalf("type = %q, want %q", msg.Type, protocol.EventDaemonRPCResponse) + } + var resp protocol.RPCResponsePayload + if err := json.Unmarshal(msg.Payload, &resp); err != nil { + t.Fatalf("unmarshal resp: %v", err) + } + return resp +} + +// TestRPCDispatch_RoundTrip pins the generic WS request/response contract +// (MUL-4257): a daemon:rpc_request is routed to the registered handler with the +// connection's identity, and the daemon:rpc_response echoes the request id and +// carries the handler's body. +func TestRPCDispatch_RoundTrip(t *testing.T) { + hub := NewHub() + var gotMethod, gotDaemonID string + hub.SetRPCHandler(func(ctx context.Context, identity ClientIdentity, method string, body json.RawMessage) (int, json.RawMessage, error) { + gotMethod = method + gotDaemonID = identity.DaemonID + return http.StatusOK, json.RawMessage(`{"ok":true}`), nil + }) + conn := dialRPCTestConn(t, hub, ClientIdentity{DaemonID: "daemon-1", RuntimeIDs: []string{"rt-1"}}) + + resp := sendRPCRequest(t, conn, protocol.RPCRequestPayload{ + RequestID: "req-1", + Method: "tasks.claim", + Body: json.RawMessage(`{"max_tasks":3}`), + }) + if resp.RequestID != "req-1" || resp.Status != http.StatusOK || string(resp.Body) != `{"ok":true}` { + t.Fatalf("resp = %+v, want req-1/200/{ok:true}", resp) + } + if gotMethod != "tasks.claim" || gotDaemonID != "daemon-1" { + t.Fatalf("handler saw method=%q daemon=%q, want tasks.claim/daemon-1", gotMethod, gotDaemonID) + } +} + +// TestRPCDispatch_HandlerError maps a handler error to a non-2xx response so +// the daemon can fall back to HTTP. +func TestRPCDispatch_HandlerError(t *testing.T) { + hub := NewHub() + hub.SetRPCHandler(func(ctx context.Context, identity ClientIdentity, method string, body json.RawMessage) (int, json.RawMessage, error) { + return 0, nil, context.DeadlineExceeded + }) + conn := dialRPCTestConn(t, hub, ClientIdentity{DaemonID: "daemon-1", RuntimeIDs: []string{"rt-1"}}) + resp := sendRPCRequest(t, conn, protocol.RPCRequestPayload{RequestID: "req-2", Method: "tasks.claim"}) + if resp.RequestID != "req-2" || resp.Status < 400 || resp.Error == "" { + t.Fatalf("resp = %+v, want req-2 with 5xx + error", resp) + } +} + +// TestRPCDispatch_NoHandler returns 503 when no RPC handler is registered so +// the daemon falls back to HTTP. +func TestRPCDispatch_NoHandler(t *testing.T) { + hub := NewHub() + conn := dialRPCTestConn(t, hub, ClientIdentity{DaemonID: "daemon-1", RuntimeIDs: []string{"rt-1"}}) + resp := sendRPCRequest(t, conn, protocol.RPCRequestPayload{RequestID: "req-3", Method: "tasks.claim"}) + if resp.RequestID != "req-3" || resp.Status != http.StatusServiceUnavailable { + t.Fatalf("resp = %+v, want req-3 with 503", resp) + } +} diff --git a/server/pkg/protocol/events.go b/server/pkg/protocol/events.go index 6d90c71ee0..0a2ca5cf40 100644 --- a/server/pkg/protocol/events.go +++ b/server/pkg/protocol/events.go @@ -118,6 +118,13 @@ const ( EventDaemonRegister = "daemon:register" EventDaemonTaskAvailable = "daemon:task_available" EventDaemonRuntimeProfilesChanged = "daemon:runtime_profiles_changed" + // Generic daemon→server request/response over the WebSocket control + // connection (MUL-4257). The daemon sends EventDaemonRPCRequest with a + // correlation id + method + body; the server replies EventDaemonRPCResponse + // with the same request id. This is the transport for WS-first claim (with + // HTTP fallback) and any future daemon→server RPC. + EventDaemonRPCRequest = "daemon:rpc_request" + EventDaemonRPCResponse = "daemon:rpc_response" // GitHub integration events EventGitHubInstallationCreated = "github_installation:created" diff --git a/server/pkg/protocol/messages.go b/server/pkg/protocol/messages.go index 66e89fcdc2..dc21a67a8e 100644 --- a/server/pkg/protocol/messages.go +++ b/server/pkg/protocol/messages.go @@ -5,8 +5,35 @@ import "encoding/json" const ( DaemonCapabilitySkillBundlesV1 = "skill-bundles-v1" DaemonCapabilityCoalescedCommentsV1 = "coalesced-comments-v1" + // DaemonCapabilityRPCV1 advertises that the daemon can carry + // request/response RPCs over the WebSocket control connection (MUL-4257). + // Gated so only daemons+servers that both support it route claim over WS; + // everyone else keeps using the HTTP claim endpoint. + DaemonCapabilityRPCV1 = "rpc-v1" ) +// RPCRequestPayload is the generic daemon→server request envelope carried in a +// protocol.Message of type EventDaemonRPCRequest. RequestID correlates the +// response; Method selects the server-side handler (e.g. "tasks.claim"); Body +// is the method-specific request JSON. +type RPCRequestPayload struct { + RequestID string `json:"request_id"` + Method string `json:"method"` + Body json.RawMessage `json:"body,omitempty"` +} + +// RPCResponsePayload is the server→daemon reply, carried in a +// protocol.Message of type EventDaemonRPCResponse. RequestID echoes the +// request. Status mirrors an HTTP status so the daemon can treat WS and HTTP +// outcomes uniformly. Exactly one of Body / Error is meaningful: Body on +// success (2xx), Error on failure. +type RPCResponsePayload struct { + RequestID string `json:"request_id"` + Status int `json:"status"` + Body json.RawMessage `json:"body,omitempty"` + Error string `json:"error,omitempty"` +} + // Message is the envelope for all WebSocket messages. type Message struct { Type string `json:"type"`