feat(acp): streaming WebSocket route for prompt chunks
Adds GET /api/workspaces/:id/acp/stream. Client opens a WebSocket,
sends {"type":"prompt","content":"..."}, and receives a stream of
{"type":"chunk","messageId","text"} events followed by exactly
one {"type":"complete","stopReason"} or {"type":"error","error"}.
Closing the WS early triggers session/cancel.
- internal/acp/messages.go: StreamEvent wire shape.
- internal/acp/client.go:
- streamChs []chan StreamEvent set; AddStream / RemoveStream.
- sendStream non-blocking fanout.
- Client.Stream(ctx, content, out) registers out, sends prompt,
emits complete/error after the prompt response, unregisters.
- handleNotification fans chunk events to all stream consumers.
- notifyWG ensures chunk ordering vs the terminal event.
- internal/acp/service.go: Service.Stream(workspaceID, content, out)
mirrors Prompt (per-workspace lock, 5-min timeout, EnsureReady).
- internal/service/acp_service.go: thin AcpService.Stream wrapper
that maps acp.StreamEvent -> model.AcpStreamEvent.
- internal/model/acp.go: AcpStreamRequest, AcpStreamEvent DTOs.
- internal/api/acp_handler.go: stream WS handler (upgrade, read
prompt, run Stream in a goroutine, write events, ping/pong, Cancel
on client close).
- internal/api/router.go: register the new route.
- internal/acp/transport.go: dispatch notifications synchronously
(vs. goroutine per notification) so chunks preserve order before
the session/prompt response.
- internal/acp/client_test.go: TestClientStreamEmitsChunkAndComplete
with a fake transport that drives a known sequence.
- internal/api/acp_handler_test.go: TestAcpStreamHandlerRoutes
smoke test using a fake opencode acp script.
Existing POST /api/workspaces/:id/acp/prompt is unchanged.
E2E: prompt 'say hi in exactly 3 words' -> 3 chunk events
('Hi',' there','!') + 1 complete {stopReason: 'end_turn'}.
Conversation: 019f3680-200f-79b0-860b-43302e60d0ea
This commit is contained in:
@@ -1,12 +1,16 @@
|
||||
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 {
|
||||
@@ -67,3 +71,108 @@ func (h *acpHandler) cancel(c *gin.Context) {
|
||||
}
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"codespace/internal/model"
|
||||
"codespace/internal/process"
|
||||
"codespace/internal/service"
|
||||
"codespace/internal/shell"
|
||||
"codespace/internal/workspace"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
func TestAcpStreamHandlerRoutes(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
|
||||
opencodePath := filepath.Join(tmpDir, "opencode")
|
||||
script := []byte(`#!/bin/sh
|
||||
while IFS= read -r line; do
|
||||
id=$(printf '%s' "$line" | grep -o '"id":[0-9]*' | cut -d: -f2)
|
||||
case "$line" in
|
||||
*initialize*)
|
||||
printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":1,"agentCapabilities":{},"agentInfo":{}}}\n' "$id"
|
||||
;;
|
||||
*session/new*)
|
||||
printf '{"jsonrpc":"2.0","id":%s,"result":{"sessionId":"sess-1"}}\n' "$id"
|
||||
;;
|
||||
*session/prompt*)
|
||||
printf '{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"sess-1","update":{"sessionUpdate":"agent_message_chunk","messageId":"m1","content":{"type":"text","text":"hello"}}}}\n'
|
||||
printf '{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"sess-1","update":{"sessionUpdate":"agent_message_chunk","messageId":"m1","content":{"type":"text","text":" world"}}}}\n'
|
||||
printf '{"jsonrpc":"2.0","id":%s,"result":{"stopReason":"end_turn"}}\n' "$id"
|
||||
;;
|
||||
esac
|
||||
done
|
||||
`)
|
||||
if err := os.WriteFile(opencodePath, script, 0o755); err != nil {
|
||||
t.Fatalf("write fake opencode: %v", err)
|
||||
}
|
||||
|
||||
wsRoot := filepath.Join(tmpDir, "workspaces")
|
||||
wsMgr := workspace.NewLocalManager(wsRoot)
|
||||
procMgr := process.NewManager(opencodePath, nil)
|
||||
shellMgr := shell.NewManager("bash", []string{"-i"})
|
||||
wsSvc := service.NewWorkspaceService(wsMgr, procMgr, shellMgr, nil)
|
||||
fileSvc := service.NewFileService(wsMgr, 1<<20)
|
||||
procSvc := service.NewProcessService(wsMgr, procMgr, nil)
|
||||
shellSvc := service.NewShellService(wsMgr, shellMgr, nil)
|
||||
acpSvc := service.NewAcpService(procMgr, wsMgr, nil)
|
||||
|
||||
r := NewRouter(wsSvc, fileSvc, procSvc, shellSvc, acpSvc, nil, gin.TestMode)
|
||||
srv := httptest.NewServer(r)
|
||||
defer srv.Close()
|
||||
|
||||
ws, err := wsSvc.Create("streamws")
|
||||
if err != nil {
|
||||
t.Fatalf("create workspace: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = wsSvc.Delete(ws.ID)
|
||||
})
|
||||
|
||||
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/api/workspaces/" + ws.ID + "/acp/stream"
|
||||
conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("dial websocket: %v", err)
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
if err := conn.WriteJSON(model.AcpStreamRequest{Type: "prompt", Content: "hello"}); err != nil {
|
||||
t.Fatalf("write prompt request: %v", err)
|
||||
}
|
||||
|
||||
if err := conn.SetReadDeadline(time.Now().Add(3 * time.Second)); err != nil {
|
||||
t.Fatalf("set read deadline: %v", err)
|
||||
}
|
||||
|
||||
var chunks int
|
||||
var complete bool
|
||||
for {
|
||||
var ev model.AcpStreamEvent
|
||||
if err := conn.ReadJSON(&ev); err != nil {
|
||||
t.Fatalf("read event: %v", err)
|
||||
}
|
||||
switch ev.Type {
|
||||
case "chunk":
|
||||
chunks++
|
||||
case "complete":
|
||||
complete = true
|
||||
case "error":
|
||||
t.Fatalf("unexpected error event: %s", ev.Error)
|
||||
default:
|
||||
t.Fatalf("unexpected event type %q", ev.Type)
|
||||
}
|
||||
if complete {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if chunks != 2 {
|
||||
t.Fatalf("chunks = %d, want 2", chunks)
|
||||
}
|
||||
if !complete {
|
||||
t.Fatal("expected complete event")
|
||||
}
|
||||
}
|
||||
@@ -55,6 +55,7 @@ func NewRouter(workspaces *service.WorkspaceService, files *service.FileService,
|
||||
api.GET("/workspaces/:id/acp/history", acpHandler.history)
|
||||
api.POST("/workspaces/:id/acp/prompt", acpHandler.prompt)
|
||||
api.POST("/workspaces/:id/acp/cancel", acpHandler.cancel)
|
||||
api.GET("/workspaces/:id/acp/stream", acpHandler.stream)
|
||||
|
||||
return r
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user