1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321

// Package bridge relays events between the ACP agent session and the
// browsers connected over WebSocket.
//
// It is the meeting point of two asynchronous worlds: the agent streams
// updates and blocks on permission requests; browsers connect, disconnect and
// answer at human speed. The bridge broadcasts agent events to every client,
// replays the session history to late joiners, and routes each permission
// decision back to the exact request waiting for it.
package bridge

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"path/filepath"
	"strconv"
	"sync"

	acp "github.com/coder/acp-go-sdk"
)

// Prompter is the slice of the agent session the bridge drives. It is
// satisfied by *agent.Session and by test fakes.
type Prompter interface {
	ID() string
	Prompt(ctx context.Context, blocks []acp.ContentBlock) (acp.StopReason, error)
	Cancel(ctx context.Context) error
}

// maxHistory bounds the number of events replayed to late-joining clients;
// beyond it the oldest events are dropped (the live broadcast is unaffected).
const maxHistory = 4096

// clientBuffer is each subscriber's outgoing queue; a client too slow to
// drain it is disconnected rather than allowed to stall the whole bridge.
const clientBuffer = 256

type permissionDecision struct {
	optionId  string
	cancelled bool
}

// Bridge broadcasts agent events to WebSocket clients and routes their
// responses back to the agent. The zero value is not usable; call New.
//
// Example:
//
//	b := bridge.New(slog.Default())
//	session, _ := agent.Start(ctx, agent.Options{Handler: b, ...})
//	b.SetSession(session)
//	http.Handle("GET /ws", b.WebSocketHandler())
type Bridge struct {
	logger *slog.Logger

	mu           sync.Mutex
	session      Prompter
	root         string
	subscribers  map[chan []byte]struct{}
	history      [][]byte
	turnActive   bool
	nextRequest  int
	pendingPerms map[string]chan permissionDecision
	// pendingRequests keeps the encoded permission_request events so they
	// can be replayed to clients that connect while a request is open.
	pendingRequests map[string][]byte
}

// New creates an empty bridge; attach the agent session with SetSession.
func New(logger *slog.Logger) *Bridge {
	if logger == nil {
		logger = slog.Default()
	}
	return &Bridge{
		logger:          logger,
		subscribers:     make(map[chan []byte]struct{}),
		pendingPerms:    make(map[string]chan permissionDecision),
		pendingRequests: make(map[string][]byte),
	}
}

// SetSession attaches the agent session the bridge drives. It must be called
// before the first client message arrives.
func (b *Bridge) SetSession(session Prompter) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.session = session
}

// SetWorkspaceRoot sets the directory relative attachment paths resolve
// against (the agent's cwd). Without it, relative paths are sent as given.
func (b *Bridge) SetWorkspaceRoot(root string) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.root = root
}

// Subscribe registers a new client. It returns the client's event channel,
// the events to replay so the client catches up (history, then any pending
// permission requests), and an unsubscribe function.
func (b *Bridge) Subscribe() (events <-chan []byte, replay [][]byte, unsubscribe func()) {
	ch := make(chan []byte, clientBuffer)
	b.mu.Lock()
	defer b.mu.Unlock()

	b.subscribers[ch] = struct{}{}

	sessionId := ""
	if b.session != nil {
		sessionId = b.session.ID()
	}
	replay = append(replay, encode(Outgoing{Type: OutgoingHello, SessionId: sessionId, TurnActive: b.turnActive}))
	replay = append(replay, b.history...)
	for _, raw := range b.pendingRequests {
		replay = append(replay, raw)
	}

	return ch, replay, func() {
		b.mu.Lock()
		defer b.mu.Unlock()
		if _, ok := b.subscribers[ch]; ok {
			delete(b.subscribers, ch)
			close(ch)
		}
	}
}

// HandleIncoming processes one message from a client. Errors are reported to
// the clients as OutgoingError events, never returned: a malformed message
// from one browser must not tear down the connection handling.
func (b *Bridge) HandleIncoming(ctx context.Context, raw []byte) {
	var msg Incoming
	if err := json.Unmarshal(raw, &msg); err != nil {
		b.broadcast(encode(Outgoing{Type: OutgoingError, Message: "malformed message: " + err.Error()}), false)
		return
	}

	switch msg.Type {
	case IncomingPrompt:
		b.startTurn(msg.Text, msg.Attachments)
	case IncomingCancel:
		b.cancelTurn(ctx)
	case IncomingPermissionResponse:
		b.resolvePermission(msg)
	default:
		b.broadcast(encode(Outgoing{Type: OutgoingError, Message: fmt.Sprintf("unknown message type %q", msg.Type)}), false)
	}
}

// startTurn launches a prompt turn in the background; updates stream through
// HandleSessionUpdate while it runs.
func (b *Bridge) startTurn(text string, attachments []Attachment) {
	b.mu.Lock()
	session := b.session
	root := b.root
	if session == nil {
		b.mu.Unlock()
		b.broadcast(encode(Outgoing{Type: OutgoingError, Message: "no agent session"}), false)
		return
	}
	if b.turnActive {
		b.mu.Unlock()
		b.broadcast(encode(Outgoing{Type: OutgoingError, Message: "a turn is already running"}), false)
		return
	}
	b.turnActive = true
	b.mu.Unlock()

	b.broadcast(encode(Outgoing{Type: OutgoingUserMessage, Text: text, Attachments: attachments}), true)
	b.broadcast(encode(Outgoing{Type: OutgoingTurnStarted}), true)

	blocks := promptBlocks(text, attachments, root)
	go func() {
		// The turn outlives the WebSocket message that started it, so it
		// runs under its own context, ended by session/cancel only.
		stop, err := session.Prompt(context.Background(), blocks)

		b.mu.Lock()
		b.turnActive = false
		b.mu.Unlock()

		if err != nil {
			b.logger.Error("prompt failed", "error", err)
			b.broadcast(encode(Outgoing{Type: OutgoingError, Message: "prompt failed: " + err.Error()}), true)
			b.broadcast(encode(Outgoing{Type: OutgoingTurnEnded, StopReason: string(acp.StopReasonRefusal)}), true)
			return
		}
		b.broadcast(encode(Outgoing{Type: OutgoingTurnEnded, StopReason: string(stop)}), true)
	}()
}

