feat: agents connect by WebSocket, falling back to long polling and trying again each minute; auto, ws or lp per machine (0.1.0)

internal/agent/session.go

@@ -54,6 +54,11 @@ func (a *Agent) bridgeStream(parent context.Context, sid string, rw io.ReadWrite
 	ctx, cancel := context.WithCancel(parent)
 	defer cancel()
 	defer rw.Close()
+	defer func() {
+		if l := a.link.Load(); l != nil {
+			l.unsubscribe(sid)
+		}
+	}()
 	// A session closed from the panel (a tab reloaded, a window shut) ends
 	// the stream at once, not on its next output: an idle shell would live
 	// on until the agent stopped.
@@ -88,6 +93,21 @@ func (a *Agent) bridgeStream(parent context.Context, sid string, rw io.ReadWrite
 // pollIn performs one downstream long-poll and applies the frames.
 // Returns false when the session is gone and the loop should stop.
 func (a *Agent) pollIn(ctx context.Context, sid string, w io.Writer, onResize func(cols, rows uint16)) bool {
+	// On the WebSocket the input comes pushed; if it drops, by a long poll.
+	if l := a.link.Load(); l != nil {
+		if ch := l.subscribe(sid); ch != nil {
+			select {
+			case in := <-ch:
+				if in.gone {
+					return false
+				}
+				return applyFrames(in.frames, w, onResize)
+			case <-l.done:
+			case <-ctx.Done():
+				return false
+			}
+		}
+	}
 	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, a.cfg.Host+"/api/agent/term/"+sid+"/in", nil)
 	a.authHeaders(req)
 	resp, err := a.http.Do(req)
@@ -105,7 +125,13 @@ func (a *Agent) pollIn(ctx context.Context, sid string, w io.Writer, onResize fu
 	if err := json.NewDecoder(resp.Body).Decode(&batch); err != nil {
 		return true
 	}
-	for _, f := range batch.Frames {
+	return applyFrames(batch.Frames, w, onResize)
+}
+
+// applyFrames writes a session's input to its stream; false once the panel
+// closed it.
+func applyFrames(frames []proto.InFrame, w io.Writer, onResize func(cols, rows uint16)) bool {
+	for _, f := range frames {
 		switch f.Type {
 		case proto.FrameStdin:
 			if data, err := b64.DecodeString(f.Data); err == nil {
@@ -124,6 +150,11 @@ func (a *Agent) pollIn(ctx context.Context, sid string, w io.Writer, onResize fu
 
 // postOut sends one upstream message. Returns false if the session is gone.
 func (a *Agent) postOut(ctx context.Context, sid string, msg proto.OutMessage) bool {
+	if l := a.link.Load(); l != nil {
+		if l.send(proto.WSMsg{T: "out", SID: sid, Out: &msg}) == nil {
+			return true
+		}
+	}
 	body, _ := json.Marshal(msg)
 	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, a.cfg.Host+"/api/agent/term/"+sid+"/out", bytes.NewReader(body))
 	a.authHeaders(req)