1package server
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "net/http"
8 "net/http/httptest"
9 "testing"
10
11 "charm.land/fantasy"
12 "github.com/charmbracelet/crush/internal/agent"
13 "github.com/charmbracelet/crush/internal/app"
14 "github.com/charmbracelet/crush/internal/backend"
15 "github.com/charmbracelet/crush/internal/message"
16 "github.com/charmbracelet/crush/internal/proto"
17 "github.com/charmbracelet/crush/internal/session"
18 "github.com/google/uuid"
19 "github.com/stretchr/testify/require"
20)
21
22// stubCoordinator is a minimal agent.Coordinator that only reports
23// per-session busy state. Every other method returns a zero value so
24// the type satisfies the interface without dragging in the full
25// coordinator dependency graph.
26type stubCoordinator struct {
27 busy map[string]bool
28}
29
30func (s *stubCoordinator) Run(ctx context.Context, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) {
31 return nil, nil
32}
33func (s *stubCoordinator) Cancel(string) {}
34func (s *stubCoordinator) CancelAll() {}
35func (s *stubCoordinator) IsBusy() bool { return false }
36func (s *stubCoordinator) IsSessionBusy(id string) bool {
37 return s.busy[id]
38}
39func (s *stubCoordinator) QueuedPrompts(string) int { return 0 }
40func (s *stubCoordinator) QueuedPromptsList(string) []string { return nil }
41func (s *stubCoordinator) ClearQueue(string) {}
42func (s *stubCoordinator) Summarize(context.Context, string) error {
43 return nil
44}
45func (s *stubCoordinator) Model() agent.Model { return agent.Model{} }
46func (s *stubCoordinator) UpdateModels(context.Context) error { return nil }
47
48// stubSessions is a minimal session.Service that returns a fixed list
49// (and supports Get by ID). All other methods return zero values; the
50// IsBusy tests do not exercise them.
51type stubSessions struct {
52 session.Service // embed nil to inherit the unexported broker methods
53 all []session.Session
54}
55
56func (s *stubSessions) List(context.Context) ([]session.Session, error) {
57 return s.all, nil
58}
59
60func (s *stubSessions) Get(_ context.Context, id string) (session.Session, error) {
61 for _, sess := range s.all {
62 if sess.ID == id {
63 return sess, nil
64 }
65 }
66 return session.Session{}, errors.New("not found")
67}
68
69// buildBusyWorkspace returns a controller wired to a backend that owns
70// a single workspace whose AgentCoordinator reports the named session
71// as busy.
72func buildBusyWorkspace(t *testing.T, sessionID string, busy bool) (*controllerV1, string) {
73 t.Helper()
74
75 b := backend.New(context.Background(), nil, nil)
76 wsID := uuid.New().String()
77 coord := &stubCoordinator{busy: map[string]bool{sessionID: busy}}
78 a := &app.App{AgentCoordinator: coord}
79 a.Sessions = &stubSessions{all: []session.Session{{ID: sessionID, Title: "t"}}}
80
81 ws := &backend.Workspace{
82 ID: wsID,
83 Path: t.TempDir(),
84 App: a,
85 }
86 backend.InsertWorkspaceForTest(b, ws)
87
88 s := &Server{backend: b}
89 return &controllerV1{backend: b, server: s}, wsID
90}
91
92func TestSessionListIncludesIsBusy(t *testing.T) {
93 t.Parallel()
94 const sid = "s-busy"
95 c, wsID := buildBusyWorkspace(t, sid, true)
96
97 req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/v1/workspaces/"+wsID+"/sessions", nil)
98 req.SetPathValue("id", wsID)
99 rec := httptest.NewRecorder()
100 c.handleGetWorkspaceSessions(rec, req)
101 require.Equal(t, http.StatusOK, rec.Code)
102
103 var got []proto.Session
104 require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
105 require.Len(t, got, 1)
106 require.Equal(t, sid, got[0].ID)
107 require.True(t, got[0].IsBusy, "expected IsBusy=true for the busy session")
108}
109
110func TestSessionListIdleSessionIsNotBusy(t *testing.T) {
111 t.Parallel()
112 const sid = "s-idle"
113 c, wsID := buildBusyWorkspace(t, sid, false)
114
115 req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/v1/workspaces/"+wsID+"/sessions", nil)
116 req.SetPathValue("id", wsID)
117 rec := httptest.NewRecorder()
118 c.handleGetWorkspaceSessions(rec, req)
119 require.Equal(t, http.StatusOK, rec.Code)
120
121 var got []proto.Session
122 require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
123 require.Len(t, got, 1)
124 require.False(t, got[0].IsBusy, "expected IsBusy=false for idle session")
125}
126
127func TestSessionGetIncludesIsBusy(t *testing.T) {
128 t.Parallel()
129 const sid = "s-busy"
130 c, wsID := buildBusyWorkspace(t, sid, true)
131
132 req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/v1/workspaces/"+wsID+"/sessions/"+sid, nil)
133 req.SetPathValue("id", wsID)
134 req.SetPathValue("sid", sid)
135 rec := httptest.NewRecorder()
136 c.handleGetWorkspaceSession(rec, req)
137 require.Equal(t, http.StatusOK, rec.Code)
138
139 var got proto.Session
140 require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
141 require.Equal(t, sid, got.ID)
142 require.True(t, got.IsBusy)
143}
144
145// TestIsSessionBusyNilSafe verifies the helper tolerates a missing
146// workspace, app, or coordinator — phase A handlers rely on this so
147// they can pass GetWorkspace's result through without an extra guard.
148func TestIsSessionBusyNilSafe(t *testing.T) {
149 t.Parallel()
150
151 require.False(t, isSessionBusy(nil, "x"))
152 require.False(t, isSessionBusy(&backend.Workspace{}, "x"))
153 require.False(t, isSessionBusy(&backend.Workspace{App: &app.App{}}, "x"))
154}
155
156// TestProtoSessionIsBusyBackwardCompat verifies older consumers that
157// unmarshal proto.Session without knowing about IsBusy still succeed
158// and ignore the new field harmlessly.
159func TestProtoSessionIsBusyBackwardCompat(t *testing.T) {
160 t.Parallel()
161
162 wire := proto.Session{ID: "s1", Title: "t", IsBusy: true}
163 raw, err := json.Marshal(wire)
164 require.NoError(t, err)
165
166 // Old client shape: same struct minus IsBusy. We model this by
167 // unmarshaling into a struct that doesn't declare the field.
168 type oldSession struct {
169 ID string `json:"id"`
170 Title string `json:"title"`
171 }
172 var old oldSession
173 require.NoError(t, json.Unmarshal(raw, &old))
174 require.Equal(t, "s1", old.ID)
175 require.Equal(t, "t", old.Title)
176}
177
178// buildMultiSessionWorkspace returns a controller wired to a backend
179// that owns a workspace with the given session IDs. Used to exercise
180// AttachedClients counts across more than one session.
181func buildMultiSessionWorkspace(t *testing.T, sessionIDs ...string) (*controllerV1, *backend.Workspace) {
182 t.Helper()
183
184 b := backend.New(context.Background(), nil, nil)
185 a := &app.App{AgentCoordinator: &stubCoordinator{}}
186 sessions := make([]session.Session, len(sessionIDs))
187 for i, sid := range sessionIDs {
188 sessions[i] = session.Session{ID: sid, Title: sid}
189 }
190 a.Sessions = &stubSessions{all: sessions}
191
192 ws := &backend.Workspace{
193 ID: uuid.New().String(),
194 Path: t.TempDir(),
195 App: a,
196 }
197 backend.InsertWorkspaceForTest(b, ws)
198 // Synthetic workspaces have an incomplete App; bypass the
199 // default teardown to avoid panics when the last client detaches.
200 backend.SetWorkspaceShutdownFnForTest(ws, func() {})
201
202 s := &Server{backend: b}
203 return &controllerV1{backend: b, server: s}, ws
204}
205
206// listSessions invokes handleGetWorkspaceSessions and returns the
207// decoded response so tests can assert per-session counts.
208func listSessions(t *testing.T, c *controllerV1, wsID string) []proto.Session {
209 t.Helper()
210 req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/v1/workspaces/"+wsID+"/sessions", nil)
211 req.SetPathValue("id", wsID)
212 rec := httptest.NewRecorder()
213 c.handleGetWorkspaceSessions(rec, req)
214 require.Equal(t, http.StatusOK, rec.Code)
215 var got []proto.Session
216 require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
217 return got
218}
219
220func countsBySessionID(sessions []proto.Session) map[string]int {
221 out := make(map[string]int, len(sessions))
222 for _, s := range sessions {
223 out[s.ID] = s.AttachedClients
224 }
225 return out
226}
227
228// TestSessionListIncludesAttachedClients walks two sessions through
229// the same lifecycle covered by TestAttachedClients_BasicLifecycle in
230// the backend package, but observed at the handler boundary.
231func TestSessionListIncludesAttachedClients(t *testing.T) {
232 t.Parallel()
233 c, ws := buildMultiSessionWorkspace(t, "S1", "S2")
234
235 // No attached clients yet.
236 counts := countsBySessionID(listSessions(t, c, ws.ID))
237 require.Equal(t, 0, counts["S1"])
238 require.Equal(t, 0, counts["S2"])
239
240 // Attach A, set to S1: S1=1.
241 cidA := uuid.New().String()
242 require.NoError(t, c.backend.AttachClient(ws.ID, cidA))
243 t.Cleanup(func() { c.backend.DetachClient(ws.ID, cidA) })
244 require.NoError(t, c.backend.SetCurrentSession(ws.ID, cidA, "S1"))
245 counts = countsBySessionID(listSessions(t, c, ws.ID))
246 require.Equal(t, 1, counts["S1"])
247 require.Equal(t, 0, counts["S2"])
248
249 // Attach B, set to S1: S1=2.
250 cidB := uuid.New().String()
251 require.NoError(t, c.backend.AttachClient(ws.ID, cidB))
252 require.NoError(t, c.backend.SetCurrentSession(ws.ID, cidB, "S1"))
253 counts = countsBySessionID(listSessions(t, c, ws.ID))
254 require.Equal(t, 2, counts["S1"])
255 require.Equal(t, 0, counts["S2"])
256
257 // B switches to S2: counts redistribute.
258 require.NoError(t, c.backend.SetCurrentSession(ws.ID, cidB, "S2"))
259 counts = countsBySessionID(listSessions(t, c, ws.ID))
260 require.Equal(t, 1, counts["S1"])
261 require.Equal(t, 1, counts["S2"])
262
263 // B detaches: S2 drops to 0.
264 c.backend.DetachClient(ws.ID, cidB)
265 counts = countsBySessionID(listSessions(t, c, ws.ID))
266 require.Equal(t, 1, counts["S1"])
267 require.Equal(t, 0, counts["S2"])
268}
269
270// TestSessionListExcludesHoldOnlyClient verifies that a registered
271// client without an SSE stream (streams == 0) does not contribute to
272// AttachedClients, even though it has an entry in the workspace's
273// clients map.
274func TestSessionListExcludesHoldOnlyClient(t *testing.T) {
275 t.Parallel()
276 c, ws := buildMultiSessionWorkspace(t, "S1")
277
278 cid := uuid.New().String()
279 require.NoError(t, backend.RegisterClientForTesting(c.backend, ws, cid))
280 t.Cleanup(func() { _ = c.backend.DeleteWorkspace(ws.ID, cid) })
281
282 counts := countsBySessionID(listSessions(t, c, ws.ID))
283 require.Equal(t, 0, counts["S1"], "hold-only client must not be counted")
284}
285
286// TestSessionListExcludesUnselectedAttachedClient verifies that a
287// client with a live SSE stream but no current session
288// (currentSessionID == "") does not show up under any session's count.
289func TestSessionListExcludesUnselectedAttachedClient(t *testing.T) {
290 t.Parallel()
291 c, ws := buildMultiSessionWorkspace(t, "S1")
292
293 cid := uuid.New().String()
294 require.NoError(t, c.backend.AttachClient(ws.ID, cid))
295 t.Cleanup(func() { c.backend.DetachClient(ws.ID, cid) })
296 // Intentionally do NOT call SetCurrentSession.
297
298 counts := countsBySessionID(listSessions(t, c, ws.ID))
299 require.Equal(t, 0, counts["S1"],
300 "attached client with no current session must not contribute to S1")
301}
302
303// TestSessionGetIncludesAttachedClients verifies the single-session
304// handler also populates AttachedClients.
305func TestSessionGetIncludesAttachedClients(t *testing.T) {
306 t.Parallel()
307 c, ws := buildMultiSessionWorkspace(t, "S1")
308
309 cid := uuid.New().String()
310 require.NoError(t, c.backend.AttachClient(ws.ID, cid))
311 t.Cleanup(func() { c.backend.DetachClient(ws.ID, cid) })
312 require.NoError(t, c.backend.SetCurrentSession(ws.ID, cid, "S1"))
313
314 req := httptest.NewRequestWithContext(t.Context(), http.MethodGet,
315 "/v1/workspaces/"+ws.ID+"/sessions/S1", nil)
316 req.SetPathValue("id", ws.ID)
317 req.SetPathValue("sid", "S1")
318 rec := httptest.NewRecorder()
319 c.handleGetWorkspaceSession(rec, req)
320 require.Equal(t, http.StatusOK, rec.Code)
321
322 var got proto.Session
323 require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
324 require.Equal(t, 1, got.AttachedClients)
325}