codespace/internal/api/acp_handler.go

package api

import (
	"encoding/json"
	"net/http"
	"sync"
	"time"

	"codespace/internal/model"
	"codespace/internal/service"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

type acpHandler struct {
	svc *service.AcpService
}

func (h *acpHandler) status(c *gin.Context) {
	id := c.Param("id")
	status, err := h.svc.Status(id)
	if err != nil {
		writeError(c, err)
		return
	}
	c.JSON(http.StatusOK, model.AcpStatusResponse{
		WorkspaceID: status.WorkspaceID,
		Ready:       status.Ready,
		SessionID:   status.SessionID,
		Running:     status.Running,
		PID:         status.PID,
		Error:       status.Error,
	})
}

func (h *acpHandler) history(c *gin.Context) {
	id := c.Param("id")
	hist, err := h.svc.History(id)
	if err != nil {
		writeError(c, err)
		return
	}
	c.JSON(http.StatusOK, hist)
}

func (h *acpHandler) prompt(c *gin.Context) {
	id := c.Param("id")
	var req model.AcpPromptRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		writeBadRequest(c, "invalid json")
		return
	}
	if req.Content == "" {
		writeBadRequest(c, "content is required")
		return
	}
	res, err := h.svc.Prompt(id, req.Content)
	if err != nil {
		writeError(c, err)
		return
	}
	c.JSON(http.StatusOK, res)
}

func (h *acpHandler) cancel(c *gin.Context) {
	id := c.Param("id")
	if err := h.svc.Cancel(id); err != nil {
		writeError(c, err)
		return
	}
	c.Status(http.StatusNoContent)
}

// stream upgrades the connection to a WebSocket and streams the prompt response.
//
// Protocol:
//
//	Client -> Server: {"type":"prompt","content":"..."}
//	Server -> Client: zero or more {"type":"chunk","messageId":"...","text":"..."}
//	Server -> Client: exactly one {"type":"complete","stopReason":"..."} or
//	                  {"type":"error","error":"..."}
//
// Closing the WebSocket from the client side cancels the in-flight prompt.
func (h *acpHandler) stream(c *gin.Context) {
	id := c.Param("id")

	upgrader := websocket.Upgrader{
		CheckOrigin:     func(r *http.Request) bool { return true },
		ReadBufferSize:  4096,
		WriteBufferSize: 4096,
	}

	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return
	}

	done := make(chan struct{})
	var closeOnce sync.Once
	closeAll := func() {
		closeOnce.Do(func() {
			conn.Close()
			close(done)
		})
	}

	conn.SetReadLimit(64 << 10)
	conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	conn.SetPongHandler(func(string) error {
		conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})

	// Read the single prompt request from the client.
	_, data, err := conn.ReadMessage()
	if err != nil {
		closeAll()
		return
	}

	var req model.AcpStreamRequest
	if err := json.Unmarshal(data, &req); err != nil || req.Type != "prompt" || req.Content == "" {
		_ = conn.WriteJSON(model.AcpStreamEvent{Type: "error", Error: "expected {type: prompt, content: ...}"})
		closeAll()
		return
	}

	out := make(chan model.AcpStreamEvent, 32)

	// Run the prompt in the background and close the event channel when done.
	go func() {
		defer close(out)
		if err := h.svc.Stream(id, req.Content, out); err != nil {
			select {
			case out <- model.AcpStreamEvent{Type: "error", Error: err.Error()}:
			default:
			}
		}
	}()

	// Write events and pings to the WebSocket.
	go func() {
		ticker := time.NewTicker(30 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-done:
				return
			case ev, ok := <-out:
				if !ok {
					conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
					closeAll()
					return
				}
				if err := conn.WriteJSON(ev); err != nil {
					closeAll()
					return
				}
			case <-ticker.C:
				if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(5*time.Second)); err != nil {
					closeAll()
					return
				}
			}
		}
	}()

	// Wait for the client to close the connection and cancel any in-flight prompt.
	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			h.svc.Cancel(id)
			closeAll()
			break
		}
	}
}