feat(acp): WebSocket session handler (history + live + ping/pong + owner-only prompt/cancel)

2026-05-07 17:08:09 +08:00
parent 72cd353757
commit f5f048083a
5 changed files with 196 additions and 4 deletions
@@ -4,6 +4,7 @@ go 1.25.0

 require (
 	github.com/alexedwards/argon2id v1.0.0
+	github.com/coder/websocket v1.8.14
 	github.com/go-chi/chi/v5 v5.2.5
 	github.com/golang-migrate/migrate/v4 v4.19.1
 	github.com/google/uuid v1.6.0
@@ -24,6 +24,8 @@ github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK3
 github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
 github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
 github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/coder/websocket v1.8.14 h1:9L0p0iKiNOibykf283eHkKUHHrpG7f65OE3BhhO7v9g=
+github.com/coder/websocket v1.8.14/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg=
 github.com/containerd/errdefs v1.0.0 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI=
 github.com/containerd/errdefs v1.0.0/go.mod h1:+YBYIdtsnF4Iw6nWZhJcqGSg/dwvV7tyJ/kCkyJ2k+M=
 github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151Xdx3ZPPE=
@@ -10,10 +10,14 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"time"

 	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"

+	ws "github.com/coder/websocket"
+	"github.com/coder/websocket/wsjson"
+
 	"github.com/yan1h/agent-coding-workflow/internal/infra/crypto"
 	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
 	httpx "github.com/yan1h/agent-coding-workflow/internal/transport/http"
@@ -25,6 +29,13 @@ type AdminLookup interface {
 	IsAdmin(ctx context.Context, userID uuid.UUID) (bool, error)
 }

+// WSConfig controls WebSocket heartbeat and buffering.
+type WSConfig struct {
+	PingInterval time.Duration
+	PongTimeout  time.Duration
+	SendBuffer   int
+}
+
 // Handler exposes ACP HTTP endpoints.
 type Handler struct {
 	akSvc       AgentKindService
@@ -34,14 +45,15 @@ type Handler struct {
 	resolver    middleware.SessionResolver
 	adminLookup AdminLookup
 	enc         *crypto.Encryptor
+	cfg         WSConfig
 }

 // NewHandler constructs Handler with full dependencies.
 func NewHandler(ak AgentKindService, ss SessionService, sup *Supervisor, repo Repository,
-	resolver middleware.SessionResolver, al AdminLookup, enc *crypto.Encryptor) *Handler {
+	resolver middleware.SessionResolver, al AdminLookup, enc *crypto.Encryptor, cfg WSConfig) *Handler {
 	return &Handler{
 		akSvc: ak, sessSvc: ss, sup: sup, repo: repo,
-		resolver: resolver, adminLookup: al, enc: enc,
+		resolver: resolver, adminLookup: al, enc: enc, cfg: cfg,
 	}
 }

@@ -62,7 +74,7 @@ func (h *Handler) Mount(r chi.Router) {
 		r.Get("/{id}", h.getSession)
 		r.Delete("/{id}", h.terminateSession)
 		r.Get("/{id}/events", h.getEvents)
-		// /{id}/ws — added in G3 (WebSocket)
+		r.Get("/{id}/ws", h.sessionWS)
 	})
 }

@@ -348,6 +360,171 @@ func (h *Handler) getEvents(w http.ResponseWriter, r *http.Request) {
 	httpx.WriteJSON(w, http.StatusOK, out)
 }

