diff --git a/go.mod b/go.mod index f7eca3d..8e73088 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index 5f207c7..61c22d3 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/internal/acp/handler.go b/internal/acp/handler.go index 948c84f..3866575 100644 --- a/internal/acp/handler.go +++ b/internal/acp/handler.go @@ -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 +} diff --git a/internal/acp/handler_test.go b/internal/acp/handler_test.go index 5ab064c..df41d60 100644 --- a/internal/acp/handler_test.go +++ b/internal/acp/handler_test.go @@ -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 diff --git a/internal/app/app.go b/internal/app/app.go index 3a09d6c..1b4d1ef 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -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)