// promptBlocks builds the ACP prompt: the text first, then one resource_link
// per attachment (file:// URI, named after the mention), which every ACP
// agent must accept.
//
// Example:
//
//	promptBlocks("Explain @src/main.go", []Attachment{{Path: "src/main.go"}}, "/work")
//	// [text "Explain @src/main.go", resource_link name=src/main.go uri=file:///work/src/main.go]
func promptBlocks(text string, attachments []Attachment, root string) []acp.ContentBlock {
	blocks := []acp.ContentBlock{acp.TextBlock(text)}
	for _, attachment := range attachments {
		if attachment.Path == "" {
			continue
		}
		path := attachment.Path
		if !filepath.IsAbs(path) && root != "" {
			path = filepath.Join(root, path)
		}
		name := attachment.Name
		if name == "" {
			name = attachment.Path
		}
		blocks = append(blocks, acp.ResourceLinkBlock(name, "file://"+filepath.ToSlash(path)))
	}
	return blocks
}

func (b *Bridge) cancelTurn(ctx context.Context) {
	b.mu.Lock()
	session := b.session
	b.mu.Unlock()
	if session == nil {
		return
	}
	if err := session.Cancel(ctx); err != nil {
		b.logger.Error("cancel failed", "error", err)
		b.broadcast(encode(Outgoing{Type: OutgoingError, Message: "cancel failed: " + err.Error()}), false)
	}
}

func (b *Bridge) resolvePermission(msg Incoming) {
	b.mu.Lock()
	waiting, ok := b.pendingPerms[msg.RequestId]
	if ok {
		delete(b.pendingPerms, msg.RequestId)
		delete(b.pendingRequests, msg.RequestId)
	}
	b.mu.Unlock()
	if !ok {
		// Either already answered by another client or unknown: ignore.
		return
	}
	waiting <- permissionDecision{optionId: msg.OptionId, cancelled: msg.Cancelled}
	b.broadcast(encode(Outgoing{Type: OutgoingPermissionResolved, RequestId: msg.RequestId}), true)
}

// HandleSessionUpdate implements agent.Handler: every ACP update is relayed
// verbatim to the clients and recorded for replay.
func (b *Bridge) HandleSessionUpdate(_ context.Context, notification acp.SessionNotification) {
	update, err := json.Marshal(notification.Update)
	if err != nil {
		b.logger.Error("marshal session update", "error", err)
		return
	}
	b.broadcast(encode(Outgoing{Type: OutgoingSessionUpdate, Update: update}), true)
}

// HandlePermissionRequest implements agent.Handler: the request is broadcast
// to the clients and the call blocks until one of them answers or the agent
// cancels the turn.
func (b *Bridge) HandlePermissionRequest(ctx context.Context, request acp.RequestPermissionRequest) (acp.RequestPermissionResponse, error) {
	rawRequest, err := json.Marshal(request)
	if err != nil {
		return acp.RequestPermissionResponse{}, fmt.Errorf("marshal permission request: %w", err)
	}

	decisionCh := make(chan permissionDecision, 1)
	b.mu.Lock()
	b.nextRequest++
	requestId := "perm-" + strconv.Itoa(b.nextRequest)
	event := encode(Outgoing{Type: OutgoingPermissionRequest, RequestId: requestId, Request: rawRequest})
	b.pendingPerms[requestId] = decisionCh
	b.pendingRequests[requestId] = event
	b.mu.Unlock()

	b.broadcast(event, false)

	select {
	case decision := <-decisionCh:
		if decision.cancelled {
			return acp.RequestPermissionResponse{Outcome: acp.NewRequestPermissionOutcomeCancelled()}, nil
		}
		return acp.RequestPermissionResponse{Outcome: acp.NewRequestPermissionOutcomeSelected(acp.PermissionOptionId(decision.optionId))}, nil
	case <-ctx.Done():
		// The agent cancelled the turn (or disconnected) while waiting.
		b.mu.Lock()
		delete(b.pendingPerms, requestId)
		delete(b.pendingRequests, requestId)
		b.mu.Unlock()
		b.broadcast(encode(Outgoing{Type: OutgoingPermissionResolved, RequestId: requestId}), true)
		return acp.RequestPermissionResponse{Outcome: acp.NewRequestPermissionOutcomeCancelled()}, nil
	}
}

// broadcast sends an event to every connected client, disconnecting any
// client whose buffer is full. With record=true the event also joins the
// replay history.
func (b *Bridge) broadcast(event []byte, record bool) {
	b.mu.Lock()
	defer b.mu.Unlock()

	if record {
		b.history = append(b.history, event)
		if len(b.history) > maxHistory {
			b.history = b.history[len(b.history)-maxHistory:]
		}
	}

	for ch := range b.subscribers {
		select {
		case ch <- event:
		default:
			// The client stopped draining; drop it instead of blocking
			// the agent's event stream.
			delete(b.subscribers, ch)
			close(ch)
			b.logger.Warn("dropped a slow websocket client")
		}
	}
}