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:
+98
-10
@@ -31,6 +31,8 @@ type Client struct {
|
||||
sessionID string
|
||||
transport *transport
|
||||
history *History
|
||||
streamChs []chan<- StreamEvent
|
||||
notifyWG sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewClient creates an ACP client for the given workspace.
|
||||
@@ -159,6 +161,73 @@ func (c *Client) Prompt(ctx context.Context, content string) (PromptResult, erro
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// sendStream writes ev to out without blocking. It is used for both chunk
|
||||
// notifications and terminal events so that a slow consumer cannot stall the
|
||||
// transport goroutine.
|
||||
func (c *Client) sendStream(out chan<- StreamEvent, ev StreamEvent) {
|
||||
select {
|
||||
case out <- ev:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// AddStream registers a new streaming consumer. It must be paired with RemoveStream.
|
||||
func (c *Client) AddStream(ch chan<- StreamEvent) int {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.streamChs = append(c.streamChs, ch)
|
||||
return len(c.streamChs)
|
||||
}
|
||||
|
||||
// RemoveStream unregisters a streaming consumer previously added with AddStream.
|
||||
func (c *Client) RemoveStream(ch chan<- StreamEvent) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
for i, existing := range c.streamChs {
|
||||
if existing == ch {
|
||||
c.streamChs[i] = c.streamChs[len(c.streamChs)-1]
|
||||
c.streamChs = c.streamChs[:len(c.streamChs)-1]
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Stream sends a user prompt and forwards streamed chunks to out as they arrive.
|
||||
// It emits exactly one terminal event, either complete or error, before returning.
|
||||
func (c *Client) Stream(ctx context.Context, content string, out chan<- StreamEvent) error {
|
||||
c.mu.RLock()
|
||||
if !c.ready {
|
||||
c.mu.RUnlock()
|
||||
c.sendStream(out, StreamEvent{Type: "error", Error: "client not ready"})
|
||||
return util.New(util.CodeConflict, "client not ready")
|
||||
}
|
||||
sessionID := c.sessionID
|
||||
tr := c.transport
|
||||
c.mu.RUnlock()
|
||||
|
||||
c.history.Add("user", content)
|
||||
|
||||
c.AddStream(out)
|
||||
defer c.RemoveStream(out)
|
||||
|
||||
params := SessionPromptParams{
|
||||
SessionID: sessionID,
|
||||
Prompt: []PromptMessage{{Type: "text", Text: content}},
|
||||
}
|
||||
var result PromptResult
|
||||
if err := tr.call(ctx, "session/prompt", params, &result); err != nil {
|
||||
c.sendStream(out, StreamEvent{Type: "error", Error: err.Error()})
|
||||
return util.Wrap(util.CodeInternal, "session/prompt failed", err)
|
||||
}
|
||||
|
||||
// Wait for any chunk notifications that were already in flight to be
|
||||
// delivered before we emit the terminal complete event.
|
||||
c.notifyWG.Wait()
|
||||
|
||||
c.sendStream(out, StreamEvent{Type: "complete", StopReason: result.StopReason})
|
||||
return nil
|
||||
}
|
||||
|
||||
// Cancel sends a best-effort session/cancel notification.
|
||||
func (c *Client) Cancel() error {
|
||||
c.mu.RLock()
|
||||
@@ -230,26 +299,45 @@ func (c *Client) handleNotification(method string, params json.RawMessage) {
|
||||
c.lg.Debug("acp drop notification", "workspace_id", c.workspaceID, "method", method)
|
||||
return
|
||||
}
|
||||
|
||||
c.notifyWG.Add(1)
|
||||
defer c.notifyWG.Done()
|
||||
|
||||
var up SessionUpdateParams
|
||||
if err := json.Unmarshal(params, &up); err != nil {
|
||||
c.lg.Error("acp invalid session/update", "workspace_id", c.workspaceID, "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
var role string
|
||||
switch up.Update.SessionUpdate {
|
||||
case "agent_message_chunk":
|
||||
if up.Update.Content.Type == "text" {
|
||||
c.history.AppendOrCreate("agent", up.Update.MessageID, up.Update.Content.Text)
|
||||
} else {
|
||||
c.lg.Debug("acp drop non-text agent chunk", "workspace_id", c.workspaceID, "type", up.Update.Content.Type)
|
||||
}
|
||||
role = "agent"
|
||||
case "user_message_chunk":
|
||||
if up.Update.Content.Type == "text" {
|
||||
c.history.AppendOrCreate("user", up.Update.MessageID, up.Update.Content.Text)
|
||||
} else {
|
||||
c.lg.Debug("acp drop non-text user chunk", "workspace_id", c.workspaceID, "type", up.Update.Content.Type)
|
||||
}
|
||||
role = "user"
|
||||
default:
|
||||
c.lg.Debug("acp drop update", "workspace_id", c.workspaceID, "type", up.Update.SessionUpdate)
|
||||
return
|
||||
}
|
||||
|
||||
if up.Update.Content.Type != "text" {
|
||||
c.lg.Debug("acp drop non-text chunk", "workspace_id", c.workspaceID, "type", up.Update.Content.Type)
|
||||
return
|
||||
}
|
||||
|
||||
c.history.AppendOrCreate(role, up.Update.MessageID, up.Update.Content.Text)
|
||||
|
||||
ev := StreamEvent{
|
||||
Type: "chunk",
|
||||
MessageID: up.Update.MessageID,
|
||||
Text: up.Update.Content.Text,
|
||||
}
|
||||
c.mu.RLock()
|
||||
streams := make([]chan<- StreamEvent, len(c.streamChs))
|
||||
copy(streams, c.streamChs)
|
||||
c.mu.RUnlock()
|
||||
for _, ch := range streams {
|
||||
c.sendStream(ch, ev)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -2,6 +2,7 @@ package acp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"os"
|
||||
@@ -351,3 +352,74 @@ func TestHistoryAppendOrCreateGroupsByMessageID(t *testing.T) {
|
||||
t.Fatalf("after clear message = %+v, want text=\"after clear\" messageId=msg-1", msgs[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientStreamEmitsChunkAndComplete(t *testing.T) {
|
||||
c := &Client{
|
||||
workspaceID: "ws",
|
||||
fs: newStubFS(),
|
||||
history: NewHistory(),
|
||||
lg: slog.Default(),
|
||||
}
|
||||
stdin := &recordingWriteCloser{}
|
||||
sub := newFakeSubscription()
|
||||
tr := newTransport(stdin, sub, slog.Default())
|
||||
tr.onNotification = c.handleNotification
|
||||
c.transport = tr
|
||||
c.ready = true
|
||||
c.sessionID = "s1"
|
||||
tr.start()
|
||||
defer tr.Close()
|
||||
|
||||
out := make(chan StreamEvent, 16)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
go func() {
|
||||
for strings.TrimSpace(stdin.String()) == "" {
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
chunk1 := `{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"s1","update":{"sessionUpdate":"agent_message_chunk","messageId":"m1","content":{"type":"text","text":"hello"}}}}`
|
||||
chunk2 := `{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"s1","update":{"sessionUpdate":"agent_message_chunk","messageId":"m1","content":{"type":"text","text":" world"}}}}`
|
||||
response := `{"jsonrpc":"2.0","id":1,"result":{"stopReason":"end_turn"}}`
|
||||
sub.send([]byte(chunk1 + "\n"))
|
||||
sub.send([]byte(chunk2 + "\n"))
|
||||
sub.send([]byte(response + "\n"))
|
||||
}()
|
||||
|
||||
if err := c.Stream(ctx, "hi", out); err != nil {
|
||||
t.Fatalf("Stream error: %v", err)
|
||||
}
|
||||
close(out)
|
||||
|
||||
var chunks int
|
||||
var complete *StreamEvent
|
||||
for ev := range out {
|
||||
switch ev.Type {
|
||||
case "chunk":
|
||||
chunks++
|
||||
if ev.MessageID != "m1" {
|
||||
t.Fatalf("chunk messageId = %q, want m1", ev.MessageID)
|
||||
}
|
||||
case "complete":
|
||||
if complete != nil {
|
||||
t.Fatal("received more than one complete event")
|
||||
}
|
||||
cp := ev
|
||||
complete = &cp
|
||||
case "error":
|
||||
t.Fatalf("unexpected error event: %q", ev.Error)
|
||||
default:
|
||||
t.Fatalf("unexpected event type %q", ev.Type)
|
||||
}
|
||||
}
|
||||
|
||||
if chunks != 2 {
|
||||
t.Fatalf("chunks = %d, want 2", chunks)
|
||||
}
|
||||
if complete == nil {
|
||||
t.Fatal("expected complete event")
|
||||
}
|
||||
if complete.StopReason != "end_turn" {
|
||||
t.Fatalf("stopReason = %q, want end_turn", complete.StopReason)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -19,6 +19,17 @@ type rpcError struct {
|
||||
Data any `json:"data,omitempty"`
|
||||
}
|
||||
|
||||
// StreamEvent is the wire shape of the streaming protocol. The Type field
|
||||
// is one of "chunk", "complete", "error"; the other fields are populated
|
||||
// based on Type.
|
||||
type StreamEvent struct {
|
||||
Type string `json:"type"`
|
||||
MessageID string `json:"messageId,omitempty"`
|
||||
Text string `json:"text,omitempty"`
|
||||
StopReason string `json:"stopReason,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// InitializeParams is sent once to negotiate the ACP session.
|
||||
type InitializeParams struct {
|
||||
ProtocolVersion int `json:"protocolVersion"`
|
||||
@@ -91,10 +102,10 @@ type SessionUpdateParams struct {
|
||||
|
||||
// Update carries a single streamed update from the agent.
|
||||
type Update struct {
|
||||
SessionUpdate string `json:"sessionUpdate"`
|
||||
MessageID string `json:"messageId,omitempty"`
|
||||
SessionUpdate string `json:"sessionUpdate"`
|
||||
MessageID string `json:"messageId,omitempty"`
|
||||
Content ContentBlock `json:"content,omitempty"`
|
||||
ToolCallID string `json:"toolCallId,omitempty"`
|
||||
ToolCallID string `json:"toolCallId,omitempty"`
|
||||
}
|
||||
|
||||
// ContentBlock is a content payload in session/update notifications.
|
||||
|
||||
@@ -16,6 +16,7 @@ type Service interface {
|
||||
Status(workspaceID string) (StatusResponse, error)
|
||||
History(workspaceID string) (HistoryResponse, error)
|
||||
Prompt(workspaceID string, content string) (PromptResponse, error)
|
||||
Stream(workspaceID string, content string, out chan<- StreamEvent) error
|
||||
Cancel(workspaceID string) error
|
||||
EnsureReady(workspaceID string) error
|
||||
}
|
||||
@@ -186,3 +187,24 @@ func (s *localService) Cancel(workspaceID string) error {
|
||||
}
|
||||
return c.Cancel()
|
||||
}
|
||||
|
||||
// Stream sends a user prompt and streams response chunks to out.
|
||||
func (s *localService) Stream(workspaceID string, content string, out chan<- StreamEvent) error {
|
||||
if _, err := s.workspaces.Get(workspaceID); err != nil {
|
||||
return err
|
||||
}
|
||||
c, err := s.clientFor(workspaceID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
lock := s.workspaceLock(workspaceID)
|
||||
lock.Lock()
|
||||
defer lock.Unlock()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
|
||||
defer cancel()
|
||||
if err := c.Initialize(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
return c.Stream(ctx, content, out)
|
||||
}
|
||||
|
||||
@@ -95,7 +95,7 @@ func (t *transport) dispatchLoop() {
|
||||
case msg.ID == nil && msg.Method != "":
|
||||
// Notification from the agent.
|
||||
if t.onNotification != nil {
|
||||
go t.onNotification(msg.Method, msg.Params)
|
||||
t.onNotification(msg.Method, msg.Params)
|
||||
}
|
||||
case msg.ID != nil && (msg.Result != nil || msg.Error != nil):
|
||||
// Response to a client-initiated request.
|
||||
|
||||
Reference in New Issue
Block a user