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:
tao.chen
2026-07-06 16:43:17 +08:00
parent 3d34de96bf
commit 031451fe3e
10 changed files with 473 additions and 14 deletions
+109
View File
@@ -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
}
}
}
+112
View File
@@ -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")
}
}
+1
View File
@@ -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
}