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