+// ===== WebSocket =====
+
+func (h *Handler) sessionWS(w http.ResponseWriter, r *http.Request) {
+	c, err := h.caller(r)
+	if err != nil {
+		writeErr(w, r, err)
+		return
+	}
+	id, perr := uuid.Parse(chi.URLParam(r, "id"))
+	if perr != nil {
+		writeErr(w, r, errs.New(errs.CodeInvalidInput, "bad id"))
+		return
+	}
+
+	sess, err := h.sessSvc.Get(r.Context(), c, id)
+	if err != nil {
+		writeErr(w, r, err)
+		return
+	}
+
+	var since int64
+	if v := r.URL.Query().Get("since"); v != "" {
+		if n, err := strconv.ParseInt(v, 10, 64); err == nil {
+			since = n
+		}
+	}
+
+	conn, err := ws.Accept(w, r, &ws.AcceptOptions{
+		InsecureSkipVerify: false,
+	})
+	if err != nil {
+		return
+	}
+	defer conn.Close(ws.StatusInternalError, "internal")
+
+	ctx, cancel := context.WithCancel(r.Context())
+	defer cancel()
+
+	// 1. State event
+	stateEv := map[string]any{
+		"kind":       "state",
+		"session_id": id.String(),
+		"status":     string(sess.Status),
+		"branch":     sess.Branch,
+		"since":      since,
+	}
+	if err := wsjson.Write(ctx, conn, stateEv); err != nil {
+		return
+	}
+
+	// 2. Push history
+	history, _ := h.repo.ListEventsSince(ctx, id, since, int32(5000))
+	var lastSentID = since
+	for _, e := range history {
+		env := EventEnvelope{
+			ID:        e.ID,
+			Direction: string(e.Direction),
+			Kind:      string(e.RPCKind),
+			Method:    derefStr(e.Method),
+			Payload:   e.Payload,
+			Truncated: e.Truncated,
+		}
+		if err := wsjson.Write(ctx, conn, env); err != nil {
+			return
+		}
+		lastSentID = e.ID
+	}
+
+	// 3. Live: subscribe + heartbeat + client messages
+	relay := h.sup.GetRelay(id)
+	if relay == nil {
+		conn.Close(ws.StatusNormalClosure, "session ended")
+		return
+	}
+	sub := relay.Subscribe(c.UserID)
+	defer relay.Unsubscribe(sub.ID)
+
+	// Heartbeat goroutine
+	pingTicker := time.NewTicker(h.cfg.PingInterval)
+	defer pingTicker.Stop()
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-pingTicker.C:
+				pingCtx, pcancel := context.WithTimeout(ctx, h.cfg.PongTimeout)
+				err := conn.Ping(pingCtx)
+				pcancel()
+				if err != nil {
+					cancel()
+					return
+				}
+			}
+		}
+	}()
+
+	// Client → server reader (goroutine)
+	clientMsgCh := make(chan *Message, 8)
+	go func() {
+		defer close(clientMsgCh)
+		for {
+			_, data, err := conn.Read(ctx)
+			if err != nil {
+				return
+			}
+			var m Message
+			if err := json.Unmarshal(data, &m); err != nil {
+				_ = wsjson.Write(ctx, conn, map[string]any{
+					"kind":    "client_error",
+					"message": "bad json",
+				})
+				continue
+			}
+			clientMsgCh <- &m
+		}
+	}()
+
+	// Main loop
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case env, ok := <-sub.Send:
+			if !ok {
+				conn.Close(ws.StatusNormalClosure, "relay closed")
+				return
+			}
+			if env.ID <= lastSentID && env.ID != 0 {
+				continue
+			}
+			if err := wsjson.Write(ctx, conn, env); err != nil {
+				return
+			}
+			if env.ID > lastSentID {
+				lastSentID = env.ID
+			}
+		case m, ok := <-clientMsgCh:
+			if !ok {
+				return
+			}
+			if c.UserID != sess.UserID {
+				_ = wsjson.Write(ctx, conn, map[string]any{
+					"kind":    "client_error",
+					"message": "only owner can prompt/cancel",
+				})
+				continue
+			}
+			if m.Method != "session/prompt" && m.Method != "session/cancel" {
+				_ = wsjson.Write(ctx, conn, map[string]any{
+					"kind":    "client_error",
+					"message": "method not allowed",
+				})
+				continue
+			}
+			if err := relay.SendClient(m); err != nil {
+				_ = wsjson.Write(ctx, conn, map[string]any{
+					"kind":    "client_error",
+					"message": err.Error(),
+				})
+			}
+		}
+	}
+}
+
 // ===== DTO converters =====

 // agentKindToAdminDTO is a Handler method because env-key extraction needs the encryptor.
@@ -390,3 +567,10 @@ func agentKindToPublicDTO(k *AgentKind) AgentKindPublicDTO {
 func writeErr(w http.ResponseWriter, r *http.Request, err error) {
 	httpx.WriteError(w, middleware.RequestIDFromContext(r.Context()), err)
 }
+
+func derefStr(p *string) string {
+	if p == nil {
+		return ""
+	}
+	return *p
+}
@@ -52,6 +52,7 @@ func mountHandlerWithRepo(t *testing.T, callerUID uuid.UUID, isAdmin bool, repo
 		fakeResolver{uid: callerUID},
 		fakeAdminLookup{admin: map[uuid.UUID]bool{callerUID: isAdmin}},
 		enc,
+		acp.WSConfig{},
 	)
 	h.Mount(r)
 	return r
@@ -327,7 +327,11 @@ func New(ctx context.Context, cfg *config.Config, dist embed.FS) (*App, error) {
 			MaxActivePerUser: cfg.Acp.MaxActivePerUser,
 		},
 	)
-	acp.NewHandler(akSvc, sessSvc, acpSup, acpRepo, userSvc, userAdminAdapter{svc: userSvc}, enc).Mount(r)
+	acp.NewHandler(akSvc, sessSvc, acpSup, acpRepo, userSvc, userAdminAdapter{svc: userSvc}, enc, acp.WSConfig{
+		PingInterval: cfg.Acp.WSPingInterval,
+		PongTimeout:  cfg.Acp.WSPongTimeout,
+		SendBuffer:   cfg.Acp.WSSendBuffer,
+	}).Mount(r)

 	// ===== jobs module wiring =====
 	jobsRepo := jobs.NewPostgresRepository(pool)