You've already forked agentic-coding-workflow
feat(acp): WebSocket session handler (history + live + ping/pong + owner-only prompt/cancel)
This commit is contained in:
@@ -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
|
||||||
|
|||||||
@@ -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
@@ -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
|
||||||
|
}
|
||||||
|
|||||||
@@ -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
@@ -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)
|
||||||
|
|||||||
Reference in New Issue
Block a user