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

This commit is contained in:
2026-05-07 17:08:09 +08:00
parent 72cd353757
commit f5f048083a
5 changed files with 196 additions and 4 deletions
+1
View File
@@ -4,6 +4,7 @@ go 1.25.0
require ( require (
github.com/alexedwards/argon2id v1.0.0 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/go-chi/chi/v5 v5.2.5
github.com/golang-migrate/migrate/v4 v4.19.1 github.com/golang-migrate/migrate/v4 v4.19.1
github.com/google/uuid v1.6.0 github.com/google/uuid v1.6.0
+2
View File
@@ -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/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 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= 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 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI=
github.com/containerd/errdefs v1.0.0/go.mod h1:+YBYIdtsnF4Iw6nWZhJcqGSg/dwvV7tyJ/kCkyJ2k+M= github.com/containerd/errdefs v1.0.0/go.mod h1:+YBYIdtsnF4Iw6nWZhJcqGSg/dwvV7tyJ/kCkyJ2k+M=
github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151Xdx3ZPPE= github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151Xdx3ZPPE=
+187 -3
View File
@@ -10,10 +10,14 @@ import (
"sort" "sort"
"strconv" "strconv"
"strings" "strings"
"time"
"github.com/go-chi/chi/v5" "github.com/go-chi/chi/v5"
"github.com/google/uuid" "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/crypto"
"github.com/yan1h/agent-coding-workflow/internal/infra/errs" "github.com/yan1h/agent-coding-workflow/internal/infra/errs"
httpx "github.com/yan1h/agent-coding-workflow/internal/transport/http" 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) 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. // Handler exposes ACP HTTP endpoints.
type Handler struct { type Handler struct {
akSvc AgentKindService akSvc AgentKindService
@@ -34,14 +45,15 @@ type Handler struct {
resolver middleware.SessionResolver resolver middleware.SessionResolver
adminLookup AdminLookup adminLookup AdminLookup
enc *crypto.Encryptor enc *crypto.Encryptor
cfg WSConfig
} }
// NewHandler constructs Handler with full dependencies. // NewHandler constructs Handler with full dependencies.
func NewHandler(ak AgentKindService, ss SessionService, sup *Supervisor, repo Repository, 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{ return &Handler{
akSvc: ak, sessSvc: ss, sup: sup, repo: repo, 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.Get("/{id}", h.getSession)
r.Delete("/{id}", h.terminateSession) r.Delete("/{id}", h.terminateSession)
r.Get("/{id}/events", h.getEvents) 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) 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 ===== // ===== DTO converters =====
// agentKindToAdminDTO is a Handler method because env-key extraction needs the encryptor. // 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) { func writeErr(w http.ResponseWriter, r *http.Request, err error) {
httpx.WriteError(w, middleware.RequestIDFromContext(r.Context()), err) httpx.WriteError(w, middleware.RequestIDFromContext(r.Context()), err)
} }
func derefStr(p *string) string {
if p == nil {
return ""
}
return *p
}
+1
View File
@@ -52,6 +52,7 @@ func mountHandlerWithRepo(t *testing.T, callerUID uuid.UUID, isAdmin bool, repo
fakeResolver{uid: callerUID}, fakeResolver{uid: callerUID},
fakeAdminLookup{admin: map[uuid.UUID]bool{callerUID: isAdmin}}, fakeAdminLookup{admin: map[uuid.UUID]bool{callerUID: isAdmin}},
enc, enc,
acp.WSConfig{},
) )
h.Mount(r) h.Mount(r)
return r return r
+5 -1
View File
@@ -327,7 +327,11 @@ func New(ctx context.Context, cfg *config.Config, dist embed.FS) (*App, error) {
MaxActivePerUser: cfg.Acp.MaxActivePerUser, 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 ===== // ===== jobs module wiring =====
jobsRepo := jobs.NewPostgresRepository(pool) jobsRepo := jobs.NewPostgresRepository(pool)